diff --git a/dd-trace-core/src/main/java/datadog/trace/common/writer/TraceProcessingWorker.java b/dd-trace-core/src/main/java/datadog/trace/common/writer/TraceProcessingWorker.java index 9b70b09bc9e..b2705c5fec4 100644 --- a/dd-trace-core/src/main/java/datadog/trace/common/writer/TraceProcessingWorker.java +++ b/dd-trace-core/src/main/java/datadog/trace/common/writer/TraceProcessingWorker.java @@ -4,7 +4,9 @@ import static datadog.trace.util.AgentThreadFactory.THREAD_JOIN_TIMOUT_MS; import static datadog.trace.util.AgentThreadFactory.newAgentThread; import static java.util.concurrent.TimeUnit.MILLISECONDS; +import static java.util.concurrent.TimeUnit.NANOSECONDS; +import datadog.common.container.ServerlessInfo; import datadog.common.queue.MessagePassingBlockingQueue; import datadog.common.queue.Queues; import datadog.communication.ddagent.DroppingPolicy; @@ -46,6 +48,11 @@ public class TraceProcessingWorker implements AutoCloseable { private final SpanSamplingWorker spanSamplingWorker; + // In Lambda the environment can freeze as soon as a flush returns, so also flush the secondary + // queue there rather than leaving sampled-out traces for a periodic flush that may never run. + private final boolean flushSecondaryQueue = + ServerlessInfo.get().isRunningInServerlessEnvironment(); + public TraceProcessingWorker( final int capacity, final HealthMetrics healthMetrics, @@ -85,20 +92,32 @@ public void start() { } public boolean flush(long timeout, TimeUnit timeUnit) { - CountDownLatch latch = new CountDownLatch(1); + long deadline = System.nanoTime() + timeUnit.toNanos(timeout); + CountDownLatch latch = new CountDownLatch(flushSecondaryQueue ? 2 : 1); FlushEvent flush = new FlushEvent(latch); - boolean offered; - do { - offered = primaryQueue.offer(flush); - } while (!offered && serializerThread.isAlive()); try { - return latch.await(timeout, timeUnit); + return offer(primaryQueue, flush, deadline) + && (!flushSecondaryQueue || offer(secondaryQueue, flush, deadline)) + && latch.await(deadline - System.nanoTime(), NANOSECONDS); } catch (InterruptedException e) { Thread.currentThread().interrupt(); return false; } } + private boolean offer(MessagePassingBlockingQueue queue, FlushEvent flush, long deadline) + throws InterruptedException { + // offer before checking the serializer: it may not have started yet (e.g. on IBM JDK 8) + while (!queue.offer(flush)) { + long remaining = deadline - System.nanoTime(); + if (remaining <= 0 || !serializerThread.isAlive()) { + return false; + } + NANOSECONDS.sleep(Math.min(remaining, MILLISECONDS.toNanos(1))); + } + return true; + } + @Override public void close() { spanSamplingWorker.close(); diff --git a/dd-trace-core/src/main/java/datadog/trace/common/writer/WriterFactory.java b/dd-trace-core/src/main/java/datadog/trace/common/writer/WriterFactory.java index a32c096c2b0..0ad61b81e39 100644 --- a/dd-trace-core/src/main/java/datadog/trace/common/writer/WriterFactory.java +++ b/dd-trace-core/src/main/java/datadog/trace/common/writer/WriterFactory.java @@ -135,6 +135,13 @@ public static Writer createWriter( } } + // In Lambda the execution environment can freeze as soon as the handler returns, before a + // writer's periodic flush timer next fires -- flush synchronously so buffered events (e.g. + // LLM Observability spans) aren't lost when that happens. + boolean isServerlessDefault = + config.isAgentConfiguredUsingDefault() + && ServerlessInfo.get().isRunningInServerlessEnvironment(); + RemoteWriter remoteWriter; if (DD_INTAKE_WRITER_TYPE.equals(configuredType)) { final TrackType trackType = DDIntakeTrackTypeResolver.resolve(config); @@ -148,6 +155,7 @@ public static Writer createWriter( .healthMetrics(healthMetrics) .monitoring(commObjects.monitoring) .singleSpanSampler(singleSpanSampler) + .alwaysFlush(isServerlessDefault) .flushIntervalMilliseconds(flushIntervalMilliseconds); if (config.isCiVisibilityEnabled()) { @@ -168,8 +176,7 @@ public static Writer createWriter( } else { // configuredType == DDAgentWriter boolean alwaysFlush = false; - if (config.isAgentConfiguredUsingDefault() - && ServerlessInfo.get().isRunningInServerlessEnvironment()) { + if (isServerlessDefault) { if (!ServerlessInfo.get().hasExtension()) { log.info( "Detected serverless environment. Serverless extension has not been detected, using PrintingWriter"); diff --git a/dd-trace-core/src/test/java/datadog/trace/common/writer/TraceProcessingWorkerTest.java b/dd-trace-core/src/test/java/datadog/trace/common/writer/TraceProcessingWorkerTest.java index 86dbfa61b22..f25fdfe3473 100644 --- a/dd-trace-core/src/test/java/datadog/trace/common/writer/TraceProcessingWorkerTest.java +++ b/dd-trace-core/src/test/java/datadog/trace/common/writer/TraceProcessingWorkerTest.java @@ -13,6 +13,7 @@ import datadog.trace.bootstrap.instrumentation.api.SpanPostProcessor; import datadog.trace.common.sampling.SingleSpanSampler; +import datadog.trace.common.writer.ddagent.FlushEvent; import datadog.trace.common.writer.ddagent.PrioritizationStrategy.PublishResult; import datadog.trace.core.CoreSpan; import datadog.trace.core.DDSpan; @@ -25,6 +26,10 @@ import java.util.ArrayList; import java.util.Collections; import java.util.List; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicInteger; @@ -612,4 +617,85 @@ public > boolean setSamplingPriority(T span) { assertEquals(sampledSingleSpans, sampledSpansCount.get()); } } + + // ------------------------------------------------------------------------- + // Test 10: flush times out when the queue is full and the serializer is blocked + // ------------------------------------------------------------------------- + + @Test + void testFlushTimesOutWhenQueueIsFullAndSerializerIsBlocked() throws Exception { + CountDownLatch flushStarted = new CountDownLatch(1); + CountDownLatch releaseFlush = new CountDownLatch(1); + PayloadDispatcher dispatcher = mock(PayloadDispatcher.class); + doAnswer( + inv -> { + flushStarted.countDown(); + releaseFlush.await(); + return null; + }) + .when(dispatcher) + .flush(); + TraceProcessingWorker worker = + new TraceProcessingWorker( + 2, + mock(HealthMetrics.class), + dispatcher, + () -> false, + FAST_LANE, + 0, // disable periodic flushes + TimeUnit.SECONDS, + null); + ExecutorService caller = Executors.newSingleThreadExecutor(); + try { + // block the serializer inside a flush, then fill the primary queue behind it + assertTrue(worker.getPrimaryQueue().offer(new FlushEvent(new CountDownLatch(1)))); + worker.start(); + assertTrue(flushStarted.await(5, TimeUnit.SECONDS)); + while (worker.getPrimaryQueue().offer(Collections.singletonList(mock(DDSpan.class)))) { + // fill the queue + } + + // the flush marker can't be enqueued, so the flush must give up once its timeout elapses + Future flushed = caller.submit(() -> worker.flush(100, TimeUnit.MILLISECONDS)); + assertFalse(flushed.get(5, TimeUnit.SECONDS)); + } finally { + releaseFlush.countDown(); + worker.close(); + caller.shutdownNow(); + } + } + + // ------------------------------------------------------------------------- + // Test 11: flush succeeds when the serializer starts after the flush begins + // ------------------------------------------------------------------------- + + @Test + void testFlushSucceedsWhenSerializerStartsAfterFlushBegins() throws Exception { + AtomicInteger flushCount = new AtomicInteger(); + TraceProcessingWorker worker = + new TraceProcessingWorker( + 10, + mock(HealthMetrics.class), + flushCountingPayloadDispatcher(flushCount), + () -> false, + FAST_LANE, + 100, + TimeUnit.SECONDS, // prevent heartbeats from helping the flush happen + null); + ExecutorService caller = Executors.newSingleThreadExecutor(); + try { + // the agent can delay starting the writer (e.g. on IBM JDK 8) past a short-lived app's exit + Future flushed = caller.submit(() -> worker.flush(10, TimeUnit.SECONDS)); + while (worker.getPrimaryQueue().isEmpty() && !flushed.isDone()) { + Thread.yield(); + } + worker.start(); + + assertTrue(flushed.get(5, TimeUnit.SECONDS)); + assertEquals(1, flushCount.get()); + } finally { + worker.close(); + caller.shutdownNow(); + } + } }