Fix autoconfigure shutdown hook (#5221)

* Fix autoconfigure shutdown hook * Add unit test
open-telemetry · Feb 15, 2023 · 219729a · 219729a
1 parent 84feb47
commit 219729a
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 4 deletions.
diff --git a/...c/main/java/io/opentelemetry/sdk/autoconfigure/AutoConfiguredOpenTelemetrySdkBuilder.java b/...c/main/java/io/opentelemetry/sdk/autoconfigure/AutoConfiguredOpenTelemetrySdkBuilder.java
@@ -378,10 +378,6 @@ public AutoConfiguredOpenTelemetrySdk build() {
         SdkLoggerProvider loggerProvider = loggerProviderBuilder.build();
         closeables.add(loggerProvider);
 
-        if (registerShutdownHook) {
-          Runtime.getRuntime().addShutdownHook(new Thread(openTelemetrySdk::close));
-        }
-
         ContextPropagators propagators =
             PropagatorConfiguration.configurePropagators(
                 config, serviceClassLoader, propagatorCustomizer);
@@ -396,6 +392,11 @@ public AutoConfiguredOpenTelemetrySdk build() {
         openTelemetrySdk = sdkBuilder.build();
       }
 
+      // NOTE: Shutdown hook registration is untested. Modify with caution.
+      if (registerShutdownHook) {
+        Runtime.getRuntime().addShutdownHook(shutdownHook(openTelemetrySdk));
+      }
+
       if (setResultAsGlobal) {
         GlobalOpenTelemetry.set(openTelemetrySdk);
         GlobalLoggerProvider.set(openTelemetrySdk.getSdkLoggerProvider());
@@ -456,6 +457,11 @@ private ConfigProperties computeConfigProperties() {
     return properties;
   }
 
+  // Visible for testing
+  Thread shutdownHook(OpenTelemetrySdk sdk) {
+    return new Thread(sdk::close);
+  }
+
   private static <I, O1, O2> BiFunction<I, ConfigProperties, O2> mergeCustomizer(
       BiFunction<? super I, ConfigProperties, ? extends O1> first,
       BiFunction<? super O1, ConfigProperties, ? extends O2> second) {

diff --git a/.../src/test/java/io/opentelemetry/sdk/autoconfigure/AutoConfiguredOpenTelemetrySdkTest.java b/.../src/test/java/io/opentelemetry/sdk/autoconfigure/AutoConfiguredOpenTelemetrySdkTest.java
@@ -10,10 +10,12 @@
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.assertThatThrownBy;
 import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.doThrow;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
@@ -374,6 +376,29 @@ void builder_setResultAsGlobalTrue() {
         .isSameAs(openTelemetry.getSdkLoggerProvider());
   }
 
+  @Test
+  void builder_registersShutdownHook() {
+    builder = spy(builder);
+    Thread thread = new Thread();
+    doReturn(thread).when(builder).shutdownHook(any());
+
+    OpenTelemetrySdk sdk = builder.setResultAsGlobal(false).build().getOpenTelemetrySdk();
+
+    verify(builder, times(1)).shutdownHook(sdk);
+    assertThat(Runtime.getRuntime().removeShutdownHook(thread)).isTrue();
+  }
+
+  @Test
+  void shutdownHook() throws InterruptedException {
+    OpenTelemetrySdk sdk = mock(OpenTelemetrySdk.class);
+
+    Thread thread = builder.shutdownHook(sdk);
+    thread.start();
+    thread.join();
+
+    verify(sdk).close();
+  }
+
   private static Supplier<Map<String, String>> disableExportPropertySupplier() {
     Map<String, String> props = new HashMap<>();
     props.put("otel.metrics.exporter", "none");