forked from opensearch-project/OpenSearch
-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
[Tracing Framework] Adding support of Head based Sampler based on tra…
…ce header & probability (opensearch-project#8918) Signed-off-by: Dev Agarwal <[email protected]>
- Loading branch information
1 parent
bb38ed4
commit 32aaf8f
Showing
10 changed files
with
382 additions
and
10 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
60 changes: 60 additions & 0 deletions
60
...ns/telemetry-otel/src/main/java/org/opensearch/telemetry/tracing/sampler/HeadSampler.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,60 @@ | ||
| /* | ||
| * SPDX-License-Identifier: Apache-2.0 | ||
| * | ||
| * The OpenSearch Contributors require contributions made to | ||
| * this file be licensed under the Apache-2.0 license or a | ||
| * compatible open source license. | ||
| */ | ||
|
|
||
| package org.opensearch.telemetry.tracing.sampler; | ||
|
|
||
| import io.opentelemetry.api.common.AttributeKey; | ||
| 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.SamplingResult; | ||
|
|
||
| import java.util.List; | ||
|
|
||
| /** | ||
| * HeadBased sampler | ||
| */ | ||
| public class HeadSampler implements Sampler { | ||
| private final Sampler defaultSampler; | ||
|
|
||
| /** | ||
| * Creates Head based sampler | ||
| * @param defaultSampler defaultSampler | ||
| */ | ||
| public HeadSampler(Sampler defaultSampler) { | ||
| this.defaultSampler = defaultSampler; | ||
| } | ||
|
|
||
| @Override | ||
| public SamplingResult shouldSample( | ||
| Context parentContext, | ||
| String traceId, | ||
| String name, | ||
| SpanKind spanKind, | ||
| Attributes attributes, | ||
| List<LinkData> parentLinks) { | ||
|
|
||
| if (attributes.get(AttributeKey.stringKey("trace")) != null && Boolean.parseBoolean(attributes.get(AttributeKey.stringKey("trace"))) ) { | ||
| return SamplingResult.recordAndSample(); | ||
| } else { | ||
| return defaultSampler.shouldSample(parentContext, traceId, name, spanKind, attributes, parentLinks); | ||
| } | ||
| } | ||
|
|
||
| @Override | ||
| public String getDescription() { | ||
| return "HeadSampler"; | ||
| } | ||
|
|
||
| @Override | ||
| public String toString() { | ||
| return getDescription(); | ||
| } | ||
| } |
75 changes: 75 additions & 0 deletions
75
...try-otel/src/main/java/org/opensearch/telemetry/tracing/sampler/ProbabilisticSampler.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,75 @@ | ||
| package org.opensearch.telemetry.tracing.sampler; | ||
|
|
||
| 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.SamplingResult; | ||
| import org.opensearch.telemetry.OTelTelemetrySettings; | ||
|
|
||
| import java.util.List; | ||
|
|
||
| /** | ||
| * ProbabilisticSampler implements a head-based sampling strategy based on provided settings. | ||
| */ | ||
| public class ProbabilisticSampler implements Sampler { | ||
| private Sampler defaultSampler; | ||
| private final OTelTelemetrySettings oTelTelemetrySettings; | ||
|
|
||
| private double tracerHeadSamplerSamplingRatio; | ||
|
|
||
| /** | ||
| * Constructor | ||
| * | ||
| * @param oTelTelemetrySettings Telemetry settings. | ||
| */ | ||
| public ProbabilisticSampler(OTelTelemetrySettings oTelTelemetrySettings) { | ||
| if (oTelTelemetrySettings == null) { | ||
| throw new IllegalArgumentException("oTelTelemetrySettings cannot be null"); | ||
| } | ||
| this.oTelTelemetrySettings = oTelTelemetrySettings; | ||
| this.tracerHeadSamplerSamplingRatio = oTelTelemetrySettings.getTracerHeadSamplerSamplingRatio(); | ||
| this.defaultSampler = createProbabilitySampler(tracerHeadSamplerSamplingRatio); | ||
| } | ||
|
|
||
| private Sampler createProbabilitySampler(double samplingRatio) { | ||
| return Sampler.traceIdRatioBased(samplingRatio); | ||
| } | ||
|
|
||
| Sampler getSampler() { | ||
| double newSamplingRatio = oTelTelemetrySettings.getTracerHeadSamplerSamplingRatio(); | ||
| if(isSamplingRatioChanged(newSamplingRatio)){ | ||
| this.tracerHeadSamplerSamplingRatio = newSamplingRatio; | ||
| defaultSampler = createProbabilitySampler(tracerHeadSamplerSamplingRatio); | ||
| } | ||
| return defaultSampler; | ||
| } | ||
|
|
||
| private boolean isSamplingRatioChanged(double newSamplingRatio) { | ||
| return Double.compare(this.tracerHeadSamplerSamplingRatio, newSamplingRatio)!=0; | ||
| } | ||
| double getTracerHeadSamplerSamplingRatio() { | ||
| return tracerHeadSamplerSamplingRatio; | ||
| } | ||
| @Override | ||
| public SamplingResult shouldSample( | ||
| Context parentContext, | ||
| String traceId, | ||
| String name, | ||
| SpanKind spanKind, | ||
| Attributes attributes, | ||
| List<LinkData> parentLinks) { | ||
| return getSampler().shouldSample(parentContext, traceId, name, spanKind, attributes, parentLinks); | ||
| } | ||
|
|
||
| @Override | ||
| public String getDescription() { | ||
| return "Probabilistic Sampler"; | ||
| } | ||
|
|
||
| @Override | ||
| public String toString() { | ||
| return getDescription(); | ||
| } | ||
| } |
12 changes: 12 additions & 0 deletions
12
...s/telemetry-otel/src/main/java/org/opensearch/telemetry/tracing/sampler/package-info.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,12 @@ | ||
| /* | ||
| * SPDX-License-Identifier: Apache-2.0 | ||
| * | ||
| * The OpenSearch Contributors require contributions made to | ||
| * this file be licensed under the Apache-2.0 license or a | ||
| * compatible open source license. | ||
| */ | ||
|
|
||
| /** | ||
| * This package contains classes needed for HeadBased Sampler. | ||
| */ | ||
| package org.opensearch.telemetry.tracing.sampler; |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
80 changes: 80 additions & 0 deletions
80
...lemetry-otel/src/test/java/org/opensearch/telemetry/tracing/sampler/HeadSamplerTests.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,80 @@ | ||
| /* | ||
| * SPDX-License-Identifier: Apache-2.0 | ||
| * | ||
| * The OpenSearch Contributors require contributions made to | ||
| * this file be licensed under the Apache-2.0 license or a | ||
| * compatible open source license. | ||
| */ | ||
|
|
||
| package org.opensearch.telemetry.tracing.sampler; | ||
|
|
||
| import io.opentelemetry.api.common.AttributeKey; | ||
| import io.opentelemetry.api.trace.SpanKind; | ||
| import org.opensearch.test.OpenSearchTestCase; | ||
| import io.opentelemetry.sdk.trace.samplers.Sampler; | ||
| import io.opentelemetry.context.Context; | ||
| import io.opentelemetry.sdk.trace.samplers.SamplingResult; | ||
|
|
||
| import static org.mockito.ArgumentMatchers.any; | ||
| import static org.mockito.ArgumentMatchers.anyString; | ||
| import static org.mockito.Mockito.mock; | ||
| import static org.mockito.Mockito.when; | ||
| import static org.mockito.Mockito.verify; | ||
| import static org.mockito.Mockito.never; | ||
| import io.opentelemetry.api.common.Attributes; | ||
|
|
||
| import java.util.Collections; | ||
|
|
||
| public class HeadSamplerTests extends OpenSearchTestCase { | ||
|
|
||
| public void testShouldSampleWithTraceAttributeAsTrue() { | ||
|
|
||
| // Create a mock default sampler | ||
| Sampler defaultSampler = mock(Sampler.class); | ||
| when(defaultSampler.shouldSample(any(), anyString(), anyString(), any(), any(), any())) | ||
| .thenReturn(SamplingResult.drop()); | ||
|
|
||
| // Create an instance of HeadSampler with the mock default sampler | ||
| HeadSampler headSampler = new HeadSampler(defaultSampler); | ||
|
|
||
| // Create a mock Context and Attributes | ||
| Context parentContext = mock(Context.class); | ||
| Attributes attributes = Attributes.of(AttributeKey.stringKey("trace"), "true"); | ||
|
|
||
| // Call shouldSample on HeadSampler | ||
| SamplingResult result = headSampler.shouldSample(parentContext, "traceId", "spanName", | ||
| SpanKind.INTERNAL, attributes, Collections.emptyList()); | ||
|
|
||
| assertEquals(SamplingResult.recordAndSample(), result); | ||
|
|
||
| // Verify that the default sampler's shouldSample method was not called | ||
| verify(defaultSampler, never()).shouldSample(any(), anyString(), anyString(), any(), any(), any()); | ||
| } | ||
|
|
||
|
|
||
| public void testShouldSampleWithoutTraceAttribute() { | ||
|
|
||
| // Create a mock default sampler | ||
| Sampler defaultSampler = mock(Sampler.class); | ||
| when(defaultSampler.shouldSample(any(), anyString(), anyString(), any(), any(), any())) | ||
| .thenReturn(SamplingResult.recordAndSample()); | ||
|
|
||
| // Create an instance of HeadSampler with the mock default sampler | ||
| HeadSampler headSampler = new HeadSampler(defaultSampler); | ||
|
|
||
| // Create a mock Context and Attributes | ||
| Context parentContext = mock(Context.class); | ||
| Attributes attributes = Attributes.empty(); | ||
|
|
||
| // Call shouldSample on HeadSampler | ||
| SamplingResult result = headSampler.shouldSample(parentContext, "traceId", "spanName", | ||
| SpanKind.INTERNAL, attributes, Collections.emptyList()); | ||
|
|
||
| // Verify that HeadSampler returned SamplingResult.recordAndSample() | ||
| assertEquals(SamplingResult.recordAndSample(), result); | ||
|
|
||
| // Verify that the default sampler's shouldSample method was called | ||
| verify(defaultSampler).shouldSample(any(), anyString(), anyString(), any(), any(), any()); | ||
| } | ||
|
|
||
| } |
Oops, something went wrong.