Priority Sampling flag in DDSpanContext
MODIFY
dd-trace-ot/src/main/java/datadog/opentracing/DDSpan.java
dd-trace-ot/src/main/java/datadog/opentracing/DDSpan.java
@@ -2,7 +2,10 @@ package datadog.opentracing;
import com.fasterxml.jackson.annotation.JsonGetter;
import com.fasterxml.jackson.annotation.JsonIgnore;
+import com.fasterxml.jackson.annotation.JsonInclude;
+import com.fasterxml.jackson.annotation.JsonInclude.Include;
import datadog.trace.api.DDTags;
+import datadog.trace.common.sampling.PrioritySampling;
import datadog.trace.common.util.Clock;
import io.opentracing.Span;
import java.io.PrintWriter;
@@ -236,6 +239,16 @@ public class DDSpan implements Span {
return this;
}
+ /**
+ * Set the sampling priority of the span.
+ *
+ * <p>Has no effect if the span priority has been propagated (injected or extracted).
+ */
+ public final DDSpan setSamplingPriority(int newPriority) {
+ this.context().setSamplingPriority(newPriority);
+ return this;
+ }
+
public final DDSpan setSpanType(final String type) {
this.context().setSpanType(type);
return this;
@@ -300,6 +313,17 @@ public class DDSpan implements Span {
return context.getOperationName();
}
+ @JsonGetter("sampling_priority")
+ @JsonInclude(Include.NON_NULL)
+ public Integer getSamplingPriority() {
+ final int samplingPriority = context.getSamplingPriority();
+ if (samplingPriority == PrioritySampling.UNSET) {
+ return null;
+ } else {
+ return samplingPriority;
+ }
+ }
+
@JsonIgnore
public Map<String, Object> getTags() {
return this.context().getTags();

MODIFY
dd-trace-ot/src/main/java/datadog/opentracing/DDSpanContext.java
dd-trace-ot/src/main/java/datadog/opentracing/DDSpanContext.java
@@ -4,6 +4,7 @@ import com.fasterxml.jackson.annotation.JsonIgnore;
import com.google.common.collect.Maps;
import datadog.opentracing.decorators.AbstractDecorator;
import datadog.trace.api.DDTags;
+import datadog.trace.common.sampling.PrioritySampling;
import io.opentracing.tag.Tags;
import java.util.Collections;
import java.util.HashMap;
@@ -49,6 +50,10 @@ public class DDSpanContext implements io.opentracing.SpanContext {
private String spanType;
/** Each span have an operation name describing the current span */
private String operationName;
+ /** The sampling priority of the trace */
+ private int samplingPriority = PrioritySampling.UNSET;
+ /** When true, the samplingPriority cannot be changed. */
+ private boolean samplingPriorityLocked = false;
// Others attributes
/** Tags are associated to the current span, they will not propagate to the children span */
private Map<String, Object> tags;
@@ -60,6 +65,7 @@ public class DDSpanContext implements io.opentracing.SpanContext {
final String serviceName,
final String operationName,
final String resourceName,
+ final int samplingPriority,
final Map<String, String> baggageItems,
final boolean errorFlag,
final String spanType,
@@ -80,6 +86,7 @@ public class DDSpanContext implements io.opentracing.SpanContext {
this.serviceName = serviceName;
this.operationName = operationName;
this.resourceName = resourceName;
+ this.samplingPriority = samplingPriority;
this.errorFlag = errorFlag;
this.spanType = spanType;
@@ -141,6 +148,38 @@ public class DDSpanContext implements io.opentracing.SpanContext {
this.spanType = spanType;
}
+ public synchronized void setSamplingPriority(int newPriority) {
+ if (samplingPriorityLocked) {
+ log.warn(
+ "samplingPriority locked at {}. Refusing to set to {}", samplingPriority, newPriority);
+ } else {
+ this.samplingPriority = newPriority;
+ }
+ }
+
+ public synchronized int getSamplingPriority() {
+ return samplingPriority;
+ }
+
+ /**
+ * Prevent future changes to the context's sampling priority.
+ *
+ * <p>Used when a span is extracted or injected for propagation.
+ *
+ * <p>Has no effect if the sampling priority is unset.
+ *
+ * @return true if the sampling priority was locked.
+ */
+ public synchronized boolean lockSamplingPriority() {
+ if (samplingPriority == PrioritySampling.UNSET) {
+ log.debug("{} : refusing to lock unset samplingPriority", this);
+ } else {
+ this.samplingPriorityLocked = true;
+ log.debug("{} : locked samplingPriority to {}", this, this.samplingPriority);
+ }
+ return samplingPriorityLocked;
+ }
+
public void setBaggageItem(final String key, final String value) {
if (this.baggageItems.isEmpty()) {
this.baggageItems = new HashMap<>();
@@ -248,6 +287,9 @@ public class DDSpanContext implements io.opentracing.SpanContext {
.append(getOperationName())
.append("/")
.append(getResourceName());
+ if (getSamplingPriority() != PrioritySampling.UNSET) {
+ s.append(" samplingPriority=").append(getSamplingPriority());
+ }
if (errorFlag) {
s.append(" *errored*");
}

MODIFY
dd-trace-ot/src/main/java/datadog/opentracing/DDTracer.java
dd-trace-ot/src/main/java/datadog/opentracing/DDTracer.java
@@ -9,7 +9,11 @@ import datadog.trace.api.DDTags;
import datadog.trace.common.DDTraceConfig;
import datadog.trace.common.Service;
import datadog.trace.common.sampling.AllSampler;
+import datadog.trace.common.sampling.PrioritySampling;
+import datadog.trace.common.sampling.RateByServiceSampler;
import datadog.trace.common.sampling.Sampler;
+import datadog.trace.common.writer.DDAgentWriter;
+import datadog.trace.common.writer.DDApi;
import datadog.trace.common.writer.Writer;
import io.opentracing.Scope;
import io.opentracing.ScopeManager;
@@ -39,6 +43,8 @@ public class DDTracer extends ThreadLocalScopeManager implements io.opentracing.
final Writer writer;
/** Sampler defines the sampling policy in order to reduce the number of traces for instance */
final Sampler sampler;
+ /** Sampler which rates based on the service name */
+ final RateByServiceSampler serviceSampler;
/** Span context decorators */
private final Map<String, List<AbstractDecorator>> spanContextDecorators = new HashMap<>();
@@ -59,7 +65,8 @@ public class DDTracer extends ThreadLocalScopeManager implements io.opentracing.
this(
config.getProperty(DDTraceConfig.SERVICE_NAME),
Writer.Builder.forConfig(config),
- Sampler.Builder.forConfig(config));
+ Sampler.Builder.forConfig(config),
+ RateByServiceSampler.Builder.forConfig(config));
log.debug("Using config: {}", config);
// Create decorators from resource files
@@ -71,13 +78,26 @@ public class DDTracer extends ThreadLocalScopeManager implements io.opentracing.
}
public DDTracer(final String serviceName, final Writer writer, final Sampler sampler) {
+ this(serviceName, writer, sampler, null);
+ }
+
+ public DDTracer(
+ final String serviceName,
+ final Writer writer,
+ final Sampler sampler,
+ RateByServiceSampler serviceSampler) {
this.serviceName = serviceName;
this.writer = writer;
this.writer.start();
this.sampler = sampler;
+ this.serviceSampler = serviceSampler;
registry = new CodecRegistry();
registry.register(Format.Builtin.HTTP_HEADERS, new HTTPCodec());
registry.register(Format.Builtin.TEXT_MAP, new HTTPCodec());
+ if (this.writer instanceof DDAgentWriter && serviceSampler != null) {
+ final DDApi api = ((DDAgentWriter) this.writer).getApi();
+ api.addResponseListener(this.serviceSampler);
+ }
log.info("New instance: {}", this);
}
@@ -139,7 +159,6 @@ public class DDTracer extends ThreadLocalScopeManager implements io.opentracing.
@Override
public <T> SpanContext extract(final Format<T> format, final T carrier) {
-
final Codec<T> codec = registry.get(format);
if (codec == null) {
log.warn("Unsupported format for propagation - {}", format.getClass().getName());
@@ -159,7 +178,9 @@ public class DDTracer extends ThreadLocalScopeManager implements io.opentracing.
if (trace.isEmpty()) {
return;
}
- if (this.sampler.sample(trace.peek())) {
+ // If priority sampling is enabled, send all traces to the agent (even traces marked to drop).
+ // Otherwise, use the sampler to drop traces.
+ if (prioritySamplingEnabled() || this.sampler.sample(trace.peek())) {
this.writer.write(new ArrayList<>(trace));
}
}
@@ -208,6 +229,10 @@ public class DDTracer extends ThreadLocalScopeManager implements io.opentracing.
return services;
}
+ private boolean prioritySamplingEnabled() {
+ return null != serviceSampler;
+ }
+
private static class CodecRegistry {
private final Map<Format<?>, Codec<?>> codecs = new HashMap<>();
@@ -250,7 +275,28 @@ public class DDTracer extends ThreadLocalScopeManager implements io.opentracing.
}
private DDSpan startSpan() {
- return new DDSpan(this.timestamp, buildSpanContext());
+ final DDSpan span = new DDSpan(this.timestamp, buildSpanContext());
+ if (DDTracer.this.prioritySamplingEnabled()) {
+ if (span.isRootSpan()) {
+ // Run the priority sampler on the new span
+ if (DDTracer.this.sampler.sample(span) && DDTracer.this.serviceSampler.sample(span)) {
+ span.setSamplingPriority(PrioritySampling.SAMPLER_KEEP);
+ } else {
+ span.setSamplingPriority(PrioritySampling.SAMPLER_DROP);
+ }
+ } else {
+ if (span.getSamplingPriority() == null) {
+ // Edge case: If the parent context did not set the priority, run the priority sampler.
+ // Happens when extracted http context did not send the priority header.
+ if (DDTracer.this.sampler.sample(span) && DDTracer.this.serviceSampler.sample(span)) {
+ span.setSamplingPriority(PrioritySampling.SAMPLER_KEEP);
+ } else {
+ span.setSamplingPriority(PrioritySampling.SAMPLER_DROP);
+ }
+ }
+ }
+ }
+ return span;
}
@Override
@@ -373,6 +419,7 @@ public class DDTracer extends ThreadLocalScopeManager implements io.opentracing.
final long parentSpanId;
final Map<String, String> baggage;
final Queue<DDSpan> parentTrace;
+ final int samplingPriority;
final DDSpanContext context;
SpanContext parentContext = this.parent;
@@ -388,6 +435,7 @@ public class DDTracer extends ThreadLocalScopeManager implements io.opentracing.
parentSpanId = ddsc.getSpanId();
baggage = ddsc.getBaggageItems();
parentTrace = ddsc.getTrace();
+ samplingPriority = ddsc.getSamplingPriority();
if (this.serviceName == null) this.serviceName = ddsc.getServiceName();
if (this.spanType == null) this.spanType = ddsc.getSpanType();
@@ -396,6 +444,7 @@ public class DDTracer extends ThreadLocalScopeManager implements io.opentracing.
parentSpanId = 0L;
baggage = null;
parentTrace = null;
+ samplingPriority = PrioritySampling.UNSET;
}
if (serviceName == null) {
@@ -416,6 +465,7 @@ public class DDTracer extends ThreadLocalScopeManager implements io.opentracing.
serviceName,
operationName,
this.resourceName,
+ samplingPriority,
baggage,
errorFlag,
spanType,

MODIFY
dd-trace-ot/src/main/java/datadog/opentracing/propagation/HTTPCodec.java
dd-trace-ot/src/main/java/datadog/opentracing/propagation/HTTPCodec.java
package datadog.opentracing.propagation;
import datadog.opentracing.DDSpanContext;
+import datadog.trace.common.sampling.PrioritySampling;
import io.opentracing.propagation.TextMap;
import java.io.UnsupportedEncodingException;
import java.net.URLDecoder;
@@ -17,12 +18,15 @@ public class HTTPCodec implements Codec<TextMap> {
private static final String OT_BAGGAGE_PREFIX = "ot-baggage-";
private static final String TRACE_ID_KEY = "x-datadog-trace-id";
private static final String SPAN_ID_KEY = "x-datadog-parent-id";
+ private static final String SAMPLING_PRIORITY_KEY = "x-datadog-sampling-priority";
@Override
public void inject(final DDSpanContext context, final TextMap carrier) {
-
carrier.put(TRACE_ID_KEY, String.valueOf(context.getTraceId()));
carrier.put(SPAN_ID_KEY, String.valueOf(context.getSpanId()));
+ if (context.lockSamplingPriority()) {
+ carrier.put(SAMPLING_PRIORITY_KEY, String.valueOf(context.getSamplingPriority()));
+ }
for (final Map.Entry<String, String> entry : context.baggageItems()) {
carrier.put(OT_BAGGAGE_PREFIX + entry.getKey(), encode(entry.getValue()));
@@ -35,9 +39,9 @@ public class HTTPCodec implements Codec<TextMap> {
Map<String, String> baggage = Collections.emptyMap();
Long traceId = 0L;
Long spanId = 0L;
+ int samplingPriority = PrioritySampling.UNSET;
for (final Map.Entry<String, String> entry : carrier) {
-
final String key = entry.getKey().toLowerCase();
if (key.equalsIgnoreCase(TRACE_ID_KEY)) {
traceId = Long.parseLong(entry.getValue());
@@ -48,14 +52,28 @@ public class HTTPCodec implements Codec<TextMap> {
baggage = new HashMap<>();
}
baggage.put(key.replace(OT_BAGGAGE_PREFIX, ""), decode(entry.getValue()));
+ } else if (key.equalsIgnoreCase(SAMPLING_PRIORITY_KEY)) {
+ samplingPriority = Integer.parseInt(entry.getValue());
}
}
DDSpanContext context = null;
if (traceId != 0L) {
-
context =
new DDSpanContext(
- traceId, spanId, 0L, null, null, null, baggage, false, null, null, null, null);
+ traceId,
+ spanId,
+ 0L,
+ null,
+ null,
+ null,
+ samplingPriority,
+ baggage,
+ false,
+ null,
+ null,
+ null,
+ null);
+ context.lockSamplingPriority();
log.debug("{} - Parent context extracted", context);
}

MODIFY
dd-trace-ot/src/main/java/datadog/trace/common/writer/DDAgentWriter.java
dd-trace-ot/src/main/java/datadog/trace/common/writer/DDAgentWriter.java
@@ -147,6 +147,10 @@ public class DDAgentWriter implements Writer {
return "DDAgentWriter { api=" + api + " }";
}
+ public DDApi getApi() {
+ return api;
+ }
+
/** Infinite tasks blocking until some spans come in the blocking queue. */
class TracesSendingTask implements Runnable {

MODIFY
dd-trace-ot/src/test/groovy/datadog/opentracing/DDSpanSerializationTest.groovy
dd-trace-ot/src/test/groovy/datadog/opentracing/DDSpanSerializationTest.groovy
@@ -2,10 +2,14 @@ package datadog.opentracing
import com.fasterxml.jackson.databind.ObjectMapper
import com.google.common.collect.Maps
+import datadog.trace.api.DDTags
+import datadog.trace.common.sampling.PrioritySampling
import spock.lang.Specification
+import spock.lang.Unroll
class DDSpanSerializationTest extends Specification {
+ @Unroll
def "serialize spans"() throws Exception {
setup:
final Map<String, String> baggage = new HashMap<>()
@@ -21,6 +25,9 @@ class DDSpanSerializationTest extends Specification {
expected.put("name", "operation")
expected.put("duration", 33000)
expected.put("resource", "operation")
+ if (samplingPriority != PrioritySampling.UNSET) {
+ expected.put("sampling_priority", samplingPriority)
+ }
expected.put("start", 100000)
expected.put("span_id", 2l)
expected.put("parent_id", 0l)
@@ -34,6 +41,7 @@ class DDSpanSerializationTest extends Specification {
"service",
"operation",
null,
+ samplingPriority,
new HashMap<>(baggage),
false,
"type",
@@ -50,5 +58,10 @@ class DDSpanSerializationTest extends Specification {
expect:
serializer.readTree(serializer.writeValueAsString(span)) == serializer.readTree(serializer.writeValueAsString(expected))
+
+ where:
+ samplingPriority | _
+ PrioritySampling.SAMPLER_KEEP | _
+ PrioritySampling.UNSET | _
}
}

MODIFY
dd-trace-ot/src/test/groovy/datadog/opentracing/DDSpanTest.groovy
dd-trace-ot/src/test/groovy/datadog/opentracing/DDSpanTest.groovy
package datadog.opentracing
+import datadog.trace.common.sampling.PrioritySampling
import spock.lang.Specification
class DDSpanTest extends Specification {
@@ -14,6 +15,7 @@ class DDSpanTest extends Specification {
"fakeService",
"fakeOperation",
"fakeResource",
+ PrioritySampling.UNSET,
Collections.<String, String>emptyMap(),
false,
"fakeType",
@@ -42,6 +44,22 @@ class DDSpanTest extends Specification {
span.setSpanType("type")
then:
span.getType() == "type"
+
+ when:
+ span.setSamplingPriority(PrioritySampling.UNSET)
+ then:
+ span.getSamplingPriority() == null
+
+ when:
+ span.setSamplingPriority(PrioritySampling.SAMPLER_KEEP)
+ then:
+ span.getSamplingPriority() == PrioritySampling.SAMPLER_KEEP
+
+ when:
+ context.lockSamplingPriority()
+ span.setSamplingPriority(PrioritySampling.USER_KEEP)
+ then:
+ span.getSamplingPriority() == PrioritySampling.SAMPLER_KEEP
}
def "resource name equals operation name if null"() {

MODIFY
dd-trace-ot/src/test/groovy/datadog/opentracing/propagation/HTTPCodecTest.groovy
dd-trace-ot/src/test/groovy/datadog/opentracing/propagation/HTTPCodecTest.groovy
package datadog.opentracing.propagation
-import com.datadoghq.trace.DDSpanContext
+import datadog.opentracing.DDSpanContext
+import datadog.trace.common.sampling.PrioritySampling
import io.opentracing.propagation.TextMapExtractAdapter
import io.opentracing.propagation.TextMapInjectAdapter
import spock.lang.Shared
import spock.lang.Specification
+import spock.lang.Unroll
class HTTPCodecTest extends Specification {
@Shared
@@ -13,7 +15,10 @@ class HTTPCodecTest extends Specification {
private static final String TRACE_ID_KEY = "x-datadog-trace-id"
@Shared
private static final String SPAN_ID_KEY = "x-datadog-parent-id"
+ @Shared
+ private static final String SAMPLING_PRIORITY_KEY = "x-datadog-sampling-priority"
+ @Unroll
def "inject http headers"() {
setup:
final DDSpanContext mockedContext =
@@ -24,6 +29,7 @@ class HTTPCodecTest extends Specification {
"fakeService",
"fakeOperation",
"fakeResource",
+ samplingPriority,
new HashMap<String, String>() {
{
put("k1", "v1")
@@ -44,10 +50,17 @@ class HTTPCodecTest extends Specification {
expect:
carrier.get(TRACE_ID_KEY) == "1"
carrier.get(SPAN_ID_KEY) == "2"
+ carrier.get(SAMPLING_PRIORITY_KEY) == (samplingPriority == PrioritySampling.UNSET ? null : String.valueOf(samplingPriority))
carrier.get(OT_BAGGAGE_PREFIX + "k1") == "v1"
carrier.get(OT_BAGGAGE_PREFIX + "k2") == "v2"
+
+ where:
+ samplingPriority | _
+ PrioritySampling.UNSET | _
+ PrioritySampling.SAMPLER_KEEP | _
}
+ @Unroll
def "extract http headers"() {
setup:
final Map<String, String> actual =
@@ -60,6 +73,10 @@ class HTTPCodecTest extends Specification {
}
}
+ if (samplingPriority != PrioritySampling.UNSET) {
+ actual.put(SAMPLING_PRIORITY_KEY, String.valueOf(samplingPriority))
+ }
+
final HTTPCodec codec = new HTTPCodec()
final DDSpanContext context = codec.extract(new TextMapExtractAdapter(actual))
@@ -68,5 +85,11 @@ class HTTPCodecTest extends Specification {
context.getSpanId() == 2l
context.getBaggageItem("k1") == "v1"
context.getBaggageItem("k2") == "v2"
+ context.getSamplingPriority() == samplingPriority
+
+ where:
+ samplingPriority | _
+ PrioritySampling.UNSET | _
+ PrioritySampling.SAMPLER_KEEP | _
}
}

MODIFY
dd-trace-ot/src/test/groovy/datadog/trace/SpanFactory.groovy
dd-trace-ot/src/test/groovy/datadog/trace/SpanFactory.groovy
@@ -3,6 +3,7 @@ package datadog.trace
import datadog.opentracing.DDSpan
import datadog.opentracing.DDSpanContext
import datadog.opentracing.DDTracer
+import datadog.trace.common.sampling.PrioritySampling
class SpanFactory {
static newSpanOf(long timestampMicro) {
@@ -13,6 +14,7 @@ class SpanFactory {
"fakeService",
"fakeOperation",
"fakeResource",
+ PrioritySampling.UNSET,
Collections.emptyMap(),
false,
"fakeType",
@@ -30,6 +32,7 @@ class SpanFactory {
"fakeService",
"fakeOperation",
"fakeResource",
+ PrioritySampling.UNSET,
Collections.emptyMap(),
false,
"fakeType",

Priority Sampling flag in DDSpanContext