Skip to content

Commit 2f6490b

Browse files
srprashsrprash
authored
Adding AwsUnsampledOnlySpanProcessor to export batches of unsampled spans (#948)
1 parent ffc6afe commit 2f6490b

File tree

4 files changed

+297
-0
lines changed

4 files changed

+297
-0
lines changed

‎awsagentprovider/src/main/java/software/amazon/opentelemetry/javaagent/providers/AwsAttributeKeys.java

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -70,6 +70,9 @@ private AwsAttributeKeys() {}
7070
static final AttributeKey<String> AWS_LAMBDA_RESOURCE_ID =
7171
AttributeKey.stringKey("aws.lambda.resource_mapping.id");
7272

73+
static final AttributeKey<Boolean> AWS_TRACE_FLAG_SAMPLED =
74+
AttributeKey.booleanKey("aws.trace.flag.sampled");
75+
7376
// use the same AWS Resource attribute name defined by OTel java auto-instr for aws_sdk_v_1_1
7477
// TODO: all AWS specific attributes should be defined in semconv package and reused cross all
7578
// otel packages. Related sim -

‎awsagentprovider/src/main/java/software/amazon/opentelemetry/javaagent/providers/AwsUnsampledOnlySpanProcessor.java

Lines changed: 85 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,85 @@
1+
/*
2+
* Copyright Amazon.com, Inc. or its affiliates.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License").
5+
* You may not use this file except in compliance with the License.
6+
* A copy of the License is located at
7+
*
8+
* http://aws.amazon.com/apache2.0
9+
*
10+
* or in the "license" file accompanying this file. This file is distributed
11+
* on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
12+
* express or implied. See the License for the specific language governing
13+
* permissions and limitations under the License.
14+
*/
15+
16+
package software.amazon.opentelemetry.javaagent.providers;
17+
18+
import io.opentelemetry.context.Context;
19+
import io.opentelemetry.sdk.common.CompletableResultCode;
20+
import io.opentelemetry.sdk.trace.ReadWriteSpan;
21+
import io.opentelemetry.sdk.trace.ReadableSpan;
22+
import io.opentelemetry.sdk.trace.SpanProcessor;
23+
24+
/**
25+
* {@link SpanProcessor} that only exports unsampled spans in a batch via a delegated @{link
26+
* BatchSpanProcessor}. The processor also adds an attribute to each processed span to indicate that
27+
* it was sampled or not.
28+
*/
29+
final class AwsUnsampledOnlySpanProcessor implements SpanProcessor {
30+
31+
private final SpanProcessor delegate;
32+
33+
AwsUnsampledOnlySpanProcessor(SpanProcessor delegate) {
34+
this.delegate = delegate;
35+
}
36+
37+
public static AwsUnsampledOnlySpanProcessorBuilder builder() {
38+
return new AwsUnsampledOnlySpanProcessorBuilder();
39+
}
40+
41+
@Override
42+
public void onStart(Context parentContext, ReadWriteSpan span) {
43+
if (!span.getSpanContext().isSampled()) {
44+
span.setAttribute(AwsAttributeKeys.AWS_TRACE_FLAG_SAMPLED, false);
45+
}
46+
delegate.onStart(parentContext, span);
47+
}
48+
49+
@Override
50+
public void onEnd(ReadableSpan span) {
51+
if (!span.getSpanContext().isSampled()) {
52+
delegate.onEnd(span);
53+
}
54+
}
55+
56+
@Override
57+
public boolean isStartRequired() {
58+
return true;
59+
}
60+
61+
@Override
62+
public boolean isEndRequired() {
63+
return true;
64+
}
65+
66+
@Override
67+
public CompletableResultCode shutdown() {
68+
return delegate.shutdown();
69+
}
70+
71+
@Override
72+
public CompletableResultCode forceFlush() {
73+
return delegate.forceFlush();
74+
}
75+
76+
@Override
77+
public void close() {
78+
delegate.close();
79+
}
80+
81+
// Visible for testing
82+
SpanProcessor getDelegate() {
83+
return delegate;
84+
}
85+
}

‎awsagentprovider/src/main/java/software/amazon/opentelemetry/javaagent/providers/AwsUnsampledOnlySpanProcessorBuilder.java

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,46 @@
1+
/*
2+
* Copyright Amazon.com, Inc. or its affiliates.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License").
5+
* You may not use this file except in compliance with the License.
6+
* A copy of the License is located at
7+
*
8+
* http://aws.amazon.com/apache2.0
9+
*
10+
* or in the "license" file accompanying this file. This file is distributed
11+
* on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
12+
* express or implied. See the License for the specific language governing
13+
* permissions and limitations under the License.
14+
*/
15+
16+
package software.amazon.opentelemetry.javaagent.providers;
17+
18+
import static java.util.Objects.requireNonNull;
19+
20+
import io.opentelemetry.sdk.trace.export.BatchSpanProcessor;
21+
import io.opentelemetry.sdk.trace.export.SpanExporter;
22+
23+
final class AwsUnsampledOnlySpanProcessorBuilder {
24+
25+
// Default exporter is OtlpUdpSpanExporter with unsampled payload prefix
26+
private SpanExporter exporter =
27+
new OtlpUdpSpanExporterBuilder()
28+
.setPayloadSampleDecision(TracePayloadSampleDecision.UNSAMPLED)
29+
.build();
30+
31+
public AwsUnsampledOnlySpanProcessorBuilder setSpanExporter(SpanExporter exporter) {
32+
requireNonNull(exporter, "exporter cannot be null");
33+
this.exporter = exporter;
34+
return this;
35+
}
36+
37+
public AwsUnsampledOnlySpanProcessor build() {
38+
BatchSpanProcessor bsp =
39+
BatchSpanProcessor.builder(exporter).setExportUnsampledSpans(true).build();
40+
return new AwsUnsampledOnlySpanProcessor(bsp);
41+
}
42+
43+
SpanExporter getSpanExporter() {
44+
return exporter;
45+
}
46+
}

‎awsagentprovider/src/test/java/software/amazon/opentelemetry/javaagent/providers/AwsUnsampledOnlySpanProcessorTest.java

Lines changed: 163 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,163 @@
1+
/*
2+
* Copyright Amazon.com, Inc. or its affiliates.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License").
5+
* You may not use this file except in compliance with the License.
6+
* A copy of the License is located at
7+
*
8+
* http://aws.amazon.com/apache2.0
9+
*
10+
* or in the "license" file accompanying this file. This file is distributed
11+
* on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
12+
* express or implied. See the License for the specific language governing
13+
* permissions and limitations under the License.
14+
*/
15+
16+
package software.amazon.opentelemetry.javaagent.providers;
17+
18+
import static org.assertj.core.api.Assertions.assertThat;
19+
import static org.mockito.Mockito.*;
20+
21+
import io.opentelemetry.api.trace.SpanContext;
22+
import io.opentelemetry.context.Context;
23+
import io.opentelemetry.sdk.common.CompletableResultCode;
24+
import io.opentelemetry.sdk.testing.exporter.InMemorySpanExporter;
25+
import io.opentelemetry.sdk.trace.ReadWriteSpan;
26+
import io.opentelemetry.sdk.trace.ReadableSpan;
27+
import io.opentelemetry.sdk.trace.SpanProcessor;
28+
import io.opentelemetry.sdk.trace.export.BatchSpanProcessor;
29+
import io.opentelemetry.sdk.trace.export.SpanExporter;
30+
import java.util.ArrayList;
31+
import java.util.Collection;
32+
import org.junit.jupiter.api.Test;
33+
34+
public class AwsUnsampledOnlySpanProcessorTest {
35+
36+
@Test
37+
public void testIsStartRequired() {
38+
SpanProcessor processor = AwsUnsampledOnlySpanProcessor.builder().build();
39+
assertThat(processor.isStartRequired()).isTrue();
40+
}
41+
42+
@Test
43+
public void testIsEndRequired() {
44+
SpanProcessor processor = AwsUnsampledOnlySpanProcessor.builder().build();
45+
assertThat(processor.isEndRequired()).isTrue();
46+
}
47+
48+
@Test
49+
public void testDefaultSpanProcessor() {
50+
AwsUnsampledOnlySpanProcessorBuilder builder = AwsUnsampledOnlySpanProcessor.builder();
51+
AwsUnsampledOnlySpanProcessor unsampledSP = builder.build();
52+
53+
assertThat(builder.getSpanExporter()).isInstanceOf(OtlpUdpSpanExporter.class);
54+
SpanProcessor delegate = unsampledSP.getDelegate();
55+
assertThat(delegate).isInstanceOf(BatchSpanProcessor.class);
56+
BatchSpanProcessor delegateBsp = (BatchSpanProcessor) delegate;
57+
String delegateBspString = delegateBsp.toString();
58+
assertThat(delegateBspString)
59+
.contains(
60+
"spanExporter=software.amazon.opentelemetry.javaagent.providers.OtlpUdpSpanExporter");
61+
assertThat(delegateBspString).contains("exportUnsampledSpans=true");
62+
}
63+
64+
@Test
65+
public void testSpanProcessorWithExporter() {
66+
AwsUnsampledOnlySpanProcessorBuilder builder =
67+
AwsUnsampledOnlySpanProcessor.builder().setSpanExporter(InMemorySpanExporter.create());
68+
AwsUnsampledOnlySpanProcessor unsampledSP = builder.build();
69+
70+
assertThat(builder.getSpanExporter()).isInstanceOf(InMemorySpanExporter.class);
71+
SpanProcessor delegate = unsampledSP.getDelegate();
72+
assertThat(delegate).isInstanceOf(BatchSpanProcessor.class);
73+
BatchSpanProcessor delegateBsp = (BatchSpanProcessor) delegate;
74+
String delegateBspString = delegateBsp.toString();
75+
assertThat(delegateBspString)
76+
.contains("spanExporter=io.opentelemetry.sdk.testing.exporter.InMemorySpanExporter");
77+
assertThat(delegateBspString).contains("exportUnsampledSpans=true");
78+
}
79+
80+
@Test
81+
public void testStartAddsAttributeToSampledSpan() {
82+
SpanContext mockSpanContext = mock(SpanContext.class);
83+
when(mockSpanContext.isSampled()).thenReturn(true);
84+
Context parentContextMock = mock(Context.class);
85+
ReadWriteSpan spanMock = mock(ReadWriteSpan.class);
86+
when(spanMock.getSpanContext()).thenReturn(mockSpanContext);
87+
88+
AwsUnsampledOnlySpanProcessor processor = AwsUnsampledOnlySpanProcessor.builder().build();
89+
processor.onStart(parentContextMock, spanMock);
90+
91+
// verify setAttribute was never called
92+
verify(spanMock, never()).setAttribute(any(), anyBoolean());
93+
}
94+
95+
@Test
96+
public void testStartAddsAttributeToUnsampledSpan() {
97+
SpanContext mockSpanContext = mock(SpanContext.class);
98+
when(mockSpanContext.isSampled()).thenReturn(false);
99+
Context parentContextMock = mock(Context.class);
100+
ReadWriteSpan spanMock = mock(ReadWriteSpan.class);
101+
when(spanMock.getSpanContext()).thenReturn(mockSpanContext);
102+
103+
AwsUnsampledOnlySpanProcessor processor = AwsUnsampledOnlySpanProcessor.builder().build();
104+
processor.onStart(parentContextMock, spanMock);
105+
106+
// verify setAttribute was called with the correct arguments
107+
verify(spanMock, times(1)).setAttribute(AwsAttributeKeys.AWS_TRACE_FLAG_SAMPLED, false);
108+
}
109+
110+
@Test
111+
public void testExportsOnlyUnsampledSpans() {
112+
SpanExporter mockExporter = mock(SpanExporter.class);
113+
when(mockExporter.export(anyCollection())).thenReturn(CompletableResultCode.ofSuccess());
114+
115+
TestDelegateProcessor delegate = new TestDelegateProcessor();
116+
AwsUnsampledOnlySpanProcessor processor = new AwsUnsampledOnlySpanProcessor(delegate);
117+
118+
// unsampled span
119+
SpanContext mockSpanContextUnsampled = mock(SpanContext.class);
120+
when(mockSpanContextUnsampled.isSampled()).thenReturn(false);
121+
ReadableSpan mockSpanUnsampled = mock(ReadableSpan.class);
122+
when(mockSpanUnsampled.getSpanContext()).thenReturn(mockSpanContextUnsampled);
123+
124+
// sampled span
125+
SpanContext mockSpanContextSampled = mock(SpanContext.class);
126+
when(mockSpanContextSampled.isSampled()).thenReturn(true);
127+
ReadableSpan mockSpanSampled = mock(ReadableSpan.class);
128+
when(mockSpanSampled.getSpanContext()).thenReturn(mockSpanContextSampled);
129+
130+
processor.onEnd(mockSpanSampled);
131+
processor.onEnd(mockSpanUnsampled);
132+
133+
// validate that only the unsampled span was delegated
134+
assertThat(delegate.getEndedSpans()).containsExactly(mockSpanUnsampled);
135+
}
136+
137+
private static class TestDelegateProcessor implements SpanProcessor {
138+
// keep a queue of Readable spans added when onEnd is called
139+
Collection<ReadableSpan> endedSpans = new ArrayList<>();
140+
141+
@Override
142+
public void onStart(Context parentContext, ReadWriteSpan span) {}
143+
144+
@Override
145+
public boolean isStartRequired() {
146+
return false;
147+
}
148+
149+
@Override
150+
public void onEnd(ReadableSpan span) {
151+
endedSpans.add(span);
152+
}
153+
154+
@Override
155+
public boolean isEndRequired() {
156+
return false;
157+
}
158+
159+
public Collection<ReadableSpan> getEndedSpans() {
160+
return endedSpans;
161+
}
162+
}
163+
}

0 commit comments

Comments
 (0)