From bf47f6d5e9cdbc6993a725695ceaae0c048e1490 Mon Sep 17 00:00:00 2001 From: Anuraag Agrawal Date: Thu, 19 Aug 2021 17:11:16 +0900 Subject: [PATCH] First try rate limited sampler for initialSampler of XRay remote sampler. --- .../awsxray/AwsXrayRemoteSamplerBuilder.java | 10 +++- .../contrib/awsxray/OrElseSampler.java | 51 +++++++++++++++++++ .../awsxray/AwsXrayRemoteSamplerTest.java | 9 +++- .../contrib/awsxray/OrElseSamplerTest.java | 47 +++++++++++++++++ 4 files changed, 114 insertions(+), 3 deletions(-) create mode 100644 aws-xray/src/main/java/io/opentelemetry/contrib/awsxray/OrElseSampler.java create mode 100644 aws-xray/src/test/java/io/opentelemetry/contrib/awsxray/OrElseSamplerTest.java diff --git a/aws-xray/src/main/java/io/opentelemetry/contrib/awsxray/AwsXrayRemoteSamplerBuilder.java b/aws-xray/src/main/java/io/opentelemetry/contrib/awsxray/AwsXrayRemoteSamplerBuilder.java index d7490ab70..a33aacab1 100644 --- a/aws-xray/src/main/java/io/opentelemetry/contrib/awsxray/AwsXrayRemoteSamplerBuilder.java +++ b/aws-xray/src/main/java/io/opentelemetry/contrib/awsxray/AwsXrayRemoteSamplerBuilder.java @@ -11,6 +11,7 @@ import io.opentelemetry.sdk.trace.samplers.Sampler; import java.time.Duration; import java.util.concurrent.TimeUnit; +import org.checkerframework.checker.nullness.qual.Nullable; /** A builder for {@link AwsXrayRemoteSampler}. */ public final class AwsXrayRemoteSamplerBuilder { @@ -22,7 +23,7 @@ public final class AwsXrayRemoteSamplerBuilder { private Clock clock = Clock.getDefault(); private String endpoint = DEFAULT_ENDPOINT; - private Sampler initialSampler = Sampler.parentBased(Sampler.traceIdRatioBased(0.05)); + @Nullable private Sampler initialSampler; private long pollingIntervalNanos = TimeUnit.SECONDS.toNanos(DEFAULT_POLLING_INTERVAL_SECS); AwsXrayRemoteSamplerBuilder(Resource resource) { @@ -84,6 +85,13 @@ public AwsXrayRemoteSamplerBuilder setClock(Clock clock) { /** Returns a {@link AwsXrayRemoteSampler} with the configuration of this builder. */ public AwsXrayRemoteSampler build() { + Sampler initialSampler = this.initialSampler; + if (initialSampler == null) { + initialSampler = + Sampler.parentBased( + new OrElseSampler( + new RateLimitingSampler(1, clock), Sampler.traceIdRatioBased(0.05))); + } return new AwsXrayRemoteSampler( resource, clock, endpoint, initialSampler, pollingIntervalNanos); } diff --git a/aws-xray/src/main/java/io/opentelemetry/contrib/awsxray/OrElseSampler.java b/aws-xray/src/main/java/io/opentelemetry/contrib/awsxray/OrElseSampler.java new file mode 100644 index 000000000..713ef82da --- /dev/null +++ b/aws-xray/src/main/java/io/opentelemetry/contrib/awsxray/OrElseSampler.java @@ -0,0 +1,51 @@ +/* + * Copyright The OpenTelemetry Authors + * SPDX-License-Identifier: Apache-2.0 + */ +package io.opentelemetry.contrib.awsxray; + +import io.opentelemetry.api.common.Attributes; +import io.opentelemetry.api.trace.SpanKind; +import io.opentelemetry.context.Context; +import io.opentelemetry.sdk.trace.data.LinkData; +import io.opentelemetry.sdk.trace.samplers.Sampler; +import io.opentelemetry.sdk.trace.samplers.SamplingDecision; +import io.opentelemetry.sdk.trace.samplers.SamplingResult; +import java.util.List; + +class OrElseSampler implements Sampler { + + private final Sampler first; + private final Sampler second; + + OrElseSampler(Sampler first, Sampler second) { + this.first = first; + this.second = second; + } + + @Override + public SamplingResult shouldSample( + Context parentContext, + String traceId, + String name, + SpanKind spanKind, + Attributes attributes, + List parentLinks) { + SamplingResult result = + first.shouldSample(parentContext, traceId, name, spanKind, attributes, parentLinks); + if (result.getDecision() != SamplingDecision.DROP) { + return result; + } + return second.shouldSample(parentContext, traceId, name, spanKind, attributes, parentLinks); + } + + @Override + public String getDescription() { + return "OrElse{" + + "first:" + + first.getDescription() + + ", second:" + + second.getDescription() + + "}"; + } +} diff --git a/aws-xray/src/test/java/io/opentelemetry/contrib/awsxray/AwsXrayRemoteSamplerTest.java b/aws-xray/src/test/java/io/opentelemetry/contrib/awsxray/AwsXrayRemoteSamplerTest.java index 008aeaa7d..a1a9d9565 100644 --- a/aws-xray/src/test/java/io/opentelemetry/contrib/awsxray/AwsXrayRemoteSamplerTest.java +++ b/aws-xray/src/test/java/io/opentelemetry/contrib/awsxray/AwsXrayRemoteSamplerTest.java @@ -236,7 +236,12 @@ void getAndUpdate() throws Exception { @Test void defaultInitialSampler() { - assertThat(AwsXrayRemoteSampler.newBuilder(Resource.empty()).build().getDescription()) - .startsWith("AwsXrayRemoteSampler{ParentBased{root:TraceIdRatioBased{0.050000}"); + try (AwsXrayRemoteSampler sampler = AwsXrayRemoteSampler.newBuilder(Resource.empty()).build()) { + assertThat(sampler.getDescription()) + .startsWith( + "AwsXrayRemoteSampler{" + + "ParentBased{root:OrElse{" + + "first:RateLimitingSampler{1}, second:TraceIdRatioBased{0.050000}"); + } } } diff --git a/aws-xray/src/test/java/io/opentelemetry/contrib/awsxray/OrElseSamplerTest.java b/aws-xray/src/test/java/io/opentelemetry/contrib/awsxray/OrElseSamplerTest.java new file mode 100644 index 000000000..b012d5237 --- /dev/null +++ b/aws-xray/src/test/java/io/opentelemetry/contrib/awsxray/OrElseSamplerTest.java @@ -0,0 +1,47 @@ +/* + * Copyright The OpenTelemetry Authors + * SPDX-License-Identifier: Apache-2.0 + */ +package io.opentelemetry.contrib.awsxray; + +import static org.assertj.core.api.Assertions.assertThat; + +import io.opentelemetry.api.common.Attributes; +import io.opentelemetry.api.trace.SpanKind; +import io.opentelemetry.api.trace.TraceId; +import io.opentelemetry.context.Context; +import io.opentelemetry.sdk.trace.samplers.Sampler; +import io.opentelemetry.sdk.trace.samplers.SamplingDecision; +import io.opentelemetry.sdk.trace.samplers.SamplingResult; +import java.util.Collections; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.junit.jupiter.MockitoExtension; + +@ExtendWith(MockitoExtension.class) +class OrElseSamplerTest { + + @Test + void firstWins() { + Sampler sampler = new OrElseSampler(Sampler.alwaysOn(), Sampler.alwaysOff()); + assertThat(doSample(sampler).getDecision()).isEqualTo(SamplingDecision.RECORD_AND_SAMPLE); + } + + @Test + void fallsBackToSecond() { + Sampler sampler = new OrElseSampler(Sampler.alwaysOff(), Sampler.alwaysOn()); + assertThat(doSample(sampler).getDecision()).isEqualTo(SamplingDecision.RECORD_AND_SAMPLE); + sampler = new OrElseSampler(Sampler.alwaysOff(), Sampler.alwaysOff()); + assertThat(doSample(sampler).getDecision()).isEqualTo(SamplingDecision.DROP); + } + + private SamplingResult doSample(Sampler sampler) { + return sampler.shouldSample( + Context.current(), + TraceId.fromLongs(1, 2), + "span", + SpanKind.CLIENT, + Attributes.empty(), + Collections.emptyList()); + } +}