Shutdown CI Visibility test event handlers before tracer (#8677)

Author: nikita-tkachenko-datadog

dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/CiVisibilitySystem.java

@@ -15,6 +15,7 @@
 import datadog.trace.api.civisibility.telemetry.NoOpMetricCollector;
 import datadog.trace.api.git.GitInfoProvider;
 import datadog.trace.bootstrap.ContextStore;
+import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
 import datadog.trace.civisibility.config.ExecutionSettings;
 import datadog.trace.civisibility.config.JvmInfo;
 import datadog.trace.civisibility.coverage.file.instrumentation.CoverageClassTransformer;
@@ -41,6 +42,7 @@
 import java.nio.file.Path;
 import java.nio.file.Paths;
 import java.util.Collection;
+import java.util.concurrent.CopyOnWriteArrayList;
 import java.util.function.Predicate;
 import javax.annotation.Nullable;
 import org.slf4j.Logger;
@@ -101,10 +103,13 @@ public static void start(Instrumentation inst, SharedCommunicationObjects sco) {
 
       CiVisibilityCoverageServices.Child coverageServices =
           new CiVisibilityCoverageServices.Child(services, repoServices, executionSettings);
-      InstrumentationBridge.registerTestEventsHandlerFactory(
-          new TestEventsHandlerFactory(
-              services, repoServices, coverageServices, executionSettings));
+      TestEventsHandlerFactory testEventsHandlerFactory =
+          new TestEventsHandlerFactory(services, repoServices, coverageServices, executionSettings);
+      InstrumentationBridge.registerTestEventsHandlerFactory(testEventsHandlerFactory);
       CoveragePerTestBridge.registerCoverageStoreRegistry(coverageServices.coverageStoreFactory);
+
+      AgentTracer.TracerAPI tracerAPI = AgentTracer.get();
+      tracerAPI.addShutdownListener(testEventsHandlerFactory::shutdown);
     } else {
       InstrumentationBridge.registerTestEventsHandlerFactory(new NoOpTestEventsHandler.Factory());
     }
@@ -147,6 +152,8 @@ private static final class TestEventsHandlerFactory implements TestEventsHandler
     private final CiVisibilityRepoServices repoServices;
     private final TestFrameworkSession.Factory sessionFactory;
 
+    private final Collection<TestEventsHandler<?, ?>> handlers = new CopyOnWriteArrayList<>();
+
     private TestEventsHandlerFactory(
         CiVisibilityServices services,
         CiVisibilityRepoServices repoServices,
@@ -174,12 +181,21 @@ public <SuiteKey, TestKey> TestEventsHandler<SuiteKey, TestKey> create(
       TestFrameworkSession testSession =
           sessionFactory.startSession(repoServices.moduleName, component, null, capabilities);
       TestFrameworkModule testModule = testSession.testModuleStart(repoServices.moduleName, null);
-      return new TestEventsHandlerImpl<>(
-          services.metricCollector,
-          testSession,
-          testModule,
-          suiteStore != null ? suiteStore : new ConcurrentHashMapContextStore<>(),
-          testStore != null ? testStore : new ConcurrentHashMapContextStore<>());
+      TestEventsHandlerImpl<SuiteKey, TestKey> handler =
+          new TestEventsHandlerImpl<>(
+              services.metricCollector,
+              testSession,
+              testModule,
+              suiteStore != null ? suiteStore : new ConcurrentHashMapContextStore<>(),
+              testStore != null ? testStore : new ConcurrentHashMapContextStore<>());
+      handlers.add(handler);
+      return handler;
+    }
+
+    public void shutdown() {
+      for (TestEventsHandler<?, ?> handler : handlers) {
+        handler.close();
+      }
     }
   }
 

dd-java-agent/instrumentation/junit-4.10/src/main/java/datadog/trace/instrumentation/junit4/TestEventsHandlerHolder.java

@@ -6,7 +6,6 @@
 import datadog.trace.api.civisibility.events.TestEventsHandler;
 import datadog.trace.api.civisibility.events.TestSuiteDescriptor;
 import datadog.trace.api.civisibility.telemetry.tag.TestFrameworkInstrumentation;
-import datadog.trace.util.AgentThreadFactory;
 import datadog.trace.util.ConcurrentEnumMap;
 import java.util.Collection;
 import java.util.Map;
@@ -18,15 +17,6 @@ public abstract class TestEventsHandlerHolder {
           TestFrameworkInstrumentation, TestEventsHandler<TestSuiteDescriptor, TestDescriptor>>
       HANDLERS = new ConcurrentEnumMap<>(TestFrameworkInstrumentation.class);
 
-  static {
-    Runtime.getRuntime()
-        .addShutdownHook(
-            AgentThreadFactory.newAgentThread(
-                AgentThreadFactory.AgentThread.CI_TEST_EVENTS_SHUTDOWN_HOOK,
-                TestEventsHandlerHolder::stop,
-                false));
-  }
-
   public static synchronized void start(
       TestFrameworkInstrumentation framework, Collection<LibraryCapability> capabilities) {
     TestEventsHandler<TestSuiteDescriptor, TestDescriptor> handler = HANDLERS.get(framework);
@@ -38,13 +28,7 @@ public static synchronized void start(
     }
   }
 
-  public static synchronized void stop() {
-    for (TestEventsHandler<TestSuiteDescriptor, TestDescriptor> handler : HANDLERS.values()) {
-      handler.close();
-    }
-    HANDLERS.clear();
-  }
-
+  /** Used by instrumentation tests */
   public static synchronized void stop(TestFrameworkInstrumentation framework) {
     TestEventsHandler<TestSuiteDescriptor, TestDescriptor> handler = HANDLERS.remove(framework);
     if (handler != null) {

dd-java-agent/instrumentation/junit-5.3/src/main/java/datadog/trace/instrumentation/junit5/TestEventsHandlerHolder.java

@@ -7,7 +7,6 @@
 import datadog.trace.api.civisibility.execution.TestExecutionHistory;
 import datadog.trace.api.civisibility.telemetry.tag.TestFrameworkInstrumentation;
 import datadog.trace.bootstrap.ContextStore;
-import datadog.trace.util.AgentThreadFactory;
 import datadog.trace.util.ConcurrentEnumMap;
 import java.util.Map;
 import org.junit.platform.engine.TestDescriptor;
@@ -23,15 +22,6 @@ public abstract class TestEventsHandlerHolder {
   private static volatile ContextStore<TestDescriptor, TestExecutionHistory>
       EXECUTION_HISTORY_STORE;
 
-  static {
-    Runtime.getRuntime()
-        .addShutdownHook(
-            AgentThreadFactory.newAgentThread(
-                AgentThreadFactory.AgentThread.CI_TEST_EVENTS_SHUTDOWN_HOOK,
-                TestEventsHandlerHolder::stop,
-                false));
-  }
-
   public static synchronized void setExecutionHistoryStore(
       ContextStore<TestDescriptor, TestExecutionHistory> executionHistoryStore) {
     if (EXECUTION_HISTORY_STORE == null) {
@@ -71,6 +61,7 @@ public static synchronized void start(
     }
   }
 
+  /** Used by instrumentation tests */
   public static synchronized void stop() {
     for (TestEventsHandler<TestDescriptor, TestDescriptor> handler : HANDLERS.values()) {
       handler.close();

dd-java-agent/instrumentation/karate/src/main/java/datadog/trace/instrumentation/karate/TestEventsHandlerHolder.java

@@ -4,20 +4,13 @@
 import datadog.trace.api.civisibility.events.TestDescriptor;
 import datadog.trace.api.civisibility.events.TestEventsHandler;
 import datadog.trace.api.civisibility.events.TestSuiteDescriptor;
-import datadog.trace.util.AgentThreadFactory;
 
 public abstract class TestEventsHandlerHolder {
 
   public static volatile TestEventsHandler<TestSuiteDescriptor, TestDescriptor> TEST_EVENTS_HANDLER;
 
   static {
     start();
-    Runtime.getRuntime()
-        .addShutdownHook(
-            AgentThreadFactory.newAgentThread(
-                AgentThreadFactory.AgentThread.CI_TEST_EVENTS_SHUTDOWN_HOOK,
-                TestEventsHandlerHolder::stop,
-                false));
   }
 
   public static void start() {
@@ -26,6 +19,7 @@ public static void start() {
             "karate", null, null, KarateUtils.capabilities(KarateTracingHook.FRAMEWORK_VERSION));
   }
 
+  /** Used by instrumentation tests */
   public static void stop() {
     if (TEST_EVENTS_HANDLER != null) {
       TEST_EVENTS_HANDLER.close();

dd-java-agent/instrumentation/testng/src/main/java/datadog/trace/instrumentation/testng/TestEventsHandlerHolder.java

@@ -5,7 +5,6 @@
 import datadog.trace.api.civisibility.events.TestEventsHandler;
 import datadog.trace.api.civisibility.events.TestSuiteDescriptor;
 import datadog.trace.bootstrap.ContextStore;
-import datadog.trace.util.AgentThreadFactory;
 import org.testng.ITestResult;
 
 public abstract class TestEventsHandlerHolder {
@@ -14,15 +13,6 @@ public abstract class TestEventsHandlerHolder {
 
   private static ContextStore<ITestResult, DDTest> TEST_STORE;
 
-  static {
-    Runtime.getRuntime()
-        .addShutdownHook(
-            AgentThreadFactory.newAgentThread(
-                AgentThreadFactory.AgentThread.CI_TEST_EVENTS_SHUTDOWN_HOOK,
-                TestEventsHandlerHolder::stop,
-                false));
-  }
-
   public static synchronized void setContextStore(ContextStore<ITestResult, DDTest> testStore) {
     if (TEST_STORE == null) {
       TEST_STORE = testStore;
@@ -35,6 +25,7 @@ public static void start() {
             "testng", null, TEST_STORE, TestNGUtils.capabilities(TestNGUtils.getTestNGVersion()));
   }
 
+  /** Used by instrumentation tests */
   public static void stop() {
     if (TEST_EVENTS_HANDLER != null) {
       TEST_EVENTS_HANDLER.close();

dd-java-agent/instrumentation/weaver/src/main/java/datadog/trace/instrumentation/weaver/DatadogWeaverReporter.java

@@ -8,7 +8,6 @@
 import datadog.trace.api.civisibility.execution.TestExecutionHistory;
 import datadog.trace.api.civisibility.telemetry.tag.TestFrameworkInstrumentation;
 import datadog.trace.api.time.SystemTimeSource;
-import datadog.trace.util.AgentThreadFactory;
 import java.lang.reflect.Method;
 import java.util.Collection;
 import java.util.Collections;
@@ -27,15 +26,6 @@ public class DatadogWeaverReporter {
   private static volatile TestEventsHandler<TestSuiteDescriptor, TestDescriptor>
       TEST_EVENTS_HANDLER;
 
-  static {
-    Runtime.getRuntime()
-        .addShutdownHook(
-            AgentThreadFactory.newAgentThread(
-                AgentThreadFactory.AgentThread.CI_TEST_EVENTS_SHUTDOWN_HOOK,
-                DatadogWeaverReporter::stop,
-                false));
-  }
-
   public static synchronized void start() {
     if (TEST_EVENTS_HANDLER == null) {
       TEST_EVENTS_HANDLER =
@@ -44,6 +34,7 @@ public static synchronized void start() {
     }
   }
 
+  /** Used by instrumentation tests */
   public static synchronized void stop() {
     if (TEST_EVENTS_HANDLER != null) {
       TEST_EVENTS_HANDLER.close();

dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java

@@ -115,6 +115,7 @@
 import java.util.ServiceLoader;
 import java.util.SortedSet;
 import java.util.concurrent.ConcurrentSkipListSet;
+import java.util.concurrent.CopyOnWriteArrayList;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.TimeoutException;
 import java.util.zip.ZipOutputStream;
@@ -210,6 +211,7 @@ public static CoreTracerBuilder builder() {
   private final ProfilingContextIntegration profilingContextIntegration;
   private final boolean injectBaggageAsTags;
   private final boolean flushOnClose;
+  private final Collection<Runnable> shutdownListeners = new CopyOnWriteArrayList<>();
 
   /**
    * JVM shutdown callback, keeping a reference to it to remove this if DDTracer gets destroyed
@@ -1124,6 +1126,11 @@ public DataStreamsCheckpointer getDataStreamsCheckpointer() {
     return this.dataStreamsMonitoring;
   }
 
+  @Override
+  public void addShutdownListener(Runnable listener) {
+    this.shutdownListeners.add(listener);
+  }
+
   @Override
   public void addScopeListener(final ScopeListener listener) {
     this.scopeManager.addScopeListener(listener);
@@ -1152,6 +1159,13 @@ public CallbackProvider getUniversalCallbackProvider() {
 
   @Override
   public void close() {
+    for (Runnable shutdownListener : shutdownListeners) {
+      try {
+        shutdownListener.run();
+      } catch (Exception e) {
+        log.error("Error while running shutdown listener", e);
+      }
+    }
     if (flushOnClose) {
       flush();
     }

internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentTracer.java

@@ -412,6 +412,8 @@ default SpanBuilder buildSpan(CharSequence spanName) {
      * @param serviceName The service name to use as default.
      */
     void updatePreferredServiceName(String serviceName);
+
+    void addShutdownListener(Runnable listener);
   }
 
   public interface SpanBuilder {
@@ -599,6 +601,9 @@ public DataStreamsCheckpointer getDataStreamsCheckpointer() {
       return getDataStreamsMonitoring();
     }
 
+    @Override
+    public void addShutdownListener(Runnable listener) {}
+
     @Override
     public void addScopeListener(final ScopeListener listener) {}
 

internal-api/src/main/java/datadog/trace/util/AgentThreadFactory.java

@@ -52,7 +52,6 @@ public enum AgentThread {
     CI_SHELL_COMMAND("dd-ci-shell-command"),
     CI_GIT_DATA_UPLOADER("dd-ci-git-data-uploader"),
     CI_GIT_DATA_SHUTDOWN_HOOK("dd-ci-git-data-shutdown-hook"),
-    CI_TEST_EVENTS_SHUTDOWN_HOOK("dd-ci-test-events-shutdown-hook"),
     CI_PROJECT_CONFIGURATOR("dd-ci-project-configurator"),
     CI_SIGNAL_SERVER("dd-ci-signal-server"),