refactor: ShutdownHookManager

inspectIT · Nov 6, 2024 · 374eb49 · 374eb49
1 parent a65f3d5
commit 374eb49
Show file tree

Hide file tree

Showing 8 changed files with 124 additions and 32 deletions.
diff --git a/...ctit-gepard-agent/src/main/java/rocks/inspectit/gepard/agent/InspectitAgentExtension.java b/...ctit-gepard-agent/src/main/java/rocks/inspectit/gepard/agent/InspectitAgentExtension.java
@@ -42,10 +42,6 @@ public AgentBuilder extend(AgentBuilder agentBuilder, ConfigProperties config) {
     BootstrapManager bootstrapManager = BootstrapManager.create();
     bootstrapManager.appendToBootstrapClassLoader();
 
-    // Notify configuration server about this agent
-    NotificationManager notificationManager = NotificationManager.create();
-    notificationManager.sendStartNotification();
-
     // Set our global OpenTelemetry instance. For now, we use the Agent SDK
     OpenTelemetryAccessor.setOpenTelemetry(GlobalOpenTelemetry.get());
 
@@ -73,6 +69,9 @@ public AgentBuilder extend(AgentBuilder agentBuilder, ConfigProperties config) {
     ConfigurationManager configurationManager = ConfigurationManager.create();
     configurationManager.loadConfiguration();
 
+    // Notify configuration server about this agent
+    NotificationManager notificationManager = NotificationManager.create();
+    notificationManager.sendStartNotification();
     // Set up shutdown notification to configuration server
     notificationManager.setUpShutdownNotification();
 

diff --git a/...ent/src/main/java/rocks/inspectit/gepard/agent/internal/shutdown/ShutdownHookManager.java b/...ent/src/main/java/rocks/inspectit/gepard/agent/internal/shutdown/ShutdownHookManager.java
@@ -1,22 +1,30 @@
 /* (C) 2024 */
 package rocks.inspectit.gepard.agent.internal.shutdown;
 
+import com.google.common.annotations.VisibleForTesting;
 import java.util.*;
+import java.util.concurrent.atomic.AtomicBoolean;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 /**
  * Responsible for executing the agents shutdown hooks in the order they are added. We should try to
  * keep the amount of shutdown hooks as well as the logic inside them minimal to prevent long
  * shutdown time.
  */
 public class ShutdownHookManager {
+  private static final Logger log = LoggerFactory.getLogger(ShutdownHookManager.class);
 
   private static ShutdownHookManager instance;
 
   /** The registered shutdown hooks */
-  private final List<Runnable> shutdownHooks;
+  private final Set<ShutdownHook> shutdownHooks;
+
+  private final AtomicBoolean isShutdown;
 
   private ShutdownHookManager() {
-    shutdownHooks = new LinkedList<>();
+    shutdownHooks = Collections.synchronizedSet(new HashSet<>());
+    isShutdown = new AtomicBoolean(false);
   }
 
   /**
@@ -30,21 +38,51 @@ public static ShutdownHookManager getInstance() {
     return instance;
   }
 
-  /** Adds the runnable to the shutdown hooks at the beginning of the list */
+  /**
+   * Adds the runnable to the shutdown hooks at the beginning of the set. During shutdown, no new
+   * hooks can be added.
+   */
   public void addShutdownHook(Runnable runnable) {
-    shutdownHooks.add(0, runnable);
+    if (!isShutdown.get()) {
+      ShutdownHook shutdownHook = new ShutdownHook(runnable, 0);
+      shutdownHooks.add(shutdownHook);
+    }
   }
 
-  /** Adds the runnable to the shutdown hooks at the end of the list */
+  /**
+   * Adds the runnable to the shutdown hooks at the end of the set. During shutdown, no new hooks
+   * can be added.
+   */
   public void addShutdownHookLast(Runnable runnable) {
-    shutdownHooks.add(shutdownHooks.size(), runnable);
+    if (!isShutdown.get()) {
+      ShutdownHook shutdownHook = new ShutdownHook(runnable, Integer.MAX_VALUE);
+      shutdownHooks.add(shutdownHook);
+    }
   }
 
   /** Sets up the registered shutdown hooks, to be executed at shutdown */
   private void setUpShutdownHooks() {
     Runtime.getRuntime()
-        .addShutdownHook(
-            new Thread(() -> shutdownHooks.forEach(Runnable::run), "inspectit-shutdown"));
+        .addShutdownHook(new Thread(this::executeShutdownHooks, "inspectit-shutdown"));
+  }
+
+  /** Executes all registered shutdown hooks by order */
+  @VisibleForTesting
+  void executeShutdownHooks() {
+    if (!isShutdown.compareAndSet(false, true)) log.info("Cannot execute shutdown hooks twice");
+    else
+      shutdownHooks.stream()
+          .sorted(Comparator.comparingInt(ShutdownHook::getOrder))
+          .forEach(this::tryRunShutdownHook);
+  }
+
+  /** Try-catch-wrapper to run a shutdown hook */
+  private void tryRunShutdownHook(ShutdownHook shutdownHook) {
+    try {
+      shutdownHook.run();
+    } catch (Exception e) {
+      log.error("Error while executing shutdown hook", e);
+    }
   }
 
   /**
@@ -57,7 +95,29 @@ public int getShutdownHookCount() {
   }
 
   /** Method for testing. */
-  public void clearShutdownHooks() {
+  public void reset() {
     shutdownHooks.clear();
+    isShutdown.set(false);
+  }
+
+  /** Internal class for ordered shutdown hooks */
+  private static class ShutdownHook {
+
+    private final Runnable shutdownHook;
+
+    private final int order;
+
+    ShutdownHook(Runnable shutdownHook, int order) {
+      this.shutdownHook = shutdownHook;
+      this.order = order;
+    }
+
+    void run() {
+      shutdownHook.run();
+    }
+
+    int getOrder() {
+      return order;
+    }
   }
 }
diff --git a/...ent/src/main/java/rocks/inspectit/gepard/agent/notification/http/NotificationFactory.java b/...ent/src/main/java/rocks/inspectit/gepard/agent/notification/http/NotificationFactory.java
@@ -32,7 +32,8 @@ private NotificationFactory() {}
    */
   public static SimpleHttpRequest createStartNotification(String baseUrl)
       throws URISyntaxException, JsonProcessingException {
-    URI uri = new URI(baseUrl + "/connections");
+    String agentId = AgentInfo.INFO.getAgentId();
+    URI uri = new URI(baseUrl + "/connections/" + agentId);
     String agentInfoString = mapper.writeValueAsString(AgentInfo.INFO);
 
     return SimpleRequestBuilder.post(uri)
@@ -55,7 +56,7 @@ public static SimpleHttpRequest createShutdownNotification(String baseUrl)
     URI uri = new URI(baseUrl + "/connections/" + agentId);
     String notificationBody = mapper.writeValueAsString(ShutdownNotification.INSTANCE);
 
-    return SimpleRequestBuilder.put(uri)
+    return SimpleRequestBuilder.patch(uri)
         .setBody(notificationBody, ContentType.APPLICATION_JSON)
         .setHeader("content-type", "application/json")
         .build();

diff --git a/...src/test/java/rocks/inspectit/gepard/agent/internal/shutdown/ShutdownHookManagerTest.java b/...src/test/java/rocks/inspectit/gepard/agent/internal/shutdown/ShutdownHookManagerTest.java
@@ -2,26 +2,51 @@
 package rocks.inspectit.gepard.agent.internal.shutdown;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.mockito.Mockito.*;
 
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.ExtendWith;
+import org.mockito.InOrder;
+import org.mockito.Mockito;
+import org.mockito.junit.jupiter.MockitoExtension;
 
+@ExtendWith(MockitoExtension.class)
 class ShutdownHookManagerTest {
 
   private final ShutdownHookManager manager = ShutdownHookManager.getInstance();
 
   @BeforeEach
   void beforeEach() {
-    manager.clearShutdownHooks();
+    manager.reset();
   }
 
   @Test
-  void shouldAddShutdownHook() {
-    Runnable shutdownHook = () -> System.out.println("shutdown");
+  void shouldRunShutdownHooksInOrder() {
+    Runnable runnable1 = mock(Runnable.class);
+    Runnable runnable2 = mock(Runnable.class);
+    manager.addShutdownHookLast(runnable1);
+    manager.addShutdownHook(runnable2);
 
-    manager.addShutdownHook(shutdownHook);
-    manager.addShutdownHookLast(shutdownHook);
+    manager.executeShutdownHooks();
 
+    InOrder inOrder = Mockito.inOrder(runnable1, runnable2);
+    inOrder.verify(runnable2).run();
+    inOrder.verify(runnable1).run();
     assertEquals(2, manager.getShutdownHookCount());
   }
+
+  @Test
+  void shouldRunShutdownHooksWithoutExceptions() {
+    Runnable runnable1 = mock(Runnable.class);
+    Runnable runnable2 = mock(Runnable.class);
+    manager.addShutdownHook(runnable1);
+    manager.addShutdownHookLast(runnable2);
+    doThrow(RuntimeException.class).when(runnable1).run();
+
+    manager.executeShutdownHooks();
+
+    verify(runnable1).run();
+    verify(runnable2).run();
+  }
 }
diff --git a/...gent/src/test/java/rocks/inspectit/gepard/agent/notification/NotificationManagerTest.java b/...gent/src/test/java/rocks/inspectit/gepard/agent/notification/NotificationManagerTest.java
@@ -11,6 +11,7 @@
 import org.junit.jupiter.api.Test;
 import org.mockserver.model.HttpError;
 import rocks.inspectit.gepard.agent.MockServerTestBase;
+import rocks.inspectit.gepard.agent.internal.identity.model.AgentInfo;
 import rocks.inspectit.gepard.agent.internal.shutdown.ShutdownHookManager;
 
 class NotificationManagerTest extends MockServerTestBase {
@@ -19,15 +20,17 @@ class NotificationManagerTest extends MockServerTestBase {
 
   private final ShutdownHookManager shutdownHookManager = ShutdownHookManager.getInstance();
 
+  private final String agentId = AgentInfo.INFO.getAgentId();
+
   @BeforeEach
   void beforeEach() {
-    shutdownHookManager.clearShutdownHooks();
+    shutdownHookManager.reset();
   }
 
   @Test
   void sendsStartNotificationIfServerUrlWasProvided() throws Exception {
     mockServer
-        .when(request().withMethod("POST").withPath("/api/v1/connections"))
+        .when(request().withMethod("POST").withPath("/api/v1/connections/" + agentId))
         .respond(response().withStatusCode(201));
 
     withEnvironmentVariable(SERVER_URL_ENV_PROPERTY, SERVER_URL)
@@ -44,7 +47,7 @@ void sendsStartNotificationIfServerUrlWasProvided() throws Exception {
   @Test
   void startNotificationFailsWithServerError() throws Exception {
     mockServer
-        .when(request().withMethod("POST").withPath("/api/v1/connections"))
+        .when(request().withMethod("POST").withPath("/api/v1/connections/" + agentId))
         .error(HttpError.error().withDropConnection(true));
 
     withEnvironmentVariable(SERVER_URL_ENV_PROPERTY, SERVER_URL)
@@ -61,7 +64,7 @@ void startNotificationFailsWithServerError() throws Exception {
   @Test
   void sendsNoStartNotificationWithoutProvidedServerUrl() {
     mockServer
-        .when(request().withMethod("POST").withPath("/api/v1/connections"))
+        .when(request().withMethod("POST").withPath("/api/v1/connections/" + agentId))
         .respond(response().withStatusCode(201));
 
     manager = NotificationManager.create();

diff --git a/...d-agent/src/test/java/rocks/inspectit/gepard/agent/notification/ShutdownNotifierTest.java b/...d-agent/src/test/java/rocks/inspectit/gepard/agent/notification/ShutdownNotifierTest.java
@@ -20,7 +20,7 @@ class ShutdownNotifierTest extends MockServerTestBase {
   @Test
   void notificationIsSentSuccessfully() {
     mockServer
-        .when(request().withMethod("PUT").withPath("/api/v1/connections/" + agentId))
+        .when(request().withMethod("PATCH").withPath("/api/v1/connections/" + agentId))
         .respond(response().withStatusCode(200));
 
     boolean successful = notifier.sendNotification(SERVER_URL);
@@ -31,7 +31,7 @@ void notificationIsSentSuccessfully() {
   @Test
   void serverIsNotAvailable() {
     mockServer
-        .when(request().withMethod("PUT").withPath("/api/v1/connections/" + agentId))
+        .when(request().withMethod("PATCH").withPath("/api/v1/connections/" + agentId))
         .respond(response().withStatusCode(503));
 
     boolean successful = notifier.sendNotification(SERVER_URL);
@@ -42,7 +42,7 @@ void serverIsNotAvailable() {
   @Test
   void serverReturnsError() {
     mockServer
-        .when(request().withMethod("PUT").withPath("/api/v1/connections/" + agentId))
+        .when(request().withMethod("PATCH").withPath("/api/v1/connections/" + agentId))
         .error(HttpError.error().withDropConnection(true));
 
     boolean successful = notifier.sendNotification(SERVER_URL);

diff --git a/...pard-agent/src/test/java/rocks/inspectit/gepard/agent/notification/StartNotifierTest.java b/...pard-agent/src/test/java/rocks/inspectit/gepard/agent/notification/StartNotifierTest.java
@@ -8,15 +8,18 @@
 import org.junit.jupiter.api.Test;
 import org.mockserver.model.HttpError;
 import rocks.inspectit.gepard.agent.MockServerTestBase;
+import rocks.inspectit.gepard.agent.internal.identity.model.AgentInfo;
 
 class StartNotifierTest extends MockServerTestBase {
 
   private final StartNotifier notifier = new StartNotifier();
 
+  private final String agentId = AgentInfo.INFO.getAgentId();
+
   @Test
   void notificationIsSentSuccessfully() {
     mockServer
-        .when(request().withMethod("POST").withPath("/api/v1/connections"))
+        .when(request().withMethod("POST").withPath("/api/v1/connections/" + agentId))
         .respond(response().withStatusCode(201));
 
     boolean successful = notifier.sendNotification(SERVER_URL);
@@ -27,7 +30,7 @@ void notificationIsSentSuccessfully() {
   @Test
   void serverIsNotAvailable() {
     mockServer
-        .when(request().withMethod("POST").withPath("/api/v1/connections"))
+        .when(request().withMethod("POST").withPath("/api/v1/connections/" + agentId))
         .respond(response().withStatusCode(503));
 
     boolean successful = notifier.sendNotification(SERVER_URL);
@@ -38,7 +41,7 @@ void serverIsNotAvailable() {
   @Test
   void serverReturnsError() {
     mockServer
-        .when(request().withMethod("POST").withPath("/api/v1/connections"))
+        .when(request().withMethod("POST").withPath("/api/v1/connections/" + agentId))
         .error(HttpError.error().withDropConnection(true));
 
     boolean successful = notifier.sendNotification(SERVER_URL);

diff --git a/...src/test/java/rocks/inspectit/gepard/agent/notification/http/NotificationFactoryTest.java b/...src/test/java/rocks/inspectit/gepard/agent/notification/http/NotificationFactoryTest.java
@@ -17,10 +17,12 @@ class NotificationFactoryTest {
   private static final ObjectMapper mapper =
       new ObjectMapper().setVisibility(PropertyAccessor.FIELD, JsonAutoDetect.Visibility.ANY);
 
+  private final String agentId = AgentInfo.INFO.getAgentId();
+
   @Test
   void validUrlCreatesStartNotification() throws Exception {
     String baseUrl = "http://localhost:8080";
-    String url = "http://localhost:8080/connections";
+    String url = "http://localhost:8080/connections/" + agentId;
     String contentType = "application/json";
     String info = mapper.writeValueAsString(AgentInfo.INFO);
 
@@ -40,7 +42,6 @@ void invalidStartNotificationUrlThrowsException() {
 
   @Test
   void validUrlCreatesShutdownNotification() throws Exception {
-    String agentId = AgentInfo.INFO.getAgentId();
     String baseUrl = "http://localhost:8080";
     String url = "http://localhost:8080/connections/" + agentId;
     String contentType = "application/json";