Skip to content

Commit cdad898

Browse files
majanjua-amznmajanjua-amzn
committed
Add AlwaysRecordSampler
1 parent 92d23b2 commit cdad898

File tree

3 files changed

+205
-0
lines changed

3 files changed

+205
-0
lines changed

‎CHANGELOG.md‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,8 @@
22

33
## Unreleased
44

5+
* Add AlwaysRecordSampler
6+
57
## Version 1.56.0 (2025-11-07)
68

79
### API

‎sdk/trace/src/main/java/io/opentelemetry/sdk/trace/internal/AlwaysRecordSampler.java‎

Lines changed: 91 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,91 @@
1+
/*
2+
* Copyright The OpenTelemetry Authors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
// Includes work from:
7+
// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
8+
// SPDX-License-Identifier: Apache-2.0
9+
10+
package io.opentelemetry.sdk.trace.internal;
11+
12+
import io.opentelemetry.api.common.Attributes;
13+
import io.opentelemetry.api.trace.SpanKind;
14+
import io.opentelemetry.api.trace.TraceState;
15+
import io.opentelemetry.context.Context;
16+
import io.opentelemetry.sdk.trace.data.LinkData;
17+
import io.opentelemetry.sdk.trace.samplers.Sampler;
18+
import io.opentelemetry.sdk.trace.samplers.SamplingDecision;
19+
import io.opentelemetry.sdk.trace.samplers.SamplingResult;
20+
import java.util.List;
21+
import javax.annotation.concurrent.Immutable;
22+
23+
/**
24+
* This sampler will return the sampling result of the provided {@link #rootSampler}, unless the
25+
* sampling result contains the sampling decision {@link SamplingDecision#DROP}, in which case, a
26+
* new sampling result will be returned that is functionally equivalent to the original, except that
27+
* it contains the sampling decision {@link SamplingDecision#RECORD_ONLY}. This ensures that all
28+
* spans are recorded, with no change to sampling.
29+
*
30+
* <p>An intended use case of this sampler is to provide a means of sending all spans to a processor
31+
* without having an impact on the sampling rate. This may be desirable if a user wishes to count or
32+
* otherwise measure all spans produced in a service, without incurring the cost of 100% sampling.
33+
*
34+
* <p>This class is internal and experimental. Its APIs are unstable and can change at any time. Its
35+
* APIs (or a version of them) may be promoted to the public stable API in the future, but no
36+
* guarantees are made.
37+
*/
38+
@Immutable
39+
public final class AlwaysRecordSampler implements Sampler {
40+
41+
private final Sampler rootSampler;
42+
43+
public static AlwaysRecordSampler create(Sampler rootSampler) {
44+
return new AlwaysRecordSampler(rootSampler);
45+
}
46+
47+
private AlwaysRecordSampler(Sampler rootSampler) {
48+
this.rootSampler = rootSampler;
49+
}
50+
51+
@Override
52+
public SamplingResult shouldSample(
53+
Context parentContext,
54+
String traceId,
55+
String name,
56+
SpanKind spanKind,
57+
Attributes attributes,
58+
List<LinkData> parentLinks) {
59+
SamplingResult result =
60+
rootSampler.shouldSample(parentContext, traceId, name, spanKind, attributes, parentLinks);
61+
if (result.getDecision() == SamplingDecision.DROP) {
62+
result = wrapResultWithRecordOnlyResult(result);
63+
}
64+
65+
return result;
66+
}
67+
68+
@Override
69+
public String getDescription() {
70+
return "AlwaysRecordSampler{" + rootSampler.getDescription() + "}";
71+
}
72+
73+
private static SamplingResult wrapResultWithRecordOnlyResult(SamplingResult result) {
74+
return new SamplingResult() {
75+
@Override
76+
public SamplingDecision getDecision() {
77+
return SamplingDecision.RECORD_ONLY;
78+
}
79+
80+
@Override
81+
public Attributes getAttributes() {
82+
return result.getAttributes();
83+
}
84+
85+
@Override
86+
public TraceState getUpdatedTraceState(TraceState parentTraceState) {
87+
return result.getUpdatedTraceState(parentTraceState);
88+
}
89+
};
90+
}
91+
}

‎sdk/trace/src/test/java/io/opentelemetry/sdk/trace/internal/AlwaysRecordSamplerTest.java‎

Lines changed: 112 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,112 @@
1+
/*
2+
* Copyright The OpenTelemetry Authors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
// Includes work from:
7+
// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
8+
// SPDX-License-Identifier: Apache-2.0
9+
10+
package io.opentelemetry.sdk.trace.internal;
11+
12+
import static org.assertj.core.api.Assertions.assertThat;
13+
import static org.mockito.ArgumentMatchers.any;
14+
import static org.mockito.ArgumentMatchers.anyString;
15+
import static org.mockito.Mockito.mock;
16+
import static org.mockito.Mockito.when;
17+
18+
import io.opentelemetry.api.common.AttributeKey;
19+
import io.opentelemetry.api.common.Attributes;
20+
import io.opentelemetry.api.trace.SpanKind;
21+
import io.opentelemetry.api.trace.TraceId;
22+
import io.opentelemetry.api.trace.TraceState;
23+
import io.opentelemetry.context.Context;
24+
import io.opentelemetry.sdk.trace.samplers.Sampler;
25+
import io.opentelemetry.sdk.trace.samplers.SamplingDecision;
26+
import io.opentelemetry.sdk.trace.samplers.SamplingResult;
27+
import java.util.Collections;
28+
import org.junit.jupiter.api.BeforeEach;
29+
import org.junit.jupiter.api.Test;
30+
31+
/** Unit tests for {@link AlwaysRecordSampler}. */
32+
class AlwaysRecordSamplerTest {
33+
34+
// Mocks
35+
private Sampler mockSampler;
36+
37+
private AlwaysRecordSampler sampler;
38+
39+
@BeforeEach
40+
void setUpSamplers() {
41+
mockSampler = mock(Sampler.class);
42+
sampler = AlwaysRecordSampler.create(mockSampler);
43+
}
44+
45+
@Test
46+
void testGetDescription() {
47+
when(mockSampler.getDescription()).thenReturn("mockDescription");
48+
assertThat(sampler.getDescription()).isEqualTo("AlwaysRecordSampler{mockDescription}");
49+
}
50+
51+
@Test
52+
void testRecordAndSampleSamplingDecision() {
53+
validateShouldSample(SamplingDecision.RECORD_AND_SAMPLE, SamplingDecision.RECORD_AND_SAMPLE);
54+
}
55+
56+
@Test
57+
void testRecordOnlySamplingDecision() {
58+
validateShouldSample(SamplingDecision.RECORD_ONLY, SamplingDecision.RECORD_ONLY);
59+
}
60+
61+
@Test
62+
void testDropSamplingDecision() {
63+
validateShouldSample(SamplingDecision.DROP, SamplingDecision.RECORD_ONLY);
64+
}
65+
66+
private void validateShouldSample(
67+
SamplingDecision rootDecision, SamplingDecision expectedDecision) {
68+
SamplingResult rootResult = buildRootSamplingResult(rootDecision);
69+
when(mockSampler.shouldSample(any(), anyString(), anyString(), any(), any(), any()))
70+
.thenReturn(rootResult);
71+
SamplingResult actualResult =
72+
sampler.shouldSample(
73+
Context.current(),
74+
TraceId.fromLongs(1, 2),
75+
"name",
76+
SpanKind.CLIENT,
77+
Attributes.empty(),
78+
Collections.emptyList());
79+
80+
if (rootDecision.equals(expectedDecision)) {
81+
assertThat(actualResult).isEqualTo(rootResult);
82+
assertThat(actualResult.getDecision()).isEqualTo(rootDecision);
83+
} else {
84+
assertThat(actualResult).isNotEqualTo(rootResult);
85+
assertThat(actualResult.getDecision()).isEqualTo(expectedDecision);
86+
}
87+
88+
assertThat(actualResult.getAttributes()).isEqualTo(rootResult.getAttributes());
89+
TraceState traceState = TraceState.builder().build();
90+
assertThat(actualResult.getUpdatedTraceState(traceState))
91+
.isEqualTo(rootResult.getUpdatedTraceState(traceState));
92+
}
93+
94+
private static SamplingResult buildRootSamplingResult(SamplingDecision samplingDecision) {
95+
return new SamplingResult() {
96+
@Override
97+
public SamplingDecision getDecision() {
98+
return samplingDecision;
99+
}
100+
101+
@Override
102+
public Attributes getAttributes() {
103+
return Attributes.of(AttributeKey.stringKey("key"), samplingDecision.name());
104+
}
105+
106+
@Override
107+
public TraceState getUpdatedTraceState(TraceState parentTraceState) {
108+
return TraceState.builder().put("key", samplingDecision.name()).build();
109+
}
110+
};
111+
}
112+
}

0 commit comments

Comments
 (0)