1 files changed, 49 insertions, 39 deletions
diff --git a/src/zenserver/hub/zenhubserver.cpp b/src/zenserver/hub/zenhubserver.cpp
index 1390d112e..ebc2cf2f1 100644
--- a/src/zenserver/hub/zenhubserver.cpp
+++ b/src/zenserver/hub/zenhubserver.cpp
@@ -468,47 +468,57 @@ ZenHubServer::OnModuleStateChanged(std::string_view					 HubInstanceId,
 		return;
 	}
 
-	if (NewState == HubInstanceState::Provisioning || NewState == HubInstanceState::Provisioned)
+	switch (NewState)
 	{
-		consul::ServiceRegistrationInfo ServiceInfo{
-			.ServiceId				= std::string(ModuleId),
-			.ServiceName			= "zen-storage",
-			.Port					= Info.Port,
-			.HealthEndpoint			= "health",
-			.Tags					= std::vector<std::pair<std::string, std::string>>{std::make_pair("module", std::string(ModuleId)),
-																	   std::make_pair("zen-hub", std::string(HubInstanceId)),
-																	   std::make_pair("version", std::string(ZEN_CFG_VERSION))},
-			.HealthIntervalSeconds	= NewState == HubInstanceState::Provisioning ? 0u : m_ConsulHealthIntervalSeconds,
-			.DeregisterAfterSeconds = NewState == HubInstanceState::Provisioning ? 0u : m_ConsulDeregisterAfterSeconds,
-			.InitialStatus			= NewState == HubInstanceState::Provisioned ? "passing" : ""};
-
-		if (!m_ConsulClient->RegisterService(ServiceInfo))
-		{
-			ZEN_WARN("Failed to register storage server instance for module '{}' with Consul, continuing anyway", ModuleId);
-		}
-		else
-		{
-			ZEN_INFO("Registered storage server instance for module '{}' at port {} with Consul as '{}'",
-					 ModuleId,
-					 Info.Port,
-					 ServiceInfo.ServiceName);
-		}
-	}
-	else if (NewState == HubInstanceState::Unprovisioned)
-	{
-		if (!m_ConsulClient->DeregisterService(ModuleId))
-		{
-			ZEN_WARN("Failed to deregister storage server instance for module '{}' at port {} from Consul, continuing anyway",
-					 ModuleId,
-					 Info.Port);
-		}
-		else
-		{
-			ZEN_INFO("Deregistered storage server instance for module '{}' at port {} from Consul", ModuleId, Info.Port);
-		}
+		case HubInstanceState::Provisioning:
+		case HubInstanceState::Waking:
+		case HubInstanceState::Recovering:
+		case HubInstanceState::Provisioned:
+			{
+				const bool IsProvisioned = NewState == HubInstanceState::Provisioned;
+
+				consul::ServiceRegistrationInfo ServiceInfo{
+					.ServiceId		= std::string(ModuleId),
+					.ServiceName	= "zen-storage",
+					.Port			= Info.Port,
+					.HealthEndpoint = "health",
+					.Tags			= std::vector<std::pair<std::string, std::string>>{std::make_pair("module", std::string(ModuleId)),
+																			   std::make_pair("zen-hub", std::string(HubInstanceId)),
+																			   std::make_pair("version", std::string(ZEN_CFG_VERSION))},
+					.HealthIntervalSeconds	= IsProvisioned ? m_ConsulHealthIntervalSeconds : 0u,
+					.DeregisterAfterSeconds = IsProvisioned ? m_ConsulDeregisterAfterSeconds : 0u,
+					.InitialStatus			= IsProvisioned ? "passing" : ""};
+
+				m_ConsulClient->RegisterService(ServiceInfo);
+				ZEN_INFO("Submitted Consul registration for storage server instance for module '{}' at port {} as '{}'",
+						 ModuleId,
+						 Info.Port,
+						 ServiceInfo.ServiceName);
+				break;
+			}
+		case HubInstanceState::Deprovisioning:
+		case HubInstanceState::Hibernating:
+		case HubInstanceState::Obliterating:
+		case HubInstanceState::Crashed:
+		case HubInstanceState::Hibernated:
+		case HubInstanceState::Unprovisioned:
+			{
+				// A Consul registration is "live" while the module is in a register-state
+				// (Provisioning / Waking / Recovering / Provisioned). Deregister once when
+				// we leave a register-state into any non-register-state
+				const bool WasRegisteredState =
+					PreviousState == HubInstanceState::Provisioning || PreviousState == HubInstanceState::Waking ||
+					PreviousState == HubInstanceState::Recovering || PreviousState == HubInstanceState::Provisioned;
+				if (WasRegisteredState)
+				{
+					m_ConsulClient->DeregisterService(ModuleId);
+					ZEN_INFO("Submitted Consul deregistration for storage server instance for module '{}' at port {}", ModuleId, Info.Port);
+				}
+			}
+			break;
+		default:
+			break;
 	}
-	// Transitional states (Waking, Recovering, Crashed) and stable states
-	// not handled above (Hibernated) are intentionally ignored by Consul.
 }
 
 int