From 2140d3e6db7aff9bfcea70bb5608921b1feb7042 Mon Sep 17 00:00:00 2001 From: Nikita Tkachenko Date: Mon, 7 Apr 2025 15:23:16 +0200 Subject: [PATCH 1/2] Shutdown CI Visibility test event handlers before tracer --- .../civisibility/CiVisibilitySystem.java | 34 ++++++++++++++----- .../junit4/TestEventsHandlerHolder.java | 18 +--------- .../junit5/TestEventsHandlerHolder.java | 11 +----- .../karate/TestEventsHandlerHolder.java | 8 +---- .../testng/TestEventsHandlerHolder.java | 11 +----- .../weaver/DatadogWeaverReporter.java | 11 +----- .../trace/api/internal/InternalTracer.java | 2 ++ .../java/datadog/trace/core/CoreTracer.java | 14 ++++++++ .../java/datadog/opentracing/DDTracer.java | 5 +++ .../instrumentation/api/AgentTracer.java | 3 ++ .../trace/util/AgentThreadFactory.java | 1 - 11 files changed, 54 insertions(+), 64 deletions(-) diff --git a/dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/CiVisibilitySystem.java b/dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/CiVisibilitySystem.java index 5bc7e85c587..575e9a38442 100644 --- a/dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/CiVisibilitySystem.java +++ b/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> handlers = new CopyOnWriteArrayList<>(); + private TestEventsHandlerFactory( CiVisibilityServices services, CiVisibilityRepoServices repoServices, @@ -174,12 +181,21 @@ public TestEventsHandler 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 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(); + } } } diff --git a/dd-java-agent/instrumentation/junit-4.10/src/main/java/datadog/trace/instrumentation/junit4/TestEventsHandlerHolder.java b/dd-java-agent/instrumentation/junit-4.10/src/main/java/datadog/trace/instrumentation/junit4/TestEventsHandlerHolder.java index fa3eacdf225..9346862ff80 100644 --- a/dd-java-agent/instrumentation/junit-4.10/src/main/java/datadog/trace/instrumentation/junit4/TestEventsHandlerHolder.java +++ b/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> 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 capabilities) { TestEventsHandler handler = HANDLERS.get(framework); @@ -38,13 +28,7 @@ public static synchronized void start( } } - public static synchronized void stop() { - for (TestEventsHandler handler : HANDLERS.values()) { - handler.close(); - } - HANDLERS.clear(); - } - + /** Used by instrumentation tests */ public static synchronized void stop(TestFrameworkInstrumentation framework) { TestEventsHandler handler = HANDLERS.remove(framework); if (handler != null) { diff --git a/dd-java-agent/instrumentation/junit-5.3/src/main/java/datadog/trace/instrumentation/junit5/TestEventsHandlerHolder.java b/dd-java-agent/instrumentation/junit-5.3/src/main/java/datadog/trace/instrumentation/junit5/TestEventsHandlerHolder.java index a53e419e020..f7c05d59ca9 100644 --- a/dd-java-agent/instrumentation/junit-5.3/src/main/java/datadog/trace/instrumentation/junit5/TestEventsHandlerHolder.java +++ b/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 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 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 handler : HANDLERS.values()) { handler.close(); diff --git a/dd-java-agent/instrumentation/karate/src/main/java/datadog/trace/instrumentation/karate/TestEventsHandlerHolder.java b/dd-java-agent/instrumentation/karate/src/main/java/datadog/trace/instrumentation/karate/TestEventsHandlerHolder.java index 40a3ad8195d..15895b8d721 100644 --- a/dd-java-agent/instrumentation/karate/src/main/java/datadog/trace/instrumentation/karate/TestEventsHandlerHolder.java +++ b/dd-java-agent/instrumentation/karate/src/main/java/datadog/trace/instrumentation/karate/TestEventsHandlerHolder.java @@ -4,7 +4,6 @@ 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 { @@ -12,12 +11,6 @@ public abstract class TestEventsHandlerHolder { 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(); diff --git a/dd-java-agent/instrumentation/testng/src/main/java/datadog/trace/instrumentation/testng/TestEventsHandlerHolder.java b/dd-java-agent/instrumentation/testng/src/main/java/datadog/trace/instrumentation/testng/TestEventsHandlerHolder.java index d8dc16d391f..6edc08fed73 100644 --- a/dd-java-agent/instrumentation/testng/src/main/java/datadog/trace/instrumentation/testng/TestEventsHandlerHolder.java +++ b/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 TEST_STORE; - static { - Runtime.getRuntime() - .addShutdownHook( - AgentThreadFactory.newAgentThread( - AgentThreadFactory.AgentThread.CI_TEST_EVENTS_SHUTDOWN_HOOK, - TestEventsHandlerHolder::stop, - false)); - } - public static synchronized void setContextStore(ContextStore 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(); diff --git a/dd-java-agent/instrumentation/weaver/src/main/java/datadog/trace/instrumentation/weaver/DatadogWeaverReporter.java b/dd-java-agent/instrumentation/weaver/src/main/java/datadog/trace/instrumentation/weaver/DatadogWeaverReporter.java index 7ba5d31906e..4c5c0f7be63 100644 --- a/dd-java-agent/instrumentation/weaver/src/main/java/datadog/trace/instrumentation/weaver/DatadogWeaverReporter.java +++ b/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 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(); diff --git a/dd-trace-api/src/main/java/datadog/trace/api/internal/InternalTracer.java b/dd-trace-api/src/main/java/datadog/trace/api/internal/InternalTracer.java index 73a825ea329..eab7132d46e 100644 --- a/dd-trace-api/src/main/java/datadog/trace/api/internal/InternalTracer.java +++ b/dd-trace-api/src/main/java/datadog/trace/api/internal/InternalTracer.java @@ -30,4 +30,6 @@ public interface InternalTracer { * @return DataStreamsCheckpointer instance. */ DataStreamsCheckpointer getDataStreamsCheckpointer(); + + void addShutdownListener(Runnable listener); } diff --git a/dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java b/dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java index d629b1820ab..f3a718a5a3e 100644 --- a/dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java +++ b/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 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(); } diff --git a/dd-trace-ot/src/main/java/datadog/opentracing/DDTracer.java b/dd-trace-ot/src/main/java/datadog/opentracing/DDTracer.java index fdea6629942..54c5764f19a 100644 --- a/dd-trace-ot/src/main/java/datadog/opentracing/DDTracer.java +++ b/dd-trace-ot/src/main/java/datadog/opentracing/DDTracer.java @@ -461,6 +461,11 @@ public DataStreamsCheckpointer getDataStreamsCheckpointer() { return tracer.getDataStreamsCheckpointer(); } + @Override + public void addShutdownListener(Runnable listener) { + tracer.addShutdownListener(listener); + } + @Override public ScopeManager scopeManager() { return scopeManager; diff --git a/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentTracer.java b/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentTracer.java index 0ea77358500..e2fdbe497e9 100644 --- a/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentTracer.java +++ b/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentTracer.java @@ -599,6 +599,9 @@ public DataStreamsCheckpointer getDataStreamsCheckpointer() { return getDataStreamsMonitoring(); } + @Override + public void addShutdownListener(Runnable listener) {} + @Override public void addScopeListener(final ScopeListener listener) {} diff --git a/internal-api/src/main/java/datadog/trace/util/AgentThreadFactory.java b/internal-api/src/main/java/datadog/trace/util/AgentThreadFactory.java index 304fb6ca9ec..67bdf576607 100644 --- a/internal-api/src/main/java/datadog/trace/util/AgentThreadFactory.java +++ b/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"), From a8881562cbd6388fabdc034f3d9196665d9b21f7 Mon Sep 17 00:00:00 2001 From: Nikita Tkachenko Date: Mon, 7 Apr 2025 17:07:56 +0200 Subject: [PATCH 2/2] Remove addShutdownListener from OT tracer --- .../main/java/datadog/trace/api/internal/InternalTracer.java | 2 -- dd-trace-ot/src/main/java/datadog/opentracing/DDTracer.java | 5 ----- .../trace/bootstrap/instrumentation/api/AgentTracer.java | 2 ++ 3 files changed, 2 insertions(+), 7 deletions(-) diff --git a/dd-trace-api/src/main/java/datadog/trace/api/internal/InternalTracer.java b/dd-trace-api/src/main/java/datadog/trace/api/internal/InternalTracer.java index eab7132d46e..73a825ea329 100644 --- a/dd-trace-api/src/main/java/datadog/trace/api/internal/InternalTracer.java +++ b/dd-trace-api/src/main/java/datadog/trace/api/internal/InternalTracer.java @@ -30,6 +30,4 @@ public interface InternalTracer { * @return DataStreamsCheckpointer instance. */ DataStreamsCheckpointer getDataStreamsCheckpointer(); - - void addShutdownListener(Runnable listener); } diff --git a/dd-trace-ot/src/main/java/datadog/opentracing/DDTracer.java b/dd-trace-ot/src/main/java/datadog/opentracing/DDTracer.java index 54c5764f19a..fdea6629942 100644 --- a/dd-trace-ot/src/main/java/datadog/opentracing/DDTracer.java +++ b/dd-trace-ot/src/main/java/datadog/opentracing/DDTracer.java @@ -461,11 +461,6 @@ public DataStreamsCheckpointer getDataStreamsCheckpointer() { return tracer.getDataStreamsCheckpointer(); } - @Override - public void addShutdownListener(Runnable listener) { - tracer.addShutdownListener(listener); - } - @Override public ScopeManager scopeManager() { return scopeManager; diff --git a/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentTracer.java b/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentTracer.java index e2fdbe497e9..96d44360d91 100644 --- a/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentTracer.java +++ b/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 {