Fix span timestamps staying stale after an AWS Lambda SnapStart restore

dd-trace-api/src/main/java/datadog/trace/api/ConfigDefaults.java

@@ -95,6 +95,8 @@ public final class ConfigDefaults {
 
   static final int DEFAULT_CLOCK_SYNC_PERIOD = 30; // seconds
 
+  static final boolean DEFAULT_TRACE_LAMBDA_SNAPSTART_CLOCK_RESYNC_ENABLED = false;
+
   static final TracePropagationBehaviorExtract DEFAULT_TRACE_PROPAGATION_BEHAVIOR_EXTRACT =
       TracePropagationBehaviorExtract.CONTINUE;
   static final boolean DEFAULT_TRACE_PROPAGATION_EXTRACT_FIRST = false;

dd-trace-api/src/main/java/datadog/trace/api/config/TracerConfig.java

@@ -133,6 +133,17 @@ public final class TracerConfig {
 
   public static final String CLOCK_SYNC_PERIOD = "trace.clock.sync.period";
 
+  /**
+   * Opt-in (defaults to disabled). When enabled, resyncs the tracer's cached clock reference on
+   * every AWS Lambda invocation, before any span for that invocation is created, instead of
+   * relying on the periodic {@link #CLOCK_SYNC_PERIOD} check - which is gated on monotonic ticks
+   * elapsed and may never trigger following an AWS Lambda SnapStart restore, since a short-lived
+   * restore-then-invoke sequence can leave that reference stale for hours or days. Only takes
+   * effect inside an instrumented AWS Lambda handler invocation; a no-op everywhere else.
+   */
+  public static final String TRACE_LAMBDA_SNAPSTART_CLOCK_RESYNC_ENABLED =
+      "trace.lambda.snapstart.clock.resync.enabled";
+
   public static final String TRACE_SPAN_ATTRIBUTE_SCHEMA = "trace.span.attribute.schema";
 
   public static final String TRACE_LONG_RUNNING_ENABLED = "trace.experimental.long-running.enabled";

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

@@ -1032,6 +1032,50 @@ long getTimeWithNanoTicks(long nanoTicks) {
     return computedNanoTime + counterDrift;
   }
 
+  /**
+   * AWS Lambda SnapStart checkpoints the JVM at snapshot-creation time and later restores it,
+   * potentially much later (a restored snapshot can go {@code Inactive} and be restored again after
+   * 14 days of no invocations). {@link System#nanoTime()} does not account for the frozen duration
+   * across that restore, so {@link #startTimeNano}/{@link #startNanoTicks} - captured once at
+   * construction time, i.e. at snapshot-creation time - go stale, and the periodic self-correction
+   * in {@link #getTimeWithNanoTicks} may never trigger to fix it: that correction is gated on
+   * {@code nanoTicks - lastSyncTicks >= clockSyncPeriod} (monotonic ticks elapsed), which a
+   * short-lived restore-then-invoke sequence can easily never reach even though wall-clock time has
+   * moved on by hours or days. Left uncorrected, every span timestamp computed via {@link
+   * #getTimeWithNanoTicks} stays anchored near the original snapshot-creation instant instead of
+   * the real restore/invocation time.
+   *
+   * <p>Rather than reset {@link #startTimeNano}/{@link #startNanoTicks} themselves - which would
+   * require making those fields {@code volatile}, since they're read on every {@link
+   * #getTimeWithNanoTicks} call from arbitrary tracing threads - this folds the entire observed
+   * drift into {@link #counterDrift} directly, the same field (already {@code volatile}) that the
+   * periodic self-correction in {@link #getTimeWithNanoTicks} uses, just without that correction's
+   * 1ms drift threshold, since here the drift can be hours or days.
+   *
+   * <p>Rather than reacting to the restore event itself (which would need a JVM-level
+   * checkpoint/restore hook), this is called from {@link #notifyLambdaStart} - already invoked once
+   * per Lambda invocation, before any span for that invocation is created. Any real restore is
+   * always followed by an invocation, so resyncing there catches it with no extra dependency. Doing
+   * this on every invocation (not just ones following a restore) is deliberate: outside SnapStart,
+   * {@link System#nanoTime()} correctly tracks elapsed time across a warm container's normal
+   * freeze/thaw between invocations (same continuously-executing process, unlike SnapStart's
+   * restore-into-a-new-context), so this is a correct no-op there - just a couple of field reads
+   * and a subtraction, negligible next to the HTTP round-trip {@link #notifyLambdaStart} already
+   * makes.
+   *
+   * <p>Gated on {@link Config#isLambdaSnapStartClockResyncEnabled()} as an escape hatch.
+   */
+  @VisibleForTesting
+  void maybeResyncClockForLambdaInvocation() {
+    if (!initialConfig.isLambdaSnapStartClockResyncEnabled()) {
+      return;
+    }
+    long nanoTicks = timeSource.getNanoTicks();
+    long computedNanoTime = startTimeNano + Math.max(0, nanoTicks - startNanoTicks);
+    counterDrift = timeSource.getCurrentTimeNanos() - computedNanoTime;
+    lastSyncTicks = nanoTicks;
+  }
+
   @Override
   public CoreSpanBuilder buildSpan(
       final String instrumentationName, final CharSequence operationName) {
@@ -1238,6 +1282,8 @@ public void closeActive() {
 
   @Override
   public AgentSpanContext notifyLambdaStart(Object event, String lambdaRequestId) {
+    maybeResyncClockForLambdaInvocation();
+
     // Get context from AppSec
     AgentSpanContext appSecContext = LambdaAppSecHandler.processRequestStart(event);
 

dd-trace-core/src/test/java/datadog/trace/core/CoreTracerTest.java

@@ -28,6 +28,7 @@
 import datadog.trace.api.config.TracerConfig;
 import datadog.trace.api.remoteconfig.ServiceNameCollector;
 import datadog.trace.api.sampling.PrioritySampling;
+import datadog.trace.api.time.ControllableTimeSource;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import datadog.trace.bootstrap.instrumentation.api.ServiceNameSources;
 import datadog.trace.common.sampling.AllSampler;
@@ -79,6 +80,94 @@ void verifyDefaultsOnTracer() {
     }
   }
 
+  @Test
+  void
+      getTimeWithNanoTicks_whenNanoTicksStaleAfterSimulatedRestore_thenTimestampStaysAnchoredToConstructionTime() {
+    // Characterizes the AWS Lambda SnapStart bug this fix addresses: System.nanoTime() does not
+    // account for the frozen duration across a checkpoint/restore, so a nanoTicks reading taken
+    // right after restore can be indistinguishable from one taken at construction (snapshot
+    // creation) time, even though wall-clock time has moved on by hours. Without a resync, span
+    // timestamps computed from that stale nanoTicks stay anchored to construction time.
+    ControllableTimeSource timeSource = new ControllableTimeSource();
+    timeSource.set(TimeUnit.SECONDS.toNanos(1000));
+    CoreTracer tracer = tracerBuilder().writer(new ListWriter()).timeSource(timeSource).build();
+    try {
+      long constructionTimeNanoTicks = timeSource.getNanoTicks();
+
+      // Wall-clock time moves on by 2 hours (the restore happens much later), but the nanoTicks
+      // value a span would be timestamped with right after restore hasn't advanced.
+      timeSource.set(TimeUnit.SECONDS.toNanos(1000) + TimeUnit.HOURS.toNanos(2));
+
+      assertEquals(
+          TimeUnit.SECONDS.toNanos(1000), tracer.getTimeWithNanoTicks(constructionTimeNanoTicks));
+    } finally {
+      tracer.close();
+    }
+  }
+
+  @Test
+  @WithConfig(key = TracerConfig.TRACE_LAMBDA_SNAPSTART_CLOCK_RESYNC_ENABLED, value = "true")
+  void
+      maybeResyncClockForLambdaInvocation_whenEnabledAndCalledAfterSimulatedRestore_thenTimestampReflectsPostRestoreTime() {
+    ControllableTimeSource timeSource = new ControllableTimeSource();
+    timeSource.set(TimeUnit.SECONDS.toNanos(1000));
+    CoreTracer tracer = tracerBuilder().writer(new ListWriter()).timeSource(timeSource).build();
+    try {
+      // The restore happens: wall-clock/tick source jumps forward by 2 hours.
+      long postRestoreNanos = TimeUnit.SECONDS.toNanos(1000) + TimeUnit.HOURS.toNanos(2);
+      timeSource.set(postRestoreNanos);
+
+      tracer.maybeResyncClockForLambdaInvocation();
+
+      // A span timestamped right after resync now reflects the real post-restore time, not the
+      // stale construction-time (pre-restore) anchor.
+      assertEquals(postRestoreNanos, tracer.getTimeWithNanoTicks(timeSource.getNanoTicks()));
+    } finally {
+      tracer.close();
+    }
+  }
+
+  @Test
+  @WithConfig(key = TracerConfig.TRACE_LAMBDA_SNAPSTART_CLOCK_RESYNC_ENABLED, value = "true")
+  void notifyLambdaStart_whenExplicitlyEnabled_thenResyncsClockToCurrentTime() {
+    // notifyLambdaStart runs once per Lambda invocation, before any span for that invocation is
+    // created - the actual trigger point for the resync in production, not just the extracted
+    // maybeResyncClockForLambdaInvocation() logic exercised directly above.
+    ControllableTimeSource timeSource = new ControllableTimeSource();
+    timeSource.set(TimeUnit.SECONDS.toNanos(1000));
+    CoreTracer tracer = tracerBuilder().writer(new ListWriter()).timeSource(timeSource).build();
+    try {
+      long postRestoreNanos = TimeUnit.SECONDS.toNanos(1000) + TimeUnit.HOURS.toNanos(2);
+      timeSource.set(postRestoreNanos);
+
+      tracer.notifyLambdaStart(new Object(), "lambda-request-123");
+
+      assertEquals(postRestoreNanos, tracer.getTimeWithNanoTicks(timeSource.getNanoTicks()));
+    } finally {
+      tracer.close();
+    }
+  }
+
+  @Test
+  void
+      notifyLambdaStart_whenResyncDefaultsToDisabled_thenTimestampStaysAnchoredToConstructionTime() {
+    ControllableTimeSource timeSource = new ControllableTimeSource();
+    timeSource.set(TimeUnit.SECONDS.toNanos(1000));
+    CoreTracer tracer = tracerBuilder().writer(new ListWriter()).timeSource(timeSource).build();
+    try {
+      long constructionTimeNanoTicks = timeSource.getNanoTicks();
+      timeSource.set(TimeUnit.SECONDS.toNanos(1000) + TimeUnit.HOURS.toNanos(2));
+
+      // No @WithConfig override here - this is an opt-in feature, off by default.
+      tracer.notifyLambdaStart(new Object(), "lambda-request-123");
+
+      assertEquals(
+          TimeUnit.SECONDS.toNanos(1000), tracer.getTimeWithNanoTicks(constructionTimeNanoTicks));
+    } finally {
+      tracer.close();
+    }
+  }
+
   @Test
   @WithConfig(key = TracerConfig.PRIORITY_SAMPLING, value = "false")
   void verifyOverridingSampler() {

internal-api/src/main/java/datadog/trace/api/Config.java

@@ -178,6 +178,7 @@
 import static datadog.trace.api.ConfigDefaults.DEFAULT_TRACE_EXPERIMENTAL_FEATURES_ENABLED;
 import static datadog.trace.api.ConfigDefaults.DEFAULT_TRACE_HTTP_RESOURCE_REMOVE_TRAILING_SLASH;
 import static datadog.trace.api.ConfigDefaults.DEFAULT_TRACE_KEEP_LATENCY_THRESHOLD_MS;
+import static datadog.trace.api.ConfigDefaults.DEFAULT_TRACE_LAMBDA_SNAPSTART_CLOCK_RESYNC_ENABLED;
 import static datadog.trace.api.ConfigDefaults.DEFAULT_TRACE_LONG_RUNNING_ENABLED;
 import static datadog.trace.api.ConfigDefaults.DEFAULT_TRACE_LONG_RUNNING_FLUSH_INTERVAL;
 import static datadog.trace.api.ConfigDefaults.DEFAULT_TRACE_LONG_RUNNING_INITIAL_FLUSH_INTERVAL;
@@ -691,6 +692,7 @@
 import static datadog.trace.api.config.TracerConfig.TRACE_HTTP_SERVER_PATH_RESOURCE_NAME_MAPPING;
 import static datadog.trace.api.config.TracerConfig.TRACE_INFERRED_PROXY_SERVICES_ENABLED;
 import static datadog.trace.api.config.TracerConfig.TRACE_KEEP_LATENCY_THRESHOLD_MS;
+import static datadog.trace.api.config.TracerConfig.TRACE_LAMBDA_SNAPSTART_CLOCK_RESYNC_ENABLED;
 import static datadog.trace.api.config.TracerConfig.TRACE_LONG_RUNNING_ENABLED;
 import static datadog.trace.api.config.TracerConfig.TRACE_LONG_RUNNING_FLUSH_INTERVAL;
 import static datadog.trace.api.config.TracerConfig.TRACE_LONG_RUNNING_INITIAL_FLUSH_INTERVAL;
@@ -1317,6 +1319,8 @@ public static String getHostName() {
 
   private final boolean secureRandom;
 
+  private final boolean lambdaSnapStartClockResyncEnabled;
+
   private final boolean trace128bitTraceIdGenerationEnabled;
   private final boolean logs128bitTraceIdEnabled;
 
@@ -1523,6 +1527,10 @@ private Config(final ConfigProvider configProvider, final InstrumenterConfig ins
     } else {
       secureRandom = configProvider.getBoolean(SECURE_RANDOM, DEFAULT_SECURE_RANDOM);
     }
+    lambdaSnapStartClockResyncEnabled =
+        configProvider.getBoolean(
+            TRACE_LAMBDA_SNAPSTART_CLOCK_RESYNC_ENABLED,
+            DEFAULT_TRACE_LAMBDA_SNAPSTART_CLOCK_RESYNC_ENABLED);
     cassandraKeyspaceStatementExtractionEnabled =
         configProvider.getBoolean(
             CASSANDRA_KEYSPACE_STATEMENT_EXTRACTION_ENABLED,
@@ -4979,6 +4987,10 @@ public boolean isAwsServerless() {
     return awsServerless;
   }
 
+  public boolean isLambdaSnapStartClockResyncEnabled() {
+    return lambdaSnapStartClockResyncEnabled;
+  }
+
   public boolean isDataStreamsEnabled() {
     return dataStreamsEnabled;
   }

metadata/supported-configurations.json

@@ -5097,6 +5097,14 @@
         "aliases": []
       }
     ],
+    "DD_TRACE_LAMBDA_SNAPSTART_CLOCK_RESYNC_ENABLED": [
+      {
+        "version": "A",
+        "type": "boolean",
+        "default": "false",
+        "aliases": []
+      }
+    ],
     "DD_TRACE_CLOUD_PAYLOAD_TAGGING_MAX_DEPTH": [
       {
         "version": "A",