Ensure priority sampling is sent in case of partial flush

dd-smoke-tests/springboot-openliberty/snapshots/datadog.smoketest.SpringBootOpenLibertySnapshotTest.exception404.snap

@@ -103,4 +103,5 @@
                            "error.stack" "org.springframework.http.converter.HttpMessageNotWritableException: No converter for [class java.util.LinkedHashMap] with preset Content-Type 'text/html;charset=UTF-8'\n\tat org.springframework.web.servlet.mvc.method.annotation.AbstractMessageConverterMethodProcessor.writeWithMessageConverters(AbstractMessageConverterMethodProcessor.java:309)\n\tat org.springframework.web.servlet.mvc.method.annotation.HttpEntityMethodProcessor.handleReturnValue(HttpEntityMethodProcessor.java:219)\n\tat org.springframework.web.method.support.HandlerMethodReturnValueHandlerComposite.handleReturnValue(HandlerMethodReturnValueHandlerComposite.java:82)\n\tat org.springframework.web.servlet.mvc.method.annotation.ServletInvocableHandlerMethod.invokeAndHandle(ServletInvocableHandlerMethod.java:123)\n\tat org.springframework.web.servlet.mvc.method.annotation.RequestMappingHandlerAdapter.invokeHandlerMethod(RequestMappingHandlerAdapter.java:879)\n\tat org.springframework.web.servlet.mvc.method.annotation.RequestMappingHandlerAdapter.handleInternal(RequestMappingHandlerAdapter.java:793)\n\tat org.springframework.web.servlet.mvc.method.AbstractHandlerMethodAdapter.handle(AbstractHandlerMethodAdapter.java:87)\n\tat org.springframework.web.servlet.DispatcherServlet.doDispatch(DispatcherServlet.java:1040)\n\tat org.springframework.web.servlet.DispatcherServlet.doService(DispatcherServlet.java:943)\n\tat org.springframework.web.servlet.FrameworkServlet.processRequest(FrameworkServlet.java:1006)\n\tat org.springframework.web.servlet.FrameworkServlet.doGet(FrameworkServlet.java:898)\n\tat javax.servlet.http.HttpServlet.service(HttpServlet.java:686)\n\tat org.springframework.web.servlet.FrameworkServlet.service(FrameworkServlet.java:883)\n\tat javax.servlet.http.HttpServlet.service(HttpServlet.java:791)\n\tat com.ibm.ws.webcontainer.servlet.ServletWrapper.service(ServletWrapper.java:1257)\n\tat com.ibm.ws.webcontainer.servlet.ServletWrapper.handleRequest(ServletWrapper.java:745)\n\tat com.ibm.ws.webcontainer.servlet.ServletWrapper.handleRequest(ServletWrapper.java:442)\n\tat com.ibm.ws.webcontainer.filter.WebAppFilterChain.invokeTarget(WebAppFilterChain.java:182)\n\tat com.ibm.ws.webcontainer.filter.WebAppFilterChain.doFilter(WebAppFilterChain.java:78)\n\tat com.ibm.ws.webcontainer.filter.WebAppFilterManager.doFilter(WebAppFilterManager.java:1001)\n\tat com.ibm.ws.webcontainer.filter.WebAppFilterManager.invokeFilters(WebAppFilterManager.java:1139)\n\tat com.ibm.ws.webcontainer.filter.WebAppFilterManager.invokeFilters(WebAppFilterManager.java:1010)\n\tat com.ibm.ws.webcontainer.webapp.WebAppRequestDispatcher.dispatch(WebAppRequestDispatcher.java:1410)\n\tat com.ibm.ws.webcontainer.webapp.WebAppRequestDispatcher.forward(WebAppRequestDispatcher.java:171)\n\tat com.ibm.ws.webcontainer.webapp.WebApp.sendError(WebApp.java:4350)\n\tat com.ibm.ws.webcontainer.webapp.WebAppDispatcherContext.sendError(WebAppDispatcherContext.java:647)\n\tat com.ibm.ws.webcontainer.webapp.WebAppDispatcherContext.sendError(WebAppDispatcherContext.java:657)\n\tat com.ibm.ws.webcontainer.srt.SRTServletResponse.sendError(SRTServletResponse.java:1334)\n\tat com.ibm.ws.webcontainer.srt.SRTServletResponse.sendError(SRTServletResponse.java:1310)\n\tat org.springframework.web.servlet.resource.ResourceHttpRequestHandler.handleRequest(ResourceHttpRequestHandler.java:488)\n\tat org.springframework.web.servlet.mvc.HttpRequestHandlerAdapter.handle(HttpRequestHandlerAdapter.java:53)\n\tat org.springframework.web.servlet.DispatcherServlet.doDispatch(DispatcherServlet.java:1040)\n\tat org.springframework.web.servlet.DispatcherServlet.doService(DispatcherServlet.java:943)\n\tat org.springframework.web.servlet.FrameworkServlet.processRequest(FrameworkServlet.java:1006)\n\tat org.springframework.web.servlet.FrameworkServlet.doGet(FrameworkServlet.java:898)\n\tat javax.servlet.http.HttpServlet.service(HttpServlet.java:686)\n\tat org.springframework.web.servlet.FrameworkServlet.service(FrameworkServlet.java:883)\n\tat javax.servlet.http.HttpServlet.service(HttpServlet.java:791)\n\tat com.ibm.ws.webcontainer.servlet.ServletWrapper.service(ServletWrapper.java:1257)\n\tat com.ibm.ws.webcontainer.servlet.ServletWrapper.handleRequest(ServletWrapper.java:745)\n\tat com.ibm.ws.webcontainer.servlet.ServletWrapper.handleRequest(ServletWrapper.java:442)\n\tat com.ibm.ws.webcontainer.filter.WebAppFilterChain.invokeTarget(WebAppFilterChain.java:182)\n\tat com.ibm.ws.webcontainer.filter.WebAppFilterChain.doFilter(WebAppFilterChain.java:93)\n\tat org.springframework.web.filter.RequestContextFilter.doFilterInternal(RequestContextFilter.java:100)\n\tat org.springframework.web.filter.OncePerRequestFilter.doFilter(OncePerRequestFilter.java:119)\n\tat com.ibm.ws.webcontainer.filter.FilterInstanceWrapper.doFilter(FilterInstanceWrapper.java:201)\n\tat com.ibm.ws.webcontainer.filter.WebAppFilterChain.doFilter(WebAppFilterChain.java:90)\n\tat org.springframework.web.filter.FormContentFilter.doFilterInternal(FormContentFilter.java:93)\n\tat org.springframework.web.filter.OncePerRequestFilter.doFilter(OncePerRequestFilter.java:119)\n\tat com.ibm.ws.webcontainer.filter.FilterInstanceWrapper.doFilter(FilterInstanceWrapper.java:201)\n\tat com.ibm.ws.webcontainer.filter.WebAppFilterChain.doFilter(WebAppFilterChain.java:90)\n\tat datadog.trace.instrumentation.springweb.HandlerMappingResourceNameFilter.doFilterInternal(HandlerMappingResourceNameFilter.java:49)\n\tat org.springframework.web.filter.OncePerRequestFilter.doFilter(OncePerRequestFilter.java:119)\n\tat com.ibm.ws.webcontainer.filter.FilterInstanceWrapper.doFilter(FilterInstanceWrapper.java:201)\n\tat com.ibm.ws.webcontainer.filter.WebAppFilterChain.doFilter(WebAppFilterChain.java:90)\n\tat org.springframework.boot.actuate.metrics.web.servlet.WebMvcMetricsFilter.doFilterInternal(WebMvcMetricsFilter.java:109)\n\tat org.springframework.web.filter.OncePerRequestFilter.doFilter(OncePerRequestFilter.java:119)\n\tat com.ibm.ws.webcontainer.filter.FilterInstanceWrapper.doFilter(FilterInstanceWrapper.java:201)\n\tat com.ibm.ws.webcontainer.filter.WebAppFilterChain.doFilter(WebAppFilterChain.java:90)\n\tat org.springframework.web.filter.CharacterEncodingFilter.doFilterInternal(CharacterEncodingFilter.java:201)\n\tat org.springframework.web.filter.OncePerRequestFilter.doFilter(OncePerRequestFilter.java:119)\n\tat com.ibm.ws.webcontainer.filter.FilterInstanceWrapper.doFilter(FilterInstanceWrapper.java:201)\n\tat com.ibm.ws.webcontainer.filter.WebAppFilterChain.doFilter(WebAppFilterChain.java:90)\n\tat com.ibm.ws.webcontainer.filter.WebAppFilterManager.doFilter(WebAppFilterManager.java:1001)\n\tat com.ibm.ws.webcontainer.filter.WebAppFilterManager.invokeFilters(WebAppFilterManager.java:1139)\n\tat com.ibm.ws.webcontainer.webapp.WebApp.handleRequest(WebApp.java:5057)\n\tat com.ibm.ws.webcontainer.osgi.DynamicVirtualHost$2.handleRequest(DynamicVirtualHost.java:314)\n\tat com.ibm.ws.webcontainer.WebContainer.handleRequest(WebContainer.java:1006)\n\tat com.ibm.ws.webcontainer.osgi.DynamicVirtualHost$2.run(DynamicVirtualHost.java:279)\n\tat com.ibm.ws.http.dispatcher.internal.channel.HttpDispatcherLink$TaskWrapper.run(HttpDispatcherLink.java:1146)\n\tat com.ibm.ws.http.dispatcher.internal.channel.HttpDispatcherLink.wrapHandlerAndExecute(HttpDispatcherLink.java:427)\n\tat com.ibm.ws.http.dispatcher.internal.channel.HttpDispatcherLink.ready(HttpDispatcherLink.java:386)\n\tat com.ibm.ws.http.channel.internal.inbound.HttpInboundLink.handleDiscrimination(HttpInboundLink.java:566)\n\tat com.ibm.ws.http.channel.internal.inbound.HttpInboundLink.handleNewRequest(HttpInboundLink.java:500)\n\tat com.ibm.ws.http.channel.internal.inbound.HttpInboundLink.processRequest(HttpInboundLink.java:360)\n\tat com.ibm.ws.http.channel.internal.inbound.HttpInboundLink.ready(HttpInboundLink.java:327)\n\tat com.ibm.ws.tcpchannel.internal.NewConnectionInitialReadCallback.sendToDiscriminators(NewConnectionInitialReadCallback.java:167)\n\tat com.ibm.ws.tcpchannel.internal.NewConnectionInitialReadCallback.complete(NewConnectionInitialReadCallback.java:75)\n\tat com.ibm.ws.tcpchannel.internal.WorkQueueManager.requestComplete(WorkQueueManager.java:504)\n\tat com.ibm.ws.tcpchannel.internal.WorkQueueManager.attemptIO(WorkQueueManager.java:574)\n\tat com.ibm.ws.tcpchannel.internal.WorkQueueManager.workerRun(WorkQueueManager.java:958)\n\tat com.ibm.ws.tcpchannel.internal.WorkQueueManager$Worker.run(WorkQueueManager.java:1047)\n\tat com.ibm.ws.threading.internal.ExecutorServiceImpl$RunnableWrapper.run(ExecutorServiceImpl.java:238)\n\tat datadog.trace.bootstrap.instrumentation.java.concurrent.Wrapper.run(Wrapper.java:25)\n\tat java.base/java.util.concurrent.ThreadPoolExecutor.runWorker(ThreadPoolExecutor.java:1128)\n\tat java.base/java.util.concurrent.ThreadPoolExecutor$Worker.run(ThreadPoolExecutor.java:628)\n\tat java.base/java.lang.Thread.run(Thread.java:834)\n"
                            "error.type" "org.springframework.http.converter.HttpMessageNotWritableException"}
                    "metrics" {"_dd.measured" 1
+                              "_sampling_priority_v1" 1
                               "thread.id" 62}}]]

dd-smoke-tests/springboot-openliberty/snapshots/datadog.smoketest.SpringBootOpenLibertySnapshotTest.nested.snap

@@ -67,6 +67,7 @@
                    "http.method" "GET"
                    "peer.hostname" "localhost"}
             "metrics" {"_dd.measured" 1
+                       "_sampling_priority_v1" 1
                       "peer.port" 57690
                       "thread.id" 58}}
               {"name" "servlet.request"
@@ -116,4 +117,5 @@
                            "component" "spring-web-controller"
                            "span.kind" "server"}
                    "metrics" {"_dd.measured" 1
+                              "_sampling_priority_v1" 1
                               "thread.id" 47}}]]

dd-smoke-tests/springboot-openliberty/snapshots/datadog.smoketest.SpringBootOpenLibertySnapshotTest.simple.snap

@@ -46,4 +46,5 @@
                "component" "spring-web-controller"
                "span.kind" "server"}
        "metrics" {"_dd.measured" 1
+                  "_sampling_priority_v1" 1
                   "thread.id" 44}}]]

dd-trace-core/src/main/java/datadog/trace/common/writer/ddagent/TraceMapperV0_4.java

@@ -44,12 +44,18 @@ public TraceMapperV0_4() {
   private static final class MetaWriter extends MetadataConsumer {
 
     private Writable writable;
+    private boolean writeSamplingPriority;
 
     MetaWriter withWritable(Writable writable) {
       this.writable = writable;
       return this;
     }
 
+    MetaWriter withWriteSamplingPriority(final boolean writeSamplingPriority) {
+      this.writeSamplingPriority = writeSamplingPriority;
+      return this;
+    }
+
     @Override
     public void accept(Metadata metadata) {
       int metaSize =
@@ -59,7 +65,7 @@ public void accept(Metadata metadata) {
               + (null == metadata.getOrigin() ? 0 : 1)
               + 1;
       int metricsSize =
-          (metadata.hasSamplingPriority() ? 1 : 0)
+          (writeSamplingPriority && metadata.hasSamplingPriority() ? 1 : 0)
               + (metadata.measured() ? 1 : 0)
               + (metadata.topLevel() ? 1 : 0)
               + 1;
@@ -71,7 +77,7 @@ public void accept(Metadata metadata) {
       }
       writable.writeUTF8(METRICS);
       writable.startMap(metricsSize);
-      if (metadata.hasSamplingPriority()) {
+      if (writeSamplingPriority && metadata.hasSamplingPriority()) {
         writable.writeUTF8(SAMPLING_PRIORITY_KEY);
         writable.writeInt(metadata.samplingPriority());
       }
@@ -125,7 +131,8 @@ public void accept(Metadata metadata) {
   @Override
   public void map(List<? extends CoreSpan<?>> trace, final Writable writable) {
     writable.startArray(trace.size());
-    for (CoreSpan<?> span : trace) {
+    for (int i = 0; i < trace.size(); i++) {
+      final CoreSpan<?> span = trace.get(i);
       writable.startMap(12);
       /* 1  */
       writable.writeUTF8(SERVICE);
@@ -158,7 +165,10 @@ public void map(List<? extends CoreSpan<?>> trace, final Writable writable) {
       writable.writeUTF8(ERROR);
       writable.writeInt(span.getError());
       /* 11, 12 */
-      span.processTagsAndBaggage(metaWriter.withWritable(writable));
+      span.processTagsAndBaggage(
+          metaWriter
+              .withWritable(writable)
+              .withWriteSamplingPriority(i == 0 || i == trace.size() - 1));
     }
   }
 

dd-trace-core/src/main/java/datadog/trace/common/writer/ddagent/TraceMapperV0_5.java

@@ -51,7 +51,8 @@ public TraceMapperV0_5(final int dictionarySize) {
   @Override
   public void map(final List<? extends CoreSpan<?>> trace, final Writable writable) {
     writable.startArray(trace.size());
-    for (final CoreSpan<?> span : trace) {
+    for (int i = 0; i < trace.size(); i++) {
+      final CoreSpan<?> span = trace.get(i);
       writable.startArray(12);
       /* 1  */
       writeDictionaryEncoded(writable, span.getServiceName());
@@ -72,7 +73,10 @@ public void map(final List<? extends CoreSpan<?>> trace, final Writable writable
       /* 9  */
       writable.writeInt(span.getError());
       /* 10, 11  */
-      span.processTagsAndBaggage(metaWriter.withWritable(writable));
+      span.processTagsAndBaggage(
+          metaWriter
+              .withWritable(writable)
+              .withWriteSamplingPriority(i == 0 || i == trace.size() - 1));
       /* 12 */
       writeDictionaryEncoded(writable, span.getType());
     }
@@ -173,12 +177,18 @@ private List<ByteBuffer> toList() {
   private final class MetaWriter extends MetadataConsumer {
 
     private Writable writable;
+    private boolean writeSamplingPriority;
 
     MetaWriter withWritable(final Writable writable) {
       this.writable = writable;
       return this;
     }
 
+    MetaWriter withWriteSamplingPriority(final boolean writeSamplingPriority) {
+      this.writeSamplingPriority = writeSamplingPriority;
+      return this;
+    }
+
     @Override
     public void accept(Metadata metadata) {
       int metaSize =
@@ -188,7 +198,7 @@ public void accept(Metadata metadata) {
               + (null == metadata.getOrigin() ? 0 : 1)
               + 1;
       int metricsSize =
-          (metadata.hasSamplingPriority() ? 1 : 0)
+          (writeSamplingPriority && metadata.hasSamplingPriority() ? 1 : 0)
               + (metadata.measured() ? 1 : 0)
               + (metadata.topLevel() ? 1 : 0)
               + 1;
@@ -223,7 +233,7 @@ public void accept(Metadata metadata) {
         }
       }
       writable.startMap(metricsSize);
-      if (metadata.hasSamplingPriority()) {
+      if (writeSamplingPriority && metadata.hasSamplingPriority()) {
         writeDictionaryEncoded(writable, SAMPLING_PRIORITY_KEY);
         writable.writeInt(metadata.samplingPriority());
       }

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

@@ -560,7 +560,9 @@ public void processTagsAndBaggage(final MetadataConsumer consumer) {
               threadName,
               unsafeTags,
               ddTagsAndBaggageItems,
-              SamplingDecision.priority(samplingDecision),
+              (samplingDecision != SamplingDecision.UNSET_UNKNOWN
+                  ? SamplingDecision.priority(samplingDecision)
+                  : getSamplingPriority()),
               measured,
               topLevel,
               httpStatusCode == 0 ? null : HTTP_STATUSES.get(httpStatusCode),

dd-trace-core/src/test/groovy/datadog/trace/common/writer/ddagent/TraceGenerator.groovy

@@ -3,6 +3,7 @@ package datadog.trace.common.writer.ddagent
 import datadog.trace.api.DDId
 import datadog.trace.api.DDTags
 import datadog.trace.api.IdGenerationStrategy
+import datadog.trace.api.sampling.PrioritySampling
 import datadog.trace.bootstrap.instrumentation.api.UTF8BytesString
 import datadog.trace.core.CoreSpan
 import datadog.trace.core.Metadata
@@ -94,6 +95,7 @@ class TraceGenerator {
       tags,
       "type-" + random.nextInt(lowCardinality ? 1 : 100),
       random.nextBoolean(),
+      PrioritySampling.SAMPLER_KEEP,
       200,
       "some-origin")
   }
@@ -131,6 +133,7 @@ class TraceGenerator {
     private final boolean measured
     private final Metadata metadata
     private short httpStatusCode
+    private final int samplingPriority
 
     PojoSpan(
     String serviceName,
@@ -146,6 +149,7 @@ class TraceGenerator {
     Map<String, Object> tags,
     String type,
     boolean measured,
+    int samplingPriority,
     int statusCode,
     CharSequence origin) {
       this.serviceName = UTF8BytesString.create(serviceName)
@@ -159,8 +163,9 @@ class TraceGenerator {
       this.error = error
       this.type = type
       this.measured = measured
+      this.samplingPriority = samplingPriority
       this.metadata = new Metadata(Thread.currentThread().getId(),
-        UTF8BytesString.create(Thread.currentThread().getName()), tags, baggage, UNSET, measured, topLevel,
+        UTF8BytesString.create(Thread.currentThread().getName()), tags, baggage, samplingPriority, measured, topLevel,
         statusCode == 0 ? null : UTF8BytesString.create(Integer.toString(statusCode)), origin)
       this.httpStatusCode = (short) statusCode
     }
@@ -355,7 +360,7 @@ class TraceGenerator {
 
     @Override
     int samplingPriority() {
-      return UNSET
+      return samplingPriority
     }
 
     @Override
@@ -385,7 +390,7 @@ class TraceGenerator {
 
     @Override
     boolean hasSamplingPriority() {
-      return false
+      return samplingPriority != UNSET
     }
   }
 }

dd-trace-core/src/test/groovy/datadog/trace/common/writer/ddagent/TraceMapperV04PayloadTest.groovy

@@ -1,10 +1,11 @@
 package datadog.trace.common.writer.ddagent
 
-import datadog.trace.api.DDTags
-import datadog.trace.bootstrap.instrumentation.api.Tags
 import datadog.communication.serialization.ByteBufferConsumer
 import datadog.communication.serialization.FlushingBuffer
 import datadog.communication.serialization.msgpack.MsgPackWriter
+import datadog.trace.api.DDTags
+import datadog.trace.bootstrap.instrumentation.api.Tags
+import datadog.trace.core.DDSpanContext
 import datadog.trace.test.util.DDSpecification
 import org.junit.Assert
 import org.msgpack.core.MessageFormat
@@ -17,18 +18,8 @@ import java.nio.channels.WritableByteChannel
 import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.DD_MEASURED
 import static datadog.trace.common.writer.ddagent.TraceGenerator.generateRandomTraces
 import static org.junit.Assert.assertEquals
-import static org.msgpack.core.MessageFormat.FLOAT32
-import static org.msgpack.core.MessageFormat.FLOAT64
-import static org.msgpack.core.MessageFormat.INT16
-import static org.msgpack.core.MessageFormat.INT32
-import static org.msgpack.core.MessageFormat.INT64
-import static org.msgpack.core.MessageFormat.INT8
-import static org.msgpack.core.MessageFormat.NEGFIXINT
-import static org.msgpack.core.MessageFormat.POSFIXINT
-import static org.msgpack.core.MessageFormat.UINT16
-import static org.msgpack.core.MessageFormat.UINT32
-import static org.msgpack.core.MessageFormat.UINT64
-import static org.msgpack.core.MessageFormat.UINT8
+import static org.junit.Assert.assertFalse
+import static org.msgpack.core.MessageFormat.*
 
 class TraceMapperV04PayloadTest extends DDSpecification {
 
@@ -174,6 +165,13 @@ class TraceMapperV04PayloadTest extends DDSpecification {
               }
               if (DD_MEASURED.toString() == key) {
                 assert ((n == 1) && expectedSpan.isMeasured()) || !expectedSpan.isMeasured()
+              } else if (DDSpanContext.PRIORITY_SAMPLING_KEY == key) {
+                //check that priority sampling is only on first and last span
+                if (k == 0 || k == spanCount -1) {
+                  assertEquals(expectedSpan.samplingPriority(), n.intValue())
+                } else {
+                  assertFalse(expectedSpan.hasSamplingPriority())
+                }
               } else {
                 metrics.put(key, n)
               }