Deregistering service instances during shutdown

Author: Auri Munoz

extensions/smallrye-stork/deployment/src/main/java/io/quarkus/stork/deployment/SmallRyeStorkProcessor.java

@@ -100,31 +100,29 @@ void checkThatTheKubernetesExtensionIsUsedWhenKubernetesServiceDiscoveryInOnTheC
     void initializeStork(SmallRyeStorkRecorder storkRecorder, ShutdownContextBuildItem shutdown, VertxBuildItem vertx,
             StorkConfiguration configuration) {
         storkRecorder.initialize(shutdown, vertx.getVertx(), configuration);
-        if (QuarkusClassLoader.isClassPresentAtRuntime(SERVICE_REGISTRAR_PROVIDER)) {
-            storkRecorder.deregisterServiceInstance(shutdown, configuration);
-        }
     }
 
     @BuildStep
     @Record(ExecutionTime.RUNTIME_INIT)
     void checkStorkConsulRegistrar(BuildProducer<StorkRegistrationBuildItem> registration,
             BuildProducer<RunTimeConfigurationDefaultBuildItem> config,
-            StorkRegistrarConfigRecorder initializerRecorder, StorkConfiguration configuration) {
+            StorkRegistrarConfigRecorder registrarConfigRecorder, StorkConfiguration configuration) {
         if (QuarkusClassLoader.isClassPresentAtRuntime(CONSUL_SERVICE_REGISTRAR_PROVIDER)) {
-            initializerRecorder.setupServiceRegistrarConfig(configuration, CONSUL_SERVICE_REGISTRAR_TYPE);
+            registrarConfigRecorder.setupServiceRegistrarConfig(configuration, CONSUL_SERVICE_REGISTRAR_TYPE);
         } else if (QuarkusClassLoader.isClassPresentAtRuntime(EUREKA_SERVICE_REGISTRAR_PROVIDER)) {
-            initializerRecorder.setupServiceRegistrarConfig(configuration, EUREKA_SERVICE_REGISTRAR_TYPE);
+            registrarConfigRecorder.setupServiceRegistrarConfig(configuration, EUREKA_SERVICE_REGISTRAR_TYPE);
         }
         registration.produce(new StorkRegistrationBuildItem());
 
     }
 
     @BuildStep
     @Record(ExecutionTime.RUNTIME_INIT)
-    void registerServiceInstance(StorkInitializedBuildItem storkInitializedBuildItem,
+    void registerServiceInstance(StorkInitializedBuildItem storkInitializedBuildItem, ShutdownContextBuildItem shutdown,
             SmallRyeStorkRegistrationRecorder registrationRecorder, StorkConfiguration configuration) {
         if (QuarkusClassLoader.isClassPresentAtRuntime(SERVICE_REGISTRAR_PROVIDER)) {
             registrationRecorder.registerServiceInstance(configuration);
+            registrationRecorder.deregisterServiceInstance(shutdown, configuration);
         }
     }
 

extensions/smallrye-stork/runtime/src/main/java/io/quarkus/stork/SmallRyeStorkRecorder.java

@@ -1,15 +1,10 @@
 package io.quarkus.stork;
 
-import java.net.InetAddress;
 import java.util.List;
-import java.util.Map;
 
 import jakarta.enterprise.inject.Instance;
 import jakarta.enterprise.inject.spi.CDI;
 
-import org.eclipse.microprofile.config.Config;
-import org.eclipse.microprofile.config.ConfigProvider;
-
 import io.quarkus.runtime.RuntimeValue;
 import io.quarkus.runtime.ShutdownContext;
 import io.quarkus.runtime.annotations.Recorder;
@@ -40,35 +35,4 @@ public void run() {
         });
     }
 
-    public void deregisterServiceInstance(ShutdownContext shutdown, StorkConfiguration configuration) {
-        shutdown.addLastShutdownTask(new Runnable() {
-            @Override
-            public void run() {
-                Stork.shutdown();
-            }
-        });
-    }
-
-    private void deregisterServiceInstance(StorkConfiguration configuration) {
-        List<ServiceConfig> serviceConfigs = StorkConfigUtil.toStorkServiceConfig(configuration);
-        List<ServiceConfig> registrationConfigs = serviceConfigs.stream()
-                .filter(serviceConfig -> serviceConfig.serviceRegistrar() != null).toList();
-
-        registrationConfigs.get(0).serviceName();
-        Config quarkusConfig = ConfigProvider.getConfig();
-        for (ServiceConfig serviceConfig : serviceConfigs) {
-            String serviceName = serviceConfig.serviceName();
-            Map<String, String> parameters = serviceConfig.serviceRegistrar().parameters();
-            String host = parameters.containsKey("ip-address") ? parameters.get("ip-address")
-                    : quarkusConfig.getValue("quarkus.http.host", String.class);
-            int port = parameters.containsKey("port") ? Integer.parseInt(parameters.get("port"))
-                    : Integer.parseInt(quarkusConfig.getValue("quarkus.http.port", String.class));
-            if (host == null || host.isEmpty()) {
-                InetAddress inetAddress = StorkConfigUtil.detectAddress();
-                host = inetAddress != null ? inetAddress.getHostAddress() : host;
-            }
-            Stork.getInstance().getService(serviceName).getServiceRegistrar().registerServiceInstance(serviceName, host,
-                    port).await().indefinitely();
-        }
-    }
 }

extensions/smallrye-stork/runtime/src/main/java/io/quarkus/stork/SmallRyeStorkRegistrationRecorder.java

@@ -7,6 +7,7 @@
 import org.eclipse.microprofile.config.Config;
 import org.eclipse.microprofile.config.ConfigProvider;
 
+import io.quarkus.runtime.ShutdownContext;
 import io.quarkus.runtime.annotations.Recorder;
 import io.smallrye.stork.Stork;
 import io.smallrye.stork.api.config.ServiceConfig;
@@ -32,4 +33,22 @@ public void registerServiceInstance(StorkConfiguration configuration) {
                     port).await().indefinitely();
         }
     }
+
+    public void deregisterServiceInstance(ShutdownContext shutdown, StorkConfiguration configuration) {
+        shutdown.addLastShutdownTask(new Runnable() {
+            @Override
+            public void run() {
+                deregisterServiceInstance(configuration);
+            }
+        });
+    }
+
+    private void deregisterServiceInstance(StorkConfiguration configuration) {
+        List<ServiceConfig> serviceConfigs = StorkConfigUtil.toStorkServiceConfig(configuration);
+        for (ServiceConfig serviceConfig : serviceConfigs) {
+            String serviceName = serviceConfig.serviceName();
+            Stork.getInstance().getService(serviceName).getServiceRegistrar().deregisterServiceInstance(serviceName).await()
+                    .indefinitely();
+        }
+    }
 }