Allow delayed start of adaptive sampler reconfiguration task. (#6134)

A preparation step towards allowing profiler in native image.
DataDog · Nov 2, 2023 · 20a70a4 · 20a70a4
1 parent 3514cd6
commit 20a70a4
Show file tree

Hide file tree

Showing 9 changed files with 68 additions and 25 deletions.
diff --git a/...-bootstrap/src/main/java11/datadog/trace/bootstrap/instrumentation/jfr/WindowSampler.java b/...-bootstrap/src/main/java11/datadog/trace/bootstrap/instrumentation/jfr/WindowSampler.java
@@ -13,10 +13,14 @@ public class WindowSampler<E extends Event> {
 
   protected WindowSampler(
       Duration windowDuration, int samplesPerWindow, int lookback, Class<E> eventType) {
-    sampler = new AdaptiveSampler(windowDuration, samplesPerWindow, lookback, 16);
+    sampler = new AdaptiveSampler(windowDuration, samplesPerWindow, lookback, 16, false);
     sampleType = EventType.getEventType(eventType);
   }
 
+  public void start() {
+    sampler.start();
+  }
+
   public boolean sample() {
     return sampleType.isEnabled() && sampler.sample();
   }

diff --git a/...ain/java11/datadog/trace/bootstrap/instrumentation/jfr/exceptions/ExceptionProfiling.java b/...ain/java11/datadog/trace/bootstrap/instrumentation/jfr/exceptions/ExceptionProfiling.java
@@ -42,6 +42,10 @@ private ExceptionProfiling(final Config config) {
     this.recordExceptionMessage = recordExceptionMessage;
   }
 
+  public void start() {
+    sampler.start();
+  }
+
   public ExceptionSampleEvent process(final Throwable t) {
     // always record the exception in histogram
     final boolean firstHit = histogram.record(t);

diff --git a/...r/debugger-bootstrap/src/main/java/datadog/trace/bootstrap/debugger/ProbeRateLimiter.java b/...r/debugger-bootstrap/src/main/java/datadog/trace/bootstrap/debugger/ProbeRateLimiter.java
@@ -78,9 +78,9 @@ private static Sampler createSampler(double rate) {
     }
     if (rate < 1) {
       int intRate = (int) Math.round(rate * 10);
-      return new AdaptiveSampler(TEN_SECONDS_WINDOW, intRate, 180, 16);
+      return new AdaptiveSampler(TEN_SECONDS_WINDOW, intRate, 180, 16, true);
     }
-    return new AdaptiveSampler(ONE_SECOND_WINDOW, (int) Math.round(rate), 180, 16);
+    return new AdaptiveSampler(ONE_SECOND_WINDOW, (int) Math.round(rate), 180, 16, true);
   }
 
   private static class RateLimitInfo {

diff --git a/dd-java-agent/agent-profiling/profiling-controller-openjdk/build.gradle b/dd-java-agent/agent-profiling/profiling-controller-openjdk/build.gradle
@@ -17,6 +17,9 @@ apply plugin: 'idea'
 dependencies {
   api deps.slf4j
   api project(':internal-api')
+  api(project(':dd-java-agent:agent-bootstrap')) {
+    exclude group: 'com.datadoghq', module: 'agent-logging'
+  }
   api project(':dd-java-agent:agent-profiling:profiling-auxiliary')
   api project(':dd-java-agent:agent-profiling:profiling-controller')
   api project(':dd-java-agent:agent-profiling:profiling-controller-jfr')

diff --git a/...ler-openjdk/src/main/java/com/datadog/profiling/controller/openjdk/OpenJdkController.java b/...ler-openjdk/src/main/java/com/datadog/profiling/controller/openjdk/OpenJdkController.java
@@ -30,6 +30,7 @@
 import datadog.trace.api.Platform;
 import datadog.trace.api.config.ProfilingConfig;
 import datadog.trace.bootstrap.config.provider.ConfigProvider;
+import datadog.trace.bootstrap.instrumentation.jfr.exceptions.ExceptionProfiling;
 import de.thetaphi.forbiddenapis.SuppressForbidden;
 import java.io.IOException;
 import java.time.Duration;
@@ -182,6 +183,10 @@ && isEventEnabled(recordingSettings, "jdk.NativeMethodSample")) {
 
     this.recordingSettings = Collections.unmodifiableMap(recordingSettings);
 
+    if (isEventEnabled(this.recordingSettings, "datadog.ExceptionSample")) {
+      ExceptionProfiling.getInstance().start();
+    }
+
     // Register periodic events
     AvailableProcessorCoresEvent.register();
   }

diff --git a/...profiling/src/main/java11/datadog/exceptions/instrumentation/ThrowableInstanceAdvice.java b/...profiling/src/main/java11/datadog/exceptions/instrumentation/ThrowableInstanceAdvice.java
@@ -26,16 +26,16 @@ public static void onExit(@Advice.This final Object t) {
     }
     try {
       /*
-       * Exclude internal agent threads from exception profiling.
+       * We may get into a situation when this is called before exception sampling is active.
        */
-      if (Config.get().isProfilingExcludeAgentThreads()
-          && AGENT_THREAD_GROUP.equals(Thread.currentThread().getThreadGroup())) {
+      if (!InstrumentationBasedProfiling.isJFRReady()) {
         return;
       }
       /*
-       * We may get into a situation when this is called before exception sampling is active.
+       * Exclude internal agent threads from exception profiling.
        */
-      if (!InstrumentationBasedProfiling.isJFRReady()) {
+      if (Config.get().isProfilingExcludeAgentThreads()
+          && AGENT_THREAD_GROUP.equals(Thread.currentThread().getThreadGroup())) {
         return;
       }
       /*

diff --git a/internal-api/src/jmh/java/datadog/trace/api/sampling/AdaptiveSamplerBenchmark.java b/internal-api/src/jmh/java/datadog/trace/api/sampling/AdaptiveSamplerBenchmark.java
@@ -44,7 +44,8 @@ public void setup() {
             Duration.of(durationWindowMillis, ChronoUnit.MILLIS),
             samplesPerWindow,
             averageLookback,
-            BUDGET_LOOKBACK);
+            BUDGET_LOOKBACK,
+            true);
   }
 
   @Threads(4)

diff --git a/internal-api/src/main/java/datadog/trace/api/sampling/AdaptiveSampler.java b/internal-api/src/main/java/datadog/trace/api/sampling/AdaptiveSampler.java
@@ -106,6 +106,9 @@ void onWindowRoll(
 
   private final ConfigListener listener;
 
+  private final Duration windowDuration;
+  private final AgentTaskScheduler taskScheduler;
+
   /**
    * Create a new sampler instance
    *
@@ -122,7 +125,8 @@ protected AdaptiveSampler(
       final int averageLookback,
       final int budgetLookback,
       final @Nullable ConfigListener listener,
-      final AgentTaskScheduler taskScheduler) {
+      final AgentTaskScheduler taskScheduler,
+      boolean startSampler) {
 
     if (averageLookback < 1) {
       throw new IllegalArgumentException("'averageLookback' argument must be at least 1");
@@ -141,12 +145,12 @@ protected AdaptiveSampler(
       listener.onWindowRoll(0, 0, samplesBudget, totalCountRunningAverage, probability);
     }
 
-    taskScheduler.weakScheduleAtFixedRate(
-        RollWindowTask.INSTANCE,
-        this,
-        windowDuration.toNanos(),
-        windowDuration.toNanos(),
-        TimeUnit.NANOSECONDS);
+    this.windowDuration = windowDuration;
+    this.taskScheduler = taskScheduler;
+
+    if (startSampler) {
+      start();
+    }
   }
 
   /**
@@ -161,12 +165,21 @@ public AdaptiveSampler(
       final Duration windowDuration,
       final int samplesPerWindow,
       final int averageLookback,
-      final int budgetLookback) {
-    this(windowDuration, samplesPerWindow, averageLookback, budgetLookback, null);
+      final int budgetLookback,
+      boolean startSampler) {
+    this(
+        windowDuration,
+        samplesPerWindow,
+        averageLookback,
+        budgetLookback,
+        null,
+        AgentTaskScheduler.INSTANCE,
+        startSampler);
   }
 
   /**
-   * Create a new sampler instance with automatic window roll.
+   * Create a new sampler instance with automatic window roll. The instance is automatically
+   * started.
    *
    * @param windowDuration the sampling window duration
    * @param samplesPerWindow the maximum number of samples in the sampling window
@@ -186,7 +199,17 @@ public AdaptiveSampler(
         averageLookback,
         budgetLookback,
         listener,
-        AgentTaskScheduler.INSTANCE);
+        AgentTaskScheduler.INSTANCE,
+        true);
+  }
+
+  public void start() {
+    taskScheduler.weakScheduleAtFixedRate(
+        RollWindowTask.INSTANCE,
+        this,
+        windowDuration.toNanos(),
+        windowDuration.toNanos(),
+        TimeUnit.NANOSECONDS);
   }
 
   @Override

diff --git a/internal-api/src/test/java/datadog/trace/api/sampling/AdaptiveSamplerTest.java b/internal-api/src/test/java/datadog/trace/api/sampling/AdaptiveSamplerTest.java
@@ -261,7 +261,7 @@ public void testRepeatingRegularStartWithLow() throws Exception {
   @Test
   public void testKeep() {
     final AdaptiveSampler sampler =
-        new AdaptiveSampler(WINDOW_DURATION, 1, 1, 1, null, taskScheduler);
+        new AdaptiveSampler(WINDOW_DURATION, 1, 1, 1, null, taskScheduler, true);
     long tests = sampler.testCount();
     long samples = sampler.sampleCount();
     assertTrue(sampler.keep());
@@ -272,7 +272,7 @@ public void testKeep() {
   @Test
   public void testDrop() {
     final AdaptiveSampler sampler =
-        new AdaptiveSampler(WINDOW_DURATION, 1, 1, 1, null, taskScheduler);
+        new AdaptiveSampler(WINDOW_DURATION, 1, 1, 1, null, taskScheduler, true);
     long tests = sampler.testCount();
     long samples = sampler.sampleCount();
     assertFalse(sampler.drop());
@@ -344,7 +344,8 @@ void testConfigListener() throws Exception {
                   }
               }
             },
-            taskScheduler);
+            taskScheduler,
+            true);
     sampler.keep();
     sampler.drop();
     rollWindow();
@@ -383,7 +384,8 @@ private void testSamplerInline(
             AVERAGE_LOOKBACK,
             BUDGET_LOOKBACK,
             null,
-            taskScheduler);
+            taskScheduler,
+            true);
 
     // simulate event generation and sampling for the given number of sampling windows
     final long expectedSamples = WINDOWS * SAMPLES_PER_WINDOW;
@@ -544,7 +546,8 @@ private void testSamplerConcurrently(
             AVERAGE_LOOKBACK,
             BUDGET_LOOKBACK,
             null,
-            taskScheduler);
+            taskScheduler,
+            true);
     final CyclicBarrier startBarrier = new CyclicBarrier(threadCount);
     final CyclicBarrier endBarrier = new CyclicBarrier(threadCount, this::rollWindow);
     final Mean[] means = new Mean[threadCount];