From 131cda28242956a2c4eb18f77256d1daba067935 Mon Sep 17 00:00:00 2001 From: James Rasell Date: Mon, 21 Mar 2022 09:49:39 +0100 Subject: [PATCH] client: modify service wrapper to accomodate restore behaviour. --- client/serviceregistration/wrapper/wrapper.go | 26 ++++++++++++------- .../wrapper/wrapper_test.go | 6 ++--- 2 files changed, 20 insertions(+), 12 deletions(-) diff --git a/client/serviceregistration/wrapper/wrapper.go b/client/serviceregistration/wrapper/wrapper.go index 37b5e4f61..1c52dfd19 100644 --- a/client/serviceregistration/wrapper/wrapper.go +++ b/client/serviceregistration/wrapper/wrapper.go @@ -66,20 +66,28 @@ func (h *HandlerWrapper) RegisterWorkload(workload *serviceregistration.Workload // workload unless the provider is unknown. func (h *HandlerWrapper) RemoveWorkload(services *serviceregistration.WorkloadServices) { - // Don't rely on callers to check there are no services to remove. - if len(services.Services) == 0 { - return + var provider string + + // It is possible the services field is empty depending on the exact + // situation which resulted in the call. + if len(services.Services) > 0 { + provider = services.RegistrationProvider() } - provider := services.RegistrationProvider() - - // Call the correct provider. In the case it is not supported, we can't do - // much apart from log, although we should never reach this point because - // of job validation. + // Call the correct provider, if we have managed to identify it. An empty + // string means you didn't find a provider, therefore default to consul. + // + // In certain situations this function is called with zero services, + // therefore meaning we make an assumption on the provider. When this + // happens, we need to ensure the allocation is removed from the Consul + // implementation. This tracking (allocRegistrations) is used by the + // allochealth tracker and so is critical to be removed. The test + // allocrunner.TestAllocRunner_Restore_RunningTerminal covers the case + // described here. switch provider { case structs.ServiceProviderNomad: h.nomadServiceProvider.RemoveWorkload(services) - case structs.ServiceProviderConsul: + case structs.ServiceProviderConsul, "": h.consulServiceProvider.RemoveWorkload(services) default: h.log.Error("unknown service registration provider", "provider", provider) diff --git a/client/serviceregistration/wrapper/wrapper_test.go b/client/serviceregistration/wrapper/wrapper_test.go index 81dfd7c1f..2acb82376 100644 --- a/client/serviceregistration/wrapper/wrapper_test.go +++ b/client/serviceregistration/wrapper/wrapper_test.go @@ -131,10 +131,10 @@ func TestHandlerWrapper_RemoveWorkload(t *testing.T) { // Generate the test wrapper and provider mocks. wrapper, consul, nomad := setupTestWrapper() - // Call the function with no services and check that nothing is - // registered. + // Call the function with no services and check that consul is + // defaulted to. wrapper.RemoveWorkload(&serviceregistration.WorkloadServices{}) - require.Len(t, consul.GetOps(), 0) + require.Len(t, consul.GetOps(), 1) require.Len(t, nomad.GetOps(), 0) }, name: "zero services",