Skip to content

Commit

Permalink
RUMM-2165: Respect sampling decision already made in the okhttp chain
Browse files Browse the repository at this point in the history
  • Loading branch information
0xnm committed May 19, 2022
1 parent efafff3 commit 2a08eae
Show file tree
Hide file tree
Showing 7 changed files with 298 additions and 1 deletion.
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@ import com.datadog.android.tracing.internal.TracingFeature
import com.datadog.opentracing.DDTracer
import com.datadog.trace.api.DDTags
import com.datadog.trace.api.interceptor.MutableSpan
import com.datadog.trace.api.sampling.PrioritySampling
import io.opentracing.Span
import io.opentracing.SpanContext
import io.opentracing.Tracer
Expand Down Expand Up @@ -195,7 +196,8 @@ internal constructor(
request: Request,
tracer: Tracer
): Response {
val span = if (traceSampler.sample()) {
val isSampled = extractSamplingDecision(request) ?: traceSampler.sample()
val span = if (isSampled) {
buildSpan(tracer, request)
} else {
null
Expand Down Expand Up @@ -275,6 +277,14 @@ internal constructor(
return span
}

private fun extractSamplingDecision(request: Request): Boolean? {
val samplingPriority = request.header(SAMPLING_PRIORITY_HEADER)?.toIntOrNull()
?: return null
if (samplingPriority == PrioritySampling.UNSET) return null
return samplingPriority == PrioritySampling.USER_KEEP ||
samplingPriority == PrioritySampling.SAMPLER_KEEP
}

private fun extractParentContext(tracer: Tracer, request: Request): SpanContext? {
val tagContext = request.tag(Span::class.java)?.context()

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@ import com.datadog.tools.unit.extensions.TestConfigurationExtension
import com.datadog.tools.unit.extensions.config.TestConfiguration
import com.datadog.tools.unit.setStaticValue
import com.datadog.trace.api.interceptor.MutableSpan
import com.datadog.trace.api.sampling.PrioritySampling
import com.nhaarman.mockitokotlin2.any
import com.nhaarman.mockitokotlin2.anyOrNull
import com.nhaarman.mockitokotlin2.argumentCaptor
Expand Down Expand Up @@ -368,6 +369,76 @@ internal open class TracingInterceptorNonDdTracerNotSendingSpanTest {
verify(mockSpanBuilder).asChildOf(parentSpanContext)
}

@Test
fun `𝕄 respect sampling decision 𝕎 intercept() {sampled in upstream interceptor}`(
@IntForgery(min = 200, max = 600) statusCode: Int,
@StringForgery key: String,
@StringForgery(type = StringForgeryType.ALPHA_NUMERICAL) value: String,
forge: Forge
) {
// Given
whenever(mockDetector.isFirstPartyUrl(HttpUrl.get(fakeUrl))).thenReturn(true)
fakeRequest = forgeRequest(forge) {
it.addHeader(
TracingInterceptor.SAMPLING_PRIORITY_HEADER,
forge.anElementFrom(
PrioritySampling.SAMPLER_KEEP.toString(),
PrioritySampling.USER_KEEP.toString()
)
)
}
stubChain(mockChain, statusCode)
whenever(mockTraceSampler.sample()).thenReturn(false)
doAnswer { invocation ->
val carrier = invocation.arguments[2] as TextMapInject
carrier.put(key, value)
}.whenever(mockTracer).inject<TextMapInject>(any(), any(), any())

// When
val response = testedInterceptor.intercept(mockChain)

// Then
assertThat(response).isSameAs(fakeResponse)
argumentCaptor<Request> {
verify(mockChain).proceed(capture())
assertThat(lastValue.header(key)).isEqualTo(value)
}
}

@Test
fun `𝕄 respect sampling decision 𝕎 intercept() {sampled out in upstream interceptor}`(
@IntForgery(min = 200, max = 600) statusCode: Int,
forge: Forge
) {
// Given
whenever(mockDetector.isFirstPartyUrl(HttpUrl.get(fakeUrl))).thenReturn(true)
fakeRequest = forgeRequest(forge) {
it.addHeader(
TracingInterceptor.SAMPLING_PRIORITY_HEADER,
forge.anElementFrom(
PrioritySampling.SAMPLER_DROP.toString(),
PrioritySampling.USER_DROP.toString()
)
)
}
stubChain(mockChain, statusCode)
whenever(mockTraceSampler.sample()).thenReturn(true)

// When
val response = testedInterceptor.intercept(mockChain)

// Then
assertThat(response).isSameAs(fakeResponse)
argumentCaptor<Request> {
verify(mockChain).proceed(capture())
assertThat(lastValue.header(TracingInterceptor.SAMPLING_PRIORITY_HEADER))
.isEqualTo("0")
assertThat(lastValue.header(TracingInterceptor.TRACE_ID_HEADER)).isNull()
assertThat(lastValue.header(TracingInterceptor.SPAN_ID_HEADER)).isNull()
}
verifyZeroInteractions(mockTracer, mockLocalTracer, mockSpanBuilder, mockSpan)
}

@Test
fun `𝕄 not create a span 𝕎 intercept() { not sampled }`(
@IntForgery(min = 200, max = 600) statusCode: Int
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@ import com.datadog.tools.unit.extensions.TestConfigurationExtension
import com.datadog.tools.unit.extensions.config.TestConfiguration
import com.datadog.tools.unit.setStaticValue
import com.datadog.trace.api.interceptor.MutableSpan
import com.datadog.trace.api.sampling.PrioritySampling
import com.nhaarman.mockitokotlin2.any
import com.nhaarman.mockitokotlin2.anyOrNull
import com.nhaarman.mockitokotlin2.argumentCaptor
Expand Down Expand Up @@ -461,6 +462,76 @@ internal open class TracingInterceptorNonDdTracerTest {
verify(mockSpanBuilder).asChildOf(parentSpanContext)
}

@Test
fun `𝕄 respect sampling decision 𝕎 intercept() {sampled in upstream interceptor}`(
@IntForgery(min = 200, max = 600) statusCode: Int,
@StringForgery key: String,
@StringForgery(type = StringForgeryType.ALPHA_NUMERICAL) value: String,
forge: Forge
) {
// Given
whenever(mockDetector.isFirstPartyUrl(HttpUrl.get(fakeUrl))).thenReturn(true)
fakeRequest = forgeRequest(forge) {
it.addHeader(
TracingInterceptor.SAMPLING_PRIORITY_HEADER,
forge.anElementFrom(
PrioritySampling.SAMPLER_KEEP.toString(),
PrioritySampling.USER_KEEP.toString()
)
)
}
stubChain(mockChain, statusCode)
whenever(mockTraceSampler.sample()).thenReturn(false)
doAnswer { invocation ->
val carrier = invocation.arguments[2] as TextMapInject
carrier.put(key, value)
}.whenever(mockTracer).inject<TextMapInject>(any(), any(), any())

// When
val response = testedInterceptor.intercept(mockChain)

// Then
assertThat(response).isSameAs(fakeResponse)
argumentCaptor<Request> {
verify(mockChain).proceed(capture())
assertThat(lastValue.header(key)).isEqualTo(value)
}
}

@Test
fun `𝕄 respect sampling decision 𝕎 intercept() {sampled out in upstream interceptor}`(
@IntForgery(min = 200, max = 600) statusCode: Int,
forge: Forge
) {
// Given
whenever(mockDetector.isFirstPartyUrl(HttpUrl.get(fakeUrl))).thenReturn(true)
fakeRequest = forgeRequest(forge) {
it.addHeader(
TracingInterceptor.SAMPLING_PRIORITY_HEADER,
forge.anElementFrom(
PrioritySampling.SAMPLER_DROP.toString(),
PrioritySampling.USER_DROP.toString()
)
)
}
stubChain(mockChain, statusCode)
whenever(mockTraceSampler.sample()).thenReturn(true)

// When
val response = testedInterceptor.intercept(mockChain)

// Then
assertThat(response).isSameAs(fakeResponse)
argumentCaptor<Request> {
verify(mockChain).proceed(capture())
assertThat(lastValue.header(TracingInterceptor.SAMPLING_PRIORITY_HEADER))
.isEqualTo("0")
assertThat(lastValue.header(TracingInterceptor.TRACE_ID_HEADER)).isNull()
assertThat(lastValue.header(TracingInterceptor.SPAN_ID_HEADER)).isNull()
}
verifyZeroInteractions(mockTracer, mockLocalTracer, mockSpanBuilder, mockSpan)
}

@Test
fun `𝕄 not create a span 𝕎 intercept() { not sampled }`(
@IntForgery(min = 200, max = 600) statusCode: Int
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@ import com.datadog.tools.unit.extensions.TestConfigurationExtension
import com.datadog.tools.unit.extensions.config.TestConfiguration
import com.datadog.tools.unit.setStaticValue
import com.datadog.trace.api.interceptor.MutableSpan
import com.datadog.trace.api.sampling.PrioritySampling
import com.nhaarman.mockitokotlin2.any
import com.nhaarman.mockitokotlin2.anyOrNull
import com.nhaarman.mockitokotlin2.argumentCaptor
Expand Down Expand Up @@ -371,6 +372,76 @@ internal open class TracingInterceptorNotSendingSpanTest {
verify(mockSpanBuilder).withOrigin(getExpectedOrigin())
}

@Test
fun `𝕄 respect sampling decision 𝕎 intercept() {sampled in upstream interceptor}`(
@IntForgery(min = 200, max = 600) statusCode: Int,
@StringForgery key: String,
@StringForgery(type = StringForgeryType.ALPHA_NUMERICAL) value: String,
forge: Forge
) {
// Given
whenever(mockDetector.isFirstPartyUrl(HttpUrl.get(fakeUrl))).thenReturn(true)
fakeRequest = forgeRequest(forge) {
it.addHeader(
TracingInterceptor.SAMPLING_PRIORITY_HEADER,
forge.anElementFrom(
PrioritySampling.SAMPLER_KEEP.toString(),
PrioritySampling.USER_KEEP.toString()
)
)
}
stubChain(mockChain, statusCode)
whenever(mockTraceSampler.sample()).thenReturn(false)
doAnswer { invocation ->
val carrier = invocation.arguments[2] as TextMapInject
carrier.put(key, value)
}.whenever(mockTracer).inject<TextMapInject>(any(), any(), any())

// When
val response = testedInterceptor.intercept(mockChain)

// Then
assertThat(response).isSameAs(fakeResponse)
argumentCaptor<Request> {
verify(mockChain).proceed(capture())
assertThat(lastValue.header(key)).isEqualTo(value)
}
}

@Test
fun `𝕄 respect sampling decision 𝕎 intercept() {sampled out in upstream interceptor}`(
@IntForgery(min = 200, max = 600) statusCode: Int,
forge: Forge
) {
// Given
whenever(mockDetector.isFirstPartyUrl(HttpUrl.get(fakeUrl))).thenReturn(true)
fakeRequest = forgeRequest(forge) {
it.addHeader(
TracingInterceptor.SAMPLING_PRIORITY_HEADER,
forge.anElementFrom(
PrioritySampling.SAMPLER_DROP.toString(),
PrioritySampling.USER_DROP.toString()
)
)
}
stubChain(mockChain, statusCode)
whenever(mockTraceSampler.sample()).thenReturn(true)

// When
val response = testedInterceptor.intercept(mockChain)

// Then
assertThat(response).isSameAs(fakeResponse)
argumentCaptor<Request> {
verify(mockChain).proceed(capture())
assertThat(lastValue.header(TracingInterceptor.SAMPLING_PRIORITY_HEADER))
.isEqualTo("0")
assertThat(lastValue.header(TracingInterceptor.TRACE_ID_HEADER)).isNull()
assertThat(lastValue.header(TracingInterceptor.SPAN_ID_HEADER)).isNull()
}
verifyZeroInteractions(mockTracer, mockLocalTracer, mockSpanBuilder, mockSpan)
}

@Test
fun `𝕄 not create a span 𝕎 intercept() { not sampled }`(
@IntForgery(min = 200, max = 600) statusCode: Int
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@ import com.datadog.tools.unit.extensions.TestConfigurationExtension
import com.datadog.tools.unit.extensions.config.TestConfiguration
import com.datadog.tools.unit.setStaticValue
import com.datadog.trace.api.interceptor.MutableSpan
import com.datadog.trace.api.sampling.PrioritySampling
import com.nhaarman.mockitokotlin2.any
import com.nhaarman.mockitokotlin2.anyOrNull
import com.nhaarman.mockitokotlin2.argumentCaptor
Expand Down Expand Up @@ -465,6 +466,76 @@ internal open class TracingInterceptorTest {
verify(mockSpanBuilder).withOrigin(getExpectedOrigin())
}

@Test
fun `𝕄 respect sampling decision 𝕎 intercept() {sampled in upstream interceptor}`(
@IntForgery(min = 200, max = 600) statusCode: Int,
@StringForgery key: String,
@StringForgery(type = StringForgeryType.ALPHA_NUMERICAL) value: String,
forge: Forge
) {
// Given
whenever(mockDetector.isFirstPartyUrl(HttpUrl.get(fakeUrl))).thenReturn(true)
fakeRequest = forgeRequest(forge) {
it.addHeader(
TracingInterceptor.SAMPLING_PRIORITY_HEADER,
forge.anElementFrom(
PrioritySampling.SAMPLER_KEEP.toString(),
PrioritySampling.USER_KEEP.toString()
)
)
}
stubChain(mockChain, statusCode)
whenever(mockTraceSampler.sample()).thenReturn(false)
doAnswer { invocation ->
val carrier = invocation.arguments[2] as TextMapInject
carrier.put(key, value)
}.whenever(mockTracer).inject<TextMapInject>(any(), any(), any())

// When
val response = testedInterceptor.intercept(mockChain)

// Then
assertThat(response).isSameAs(fakeResponse)
argumentCaptor<Request> {
verify(mockChain).proceed(capture())
assertThat(lastValue.header(key)).isEqualTo(value)
}
}

@Test
fun `𝕄 respect sampling decision 𝕎 intercept() {sampled out in upstream interceptor}`(
@IntForgery(min = 200, max = 600) statusCode: Int,
forge: Forge
) {
// Given
whenever(mockDetector.isFirstPartyUrl(HttpUrl.get(fakeUrl))).thenReturn(true)
fakeRequest = forgeRequest(forge) {
it.addHeader(
TracingInterceptor.SAMPLING_PRIORITY_HEADER,
forge.anElementFrom(
PrioritySampling.SAMPLER_DROP.toString(),
PrioritySampling.USER_DROP.toString()
)
)
}
stubChain(mockChain, statusCode)
whenever(mockTraceSampler.sample()).thenReturn(true)

// When
val response = testedInterceptor.intercept(mockChain)

// Then
assertThat(response).isSameAs(fakeResponse)
argumentCaptor<Request> {
verify(mockChain).proceed(capture())
assertThat(lastValue.header(TracingInterceptor.SAMPLING_PRIORITY_HEADER))
.isEqualTo("0")
assertThat(lastValue.header(TracingInterceptor.TRACE_ID_HEADER)).isNull()
assertThat(lastValue.header(TracingInterceptor.SPAN_ID_HEADER)).isNull()
}
verifyZeroInteractions(mockTracer, mockLocalTracer, mockSpanBuilder, mockSpan)
}

@Test
fun `𝕄 not create a span 𝕎 intercept() { not sampled }`(
@IntForgery(min = 200, max = 600) statusCode: Int
Expand Down
1 change: 1 addition & 0 deletions detekt.yml
Original file line number Diff line number Diff line change
Expand Up @@ -1074,6 +1074,7 @@ datadog:
- "kotlin.String.substringBefore(kotlin.Char, kotlin.String)"
- "kotlin.String.toByteArray(java.nio.charset.Charset) "
- "kotlin.String.toByteArray(java.nio.charset.Charset)"
- "kotlin.String.toIntOrNull()"
- "kotlin.String.toDoubleOrNull()"
- "kotlin.String.toLongOrNull()"
- "kotlin.String.toMethod()"
Expand Down
2 changes: 2 additions & 0 deletions docs/trace_collection.md
Original file line number Diff line number Diff line change
Expand Up @@ -482,6 +482,8 @@ final OkHttpClient okHttpClient = new OkHttpClient.Builder()
{{% /tab %}}
{{< /tabs >}}

In this case trace sampling decision made by the upstream interceptor for a particular request will be respected by the downstream interceptor.

Because the way the OkHttp Request is executed (using a Thread pool), the request span won't be automatically linked with the span that triggered the request. You can manually provide a parent span in the `OkHttp Request.Builder` as follows:

{{< tabs >}}
Expand Down

0 comments on commit 2a08eae

Please sign in to comment.