Demonstrate four tracer lifecycle failures with formal models

dd-trace-core/src/test/java/datadog/trace/common/writer/TraceProcessingWorkerTest.java

@@ -11,6 +11,7 @@
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 
+import datadog.trace.api.sampling.PrioritySampling;
 import datadog.trace.bootstrap.instrumentation.api.SpanPostProcessor;
 import datadog.trace.common.sampling.SingleSpanSampler;
 import datadog.trace.common.writer.ddagent.PrioritizationStrategy.PublishResult;
@@ -25,6 +26,7 @@
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
+import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
@@ -148,6 +150,64 @@ void testAFlushShouldClearThePrimaryQueue() {
     }
   }
 
+  @Test
+  void flushWaitsForInFlightSingleSpanSampling() throws Exception {
+    CountDownLatch samplingStarted = new CountDownLatch(1);
+    CountDownLatch allowSampling = new CountDownLatch(1);
+    SingleSpanSampler blockingSampler =
+        new SingleSpanSampler() {
+          @Override
+          public <T extends CoreSpan<T>> boolean setSamplingPriority(T span) {
+            samplingStarted.countDown();
+            try {
+              allowSampling.await();
+            } catch (InterruptedException e) {
+              Thread.currentThread().interrupt();
+            }
+            return true;
+          }
+        };
+    AtomicInteger serializedTraces = new AtomicInteger();
+    PayloadDispatcher dispatcher = mock(PayloadDispatcher.class);
+    doAnswer(
+            invocation -> {
+              serializedTraces.incrementAndGet();
+              return null;
+            })
+        .when(dispatcher)
+        .addTrace(any());
+
+    try (TraceProcessingWorker worker =
+        new TraceProcessingWorker(
+            10,
+            mock(HealthMetrics.class),
+            dispatcher,
+            () -> false,
+            FAST_LANE,
+            100,
+            TimeUnit.SECONDS,
+            blockingSampler)) {
+      try {
+        worker.start();
+        DDSpan span = mock(DDSpan.class);
+        when(span.samplingPriority()).thenReturn((int) PrioritySampling.USER_DROP);
+
+        assertEquals(
+            PublishResult.ENQUEUED_FOR_SINGLE_SPAN_SAMPLING,
+            worker.publish(span, PrioritySampling.USER_DROP, Collections.singletonList(span)));
+        assertTrue(samplingStarted.await(3, TimeUnit.SECONDS));
+
+        assertTrue(worker.flush(3, TimeUnit.SECONDS));
+        assertEquals(
+            1,
+            serializedTraces.get(),
+            "flush returned before an accepted single-span-sampling candidate was serialized");
+      } finally {
+        allowSampling.countDown();
+      }
+    }
+  }
+
   // -------------------------------------------------------------------------
   // Test 4: should report failure if serialization fails
   // -------------------------------------------------------------------------

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

@@ -1,6 +1,7 @@
 package datadog.trace.core;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyLong;
 import static org.mockito.Mockito.clearInvocations;
@@ -128,6 +129,20 @@ void maxTrackedTracesIsEnforced() {
     assertEquals(1, tracker.getDropped());
   }
 
+  @Test
+  void fullTrackerLeavesOverflowTraceForBufferFallback() {
+    for (int i = 0; i < MAX_TRACKED_TRACES; i++) {
+      tracker.add(newTraceToTrack());
+    }
+
+    PendingTrace overflowTrace = newTraceToTrack();
+
+    assertFalse(
+        tracker.add(overflowTrace),
+        "a trace rejected at capacity must remain available to the buffer's fallback path");
+    assertEquals(LongRunningTracesTracker.NOT_TRACKED, overflowTrace.getLongRunningTrackedState());
+  }
+
   @Test
   void expiredTraces() {
     PendingTrace trace = newTraceToTrack();

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

@@ -45,6 +45,9 @@
 import java.util.List;
 import java.util.Map;
 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.AtomicInteger;
 import java.util.zip.ZipEntry;
@@ -379,6 +382,168 @@ public boolean writeOnBufferFull() {
     assertEquals(3, counter.get());
   }
 
+  @Test
+  void flushDrainClearsEnqueuedStateBeforeWriting() throws Exception {
+    CountDownLatch workerBlocked = new CountDownLatch(1);
+    CountDownLatch allowWorker = new CountDownLatch(1);
+    CountDownLatch traceWrites = new CountDownLatch(2);
+    doAnswer(
+            invocation -> {
+              traceWrites.countDown();
+              return null;
+            })
+        .when(tracer)
+        .write(any());
+    PendingTraceBuffer.Element blocker =
+        new PendingTraceBuffer.Element() {
+          @Override
+          public long oldestFinishedTime() {
+            workerBlocked.countDown();
+            try {
+              allowWorker.await();
+            } catch (InterruptedException e) {
+              Thread.currentThread().interrupt();
+            }
+            return 0;
+          }
+
+          @Override
+          public boolean lastReferencedNanosAgo(long nanos) {
+            return false;
+          }
+
+          @Override
+          public void write() {}
+
+          @Override
+          public DDSpan getRootSpan() {
+            return null;
+          }
+
+          @Override
+          public boolean setEnqueued(boolean enqueued) {
+            return true;
+          }
+
+          @Override
+          public boolean writeOnBufferFull() {
+            return true;
+          }
+        };
+
+    buffer.start();
+    bufferSpy.enqueue(blocker);
+    assertTrue(workerBlocked.await(3, TimeUnit.SECONDS));
+
+    ExecutorService executor = Executors.newSingleThreadExecutor();
+    try {
+      Future<?> flush = executor.submit(buffer::flush);
+      long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(3);
+      while (delayingBuffer.getQueue().isEmpty() && System.nanoTime() < deadline) {
+        Thread.yield();
+      }
+      assertEquals(1, delayingBuffer.getQueue().size(), "flush marker was not queued");
+
+      PendingTrace trace = factory.create(DDTraceId.ONE);
+      DDSpan parent = newSpanOf(trace);
+      DDSpan child = newSpanOf(parent);
+      parent.finish();
+
+      assertEquals(1, trace.getIsEnqueued());
+      allowWorker.countDown();
+      flush.get(3, TimeUnit.SECONDS);
+
+      assertTrue(trace.isRootSpanWritten());
+      assertEquals(
+          0,
+          trace.getIsEnqueued(),
+          "a trace written by the flush drain must be eligible for re-enqueue");
+
+      child.finish();
+      assertTrue(traceWrites.await(3, TimeUnit.SECONDS));
+      assertEquals(0, trace.size());
+      verify(tracer, times(2)).write(any());
+    } finally {
+      allowWorker.countDown();
+      executor.shutdownNow();
+    }
+  }
+
+  @Test
+  void flushDrainDoesNotConsumeTracerDumpMarker() throws Exception {
+    CountDownLatch workerBlocked = new CountDownLatch(1);
+    CountDownLatch allowWorker = new CountDownLatch(1);
+    PendingTraceBuffer.Element blocker =
+        new PendingTraceBuffer.Element() {
+          @Override
+          public long oldestFinishedTime() {
+            workerBlocked.countDown();
+            try {
+              allowWorker.await();
+            } catch (InterruptedException e) {
+              Thread.currentThread().interrupt();
+            }
+            return 0;
+          }
+
+          @Override
+          public boolean lastReferencedNanosAgo(long nanos) {
+            return false;
+          }
+
+          @Override
+          public void write() {}
+
+          @Override
+          public DDSpan getRootSpan() {
+            return null;
+          }
+
+          @Override
+          public boolean setEnqueued(boolean enqueued) {
+            return true;
+          }
+
+          @Override
+          public boolean writeOnBufferFull() {
+            return true;
+          }
+        };
+
+    buffer.start();
+    bufferSpy.enqueue(blocker);
+    assertTrue(workerBlocked.await(3, TimeUnit.SECONDS));
+
+    ExecutorService executor = Executors.newFixedThreadPool(2);
+    Future<?> flare = null;
+    try {
+      Future<?> flush = executor.submit(buffer::flush);
+      awaitQueueSize(1);
+
+      flare = executor.submit(TracerFlare::prepareForFlare);
+      awaitQueueSize(2);
+
+      allowWorker.countDown();
+      flush.get(3, TimeUnit.SECONDS);
+      flare.get(3, TimeUnit.SECONDS);
+    } finally {
+      allowWorker.countDown();
+      buffer.close();
+      if (flare != null) {
+        flare.get(3, TimeUnit.SECONDS);
+      }
+      executor.shutdownNow();
+    }
+  }
+
+  private void awaitQueueSize(int expectedSize) throws InterruptedException {
+    long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(3);
+    while (delayingBuffer.getQueue().size() != expectedSize && System.nanoTime() < deadline) {
+      Thread.yield();
+    }
+    assertEquals(expectedSize, delayingBuffer.getQueue().size());
+  }
+
   @Test
   void samePendingTraceIsNotEnqueuedMultipleTimes() {
     when(tracer.getPartialFlushMinSpans()).thenReturn(10000);

formal-methods/flush-dump-marker/FlushDumpMarker.cfg

@@ -0,0 +1,7 @@
+SPECIFICATION Spec
+
+INVARIANTS
+  TypeOK
+  WaitingDumpHasOwner
+
+CHECK_DEADLOCK FALSE

formal-methods/flush-dump-marker/FlushDumpMarker.tla

@@ -0,0 +1,66 @@
+--------------------------- MODULE FlushDumpMarker ---------------------------
+
+(**
+ * Models a tracer-flare dump command queued behind a flush command while the
+ * PendingTraceBuffer worker is busy. Processing the flush drains the dump
+ * command through CommandElement.write(), which is a no-op, so the waiter
+ * loses both the queued command and its acknowledgement.
+ *)
+
+VARIABLES
+  flushQueued,
+  dumpQueued,
+  flushAcknowledged,
+  dumpAcknowledged,
+  dumpWaiting,
+  pc
+
+vars ==
+  <<flushQueued, dumpQueued, flushAcknowledged, dumpAcknowledged, dumpWaiting, pc>>
+
+Init ==
+  /\ flushQueued = FALSE
+  /\ dumpQueued = FALSE
+  /\ flushAcknowledged = FALSE
+  /\ dumpAcknowledged = FALSE
+  /\ dumpWaiting = FALSE
+  /\ pc = "WorkerBlocked"
+
+QueueFlush ==
+  /\ pc = "WorkerBlocked"
+  /\ flushQueued' = TRUE
+  /\ pc' = "FlushQueued"
+  /\ UNCHANGED <<dumpQueued, flushAcknowledged, dumpAcknowledged, dumpWaiting>>
+
+QueueDump ==
+  /\ pc = "FlushQueued"
+  /\ dumpQueued' = TRUE
+  /\ dumpWaiting' = TRUE
+  /\ pc' = "DumpQueued"
+  /\ UNCHANGED <<flushQueued, flushAcknowledged, dumpAcknowledged>>
+
+(** Current behavior: WriteDrain consumes the dump marker without acknowledging it. *)
+FlushDrainsDump ==
+  /\ pc = "DumpQueued"
+  /\ flushQueued' = FALSE
+  /\ dumpQueued' = FALSE
+  /\ flushAcknowledged' = TRUE
+  /\ pc' = "DumpMarkerLost"
+  /\ UNCHANGED <<dumpAcknowledged, dumpWaiting>>
+
+Next == QueueFlush \/ QueueDump \/ FlushDrainsDump
+
+Spec == Init /\ [][Next]_vars
+
+TypeOK ==
+  /\ flushQueued \in BOOLEAN
+  /\ dumpQueued \in BOOLEAN
+  /\ flushAcknowledged \in BOOLEAN
+  /\ dumpAcknowledged \in BOOLEAN
+  /\ dumpWaiting \in BOOLEAN
+  /\ pc \in {"WorkerBlocked", "FlushQueued", "DumpQueued", "DumpMarkerLost"}
+
+WaitingDumpHasOwner ==
+  dumpWaiting => (dumpQueued \/ dumpAcknowledged)
+
+=============================================================================

formal-methods/flush-dump-marker/README.md

@@ -0,0 +1,30 @@
+# Flush/dump command ownership model
+
+This model checks the ownership of a tracer-flare dump command when it is
+queued behind a `PendingTraceBuffer` flush command.
+
+Both commands share the pending-trace queue. When the worker processes the
+flush marker, `WriteDrain` drains every following element by calling
+`Element.write()`. A dump marker is a `CommandElement`, whose `write()` method
+is a no-op. The marker therefore disappears without incrementing
+`dumpCounter`, leaving `TracerDump.prepareForFlare()` waiting indefinitely.
+
+Run from the repository root:
+
+```shell
+mkdir -p /tmp/dd-trace-java-tlc/flush-dump-marker
+java -XX:+UseParallelGC -jar /tmp/tla2tools-1.7.4.jar \
+  -config formal-methods/flush-dump-marker/FlushDumpMarker.cfg \
+  -metadir /tmp/dd-trace-java-tlc/flush-dump-marker \
+  formal-methods/flush-dump-marker/FlushDumpMarker.tla
+```
+
+The current implementation produces a four-state `WaitingDumpHasOwner`
+invariant violation.
+
+The corresponding deterministic Java regression is
+`PendingTraceBufferTest.flushDrainDoesNotConsumeTracerDumpMarker`. It blocks
+the worker, queues a flush and then a flare dump, and releases the worker. The
+flush completes, but flare preparation remains stuck because its marker was
+silently consumed. On the unfixed implementation it fails with a
+`TimeoutException` while waiting for flare preparation to complete.