Skip to content

Commit c0754c2

Browse files
anuraagaanuraaga
authored
Add an AWS composite propagator that can be used in AWS-service compatible projects. (#50)
1 parent 488ca26 commit c0754c2

File tree

5 files changed

+308
-0
lines changed

5 files changed

+308
-0
lines changed

‎awspropagator/build.gradle.kts‎

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,28 @@
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+
plugins {
17+
`java-library`
18+
}
19+
20+
base {
21+
archivesBaseName = "aws-opentelemetry-propagator"
22+
}
23+
24+
dependencies {
25+
api("io.opentelemetry:opentelemetry-context")
26+
implementation("io.opentelemetry:opentelemetry-extension-aws")
27+
implementation("io.opentelemetry:opentelemetry-extension-trace-propagators")
28+
}

‎awspropagator/src/main/java/software/amazon/opentelemetry/awspropagator/AwsCompositePropagator.java‎

Lines changed: 124 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,124 @@
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.awspropagator;
17+
18+
import io.opentelemetry.api.baggage.Baggage;
19+
import io.opentelemetry.api.baggage.propagation.W3CBaggagePropagator;
20+
import io.opentelemetry.api.trace.Span;
21+
import io.opentelemetry.api.trace.propagation.W3CTraceContextPropagator;
22+
import io.opentelemetry.context.Context;
23+
import io.opentelemetry.context.ContextKey;
24+
import io.opentelemetry.context.propagation.TextMapGetter;
25+
import io.opentelemetry.context.propagation.TextMapPropagator;
26+
import io.opentelemetry.context.propagation.TextMapSetter;
27+
import io.opentelemetry.extension.aws.AwsXrayPropagator;
28+
import io.opentelemetry.extension.trace.propagation.B3Propagator;
29+
import java.util.Arrays;
30+
import java.util.Collection;
31+
import java.util.List;
32+
import java.util.stream.Collectors;
33+
34+
/**
35+
* A {@link TextMapPropagator} for use by AWS services or compatible ones such as localstack. As AWS
36+
* is generic cloud infrastructure, rather than enforcing a specific injection format such as
37+
* W3CTraceContext, it is intended to inject the same format as was extracted to allow compatibility
38+
* with client applications that may not understand newer formats like W3C.
39+
*
40+
* <p>To allow rolling this propagator out to independent services, it also supports a mode that
41+
* always injects in AWS format. This is intended for a transition period where most services will
42+
* only understand that format.
43+
*/
44+
public final class AwsCompositePropagator implements TextMapPropagator {
45+
46+
// The currently supported propagators. This may grow or be made configurable in the future.
47+
private static final List<TextMapPropagator> PROPAGATORS =
48+
Arrays.asList(
49+
W3CBaggagePropagator.getInstance(),
50+
AwsXrayPropagator.getInstance(),
51+
W3CTraceContextPropagator.getInstance(),
52+
B3Propagator.injectingMultiHeaders());
53+
54+
/**
55+
* Returns a {@link AwsCompositePropagator} which injects in the same format it extracted a trace
56+
* context with. For example, if an incoming request contained a trace context in B3 format, an
57+
* outgoing request will have B3 format injected.
58+
*/
59+
public static AwsCompositePropagator injectingExtractedFormat() {
60+
return new AwsCompositePropagator(true);
61+
}
62+
63+
/**
64+
* Returns a {@link AwsCompositePropagator} which injects in the AWS format (X-Amzn-Trace-Id). For
65+
* example, if an incoming request contained a trace context in B3 format, an outgoing request
66+
* will have AWS format injected.
67+
*/
68+
public static AwsCompositePropagator injectingAwsFormat() {
69+
return new AwsCompositePropagator(false);
70+
}
71+
72+
private static final ContextKey<TextMapPropagator> EXTRACTED_PROPAGATOR =
73+
ContextKey.named("extracted-propagator");
74+
75+
private final List<String> fields;
76+
private final boolean injectExtractedFormat;
77+
78+
private AwsCompositePropagator(boolean injectExtractedFormat) {
79+
this.injectExtractedFormat = injectExtractedFormat;
80+
81+
fields = PROPAGATORS.stream().flatMap(p -> p.fields().stream()).collect(Collectors.toList());
82+
}
83+
84+
@Override
85+
public Collection<String> fields() {
86+
return fields;
87+
}
88+
89+
@Override
90+
public <C> void inject(Context context, C carrier, TextMapSetter<C> setter) {
91+
if (injectExtractedFormat) {
92+
TextMapPropagator extractedPropagator = context.get(EXTRACTED_PROPAGATOR);
93+
if (extractedPropagator != null) {
94+
extractedPropagator.inject(context, carrier, setter);
95+
Baggage baggage = Baggage.fromContextOrNull(context);
96+
if (baggage != null && extractedPropagator != AwsXrayPropagator.getInstance()) {
97+
// We extracted a span from a format not supporting baggage within the trace context
98+
// itself, for example b3. if we have baggage we just propagate using w3c
99+
// baggage.
100+
W3CBaggagePropagator.getInstance().inject(context, carrier, setter);
101+
}
102+
return;
103+
}
104+
}
105+
// Unless injecting in the same format as extracted, always inject X-Amzn-Trace-Id, the only
106+
// format recognized by all AWS services as of now.
107+
AwsXrayPropagator.getInstance().inject(context, carrier, setter);
108+
}
109+
110+
@Override
111+
public <C> Context extract(Context context, C carrier, TextMapGetter<C> getter) {
112+
for (TextMapPropagator propagator : PROPAGATORS) {
113+
context = propagator.extract(context, carrier, getter);
114+
if (Span.fromContextOrNull(context) != null) {
115+
if (injectExtractedFormat) {
116+
context = context.with(EXTRACTED_PROPAGATOR, propagator);
117+
}
118+
return context;
119+
}
120+
}
121+
122+
return context;
123+
}
124+
}

‎awspropagator/src/test/java/software/amazon/opentelemetry/awspropagator/AwsCompositePropagatorTest.java‎

Lines changed: 154 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,154 @@
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.awspropagator;
17+
18+
import static org.assertj.core.api.Assertions.assertThat;
19+
20+
import io.opentelemetry.api.baggage.Baggage;
21+
import io.opentelemetry.api.baggage.propagation.W3CBaggagePropagator;
22+
import io.opentelemetry.api.trace.Span;
23+
import io.opentelemetry.api.trace.SpanContext;
24+
import io.opentelemetry.api.trace.TraceFlags;
25+
import io.opentelemetry.api.trace.TraceState;
26+
import io.opentelemetry.api.trace.propagation.W3CTraceContextPropagator;
27+
import io.opentelemetry.context.Context;
28+
import io.opentelemetry.context.propagation.TextMapGetter;
29+
import io.opentelemetry.context.propagation.TextMapPropagator;
30+
import io.opentelemetry.extension.aws.AwsXrayPropagator;
31+
import io.opentelemetry.extension.trace.propagation.B3Propagator;
32+
import java.util.HashMap;
33+
import java.util.Map;
34+
import java.util.stream.Stream;
35+
import org.junit.jupiter.api.Test;
36+
import org.junit.jupiter.api.extension.ExtensionContext;
37+
import org.junit.jupiter.params.ParameterizedTest;
38+
import org.junit.jupiter.params.provider.Arguments;
39+
import org.junit.jupiter.params.provider.ArgumentsProvider;
40+
import org.junit.jupiter.params.provider.ArgumentsSource;
41+
42+
class AwsCompositePropagatorTest {
43+
44+
private static final Span SPAN1 =
45+
Span.wrap(
46+
SpanContext.createFromRemoteParent(
47+
"ff000000000000000000000000000041",
48+
"ff00000000000041",
49+
TraceFlags.getDefault(),
50+
TraceState.getDefault()));
51+
private static final Span SPAN2 =
52+
Span.wrap(
53+
SpanContext.createFromRemoteParent(
54+
"ff000000000000000000000000000042",
55+
"ff00000000000042",
56+
TraceFlags.getDefault(),
57+
TraceState.getDefault()));
58+
59+
private static final TextMapGetter<Map<String, String>> MAP_GETTER =
60+
new TextMapGetter<>() {
61+
@Override
62+
public Iterable<String> keys(Map<String, String> carrier) {
63+
return carrier.keySet();
64+
}
65+
66+
@Override
67+
public String get(Map<String, String> carrier, String key) {
68+
return carrier.get(key);
69+
}
70+
};
71+
72+
@ParameterizedTest
73+
@ArgumentsSource(SupportedPropagators.class)
74+
void extractsAndInjectAws(TextMapPropagator propagator) {
75+
Map<String, String> map = new HashMap<>();
76+
propagator.inject(Context.root().with(SPAN1), map, Map::put);
77+
TextMapPropagator awsPropagator = AwsCompositePropagator.injectingAwsFormat();
78+
Context context = awsPropagator.extract(Context.root(), map, MAP_GETTER);
79+
assertThat(Span.fromContext(context).getSpanContext()).isEqualTo(SPAN1.getSpanContext());
80+
81+
Map<String, String> map2 = new HashMap<>();
82+
awsPropagator.inject(context, map2, Map::put);
83+
84+
Map<String, String> map3 = new HashMap<>();
85+
AwsXrayPropagator.getInstance().inject(context, map3, Map::put);
86+
assertThat(map2).containsExactlyInAnyOrderEntriesOf(map3);
87+
}
88+
89+
@ParameterizedTest
90+
@ArgumentsSource(SupportedPropagators.class)
91+
void extractsAndInjectExtracted(TextMapPropagator propagator) {
92+
Map<String, String> map = new HashMap<>();
93+
propagator.inject(Context.root().with(SPAN1), map, Map::put);
94+
TextMapPropagator awsPropagator = AwsCompositePropagator.injectingExtractedFormat();
95+
Context context = awsPropagator.extract(Context.root(), map, MAP_GETTER);
96+
assertThat(Span.fromContext(context).getSpanContext()).isEqualTo(SPAN1.getSpanContext());
97+
98+
Map<String, String> map2 = new HashMap<>();
99+
awsPropagator.inject(context, map2, Map::put);
100+
101+
Map<String, String> map3 = new HashMap<>();
102+
propagator.inject(context, map3, Map::put);
103+
assertThat(map2).containsExactlyInAnyOrderEntriesOf(map3);
104+
}
105+
106+
@ParameterizedTest
107+
@ArgumentsSource(SupportedPropagators.class)
108+
void baggageExtractedFormat(TextMapPropagator propagator) {
109+
Map<String, String> map = new HashMap<>();
110+
propagator.inject(Context.root().with(SPAN1), map, Map::put);
111+
Baggage baggage = Baggage.builder().put("cat", "meow").build();
112+
W3CBaggagePropagator.getInstance().inject(Context.root().with(baggage), map, Map::put);
113+
TextMapPropagator awsPropagator = AwsCompositePropagator.injectingExtractedFormat();
114+
Context context = awsPropagator.extract(Context.root(), map, MAP_GETTER);
115+
assertThat(Span.fromContext(context).getSpanContext()).isEqualTo(SPAN1.getSpanContext());
116+
assertThat(Baggage.fromContext(context)).isEqualTo(baggage);
117+
118+
Map<String, String> map2 = new HashMap<>();
119+
awsPropagator.inject(context, map2, Map::put);
120+
121+
Map<String, String> map3 = new HashMap<>();
122+
propagator.inject(context, map3, Map::put);
123+
124+
if (propagator != AwsXrayPropagator.getInstance()) {
125+
W3CBaggagePropagator.getInstance().inject(context, map3, Map::put);
126+
}
127+
128+
assertThat(map2).containsExactlyInAnyOrderEntriesOf(map3);
129+
}
130+
131+
@Test
132+
void awsPrioritized() {
133+
Map<String, String> map = new HashMap<>();
134+
AwsXrayPropagator.getInstance().inject(Context.root().with(SPAN1), map, Map::put);
135+
W3CTraceContextPropagator.getInstance().inject(Context.root().with(SPAN2), map, Map::put);
136+
assertThat(
137+
Span.fromContext(
138+
AwsCompositePropagator.injectingExtractedFormat()
139+
.extract(Context.root(), map, MAP_GETTER))
140+
.getSpanContext())
141+
.isEqualTo(SPAN1.getSpanContext());
142+
}
143+
144+
private static class SupportedPropagators implements ArgumentsProvider {
145+
@Override
146+
public Stream<? extends Arguments> provideArguments(ExtensionContext context) {
147+
return Stream.of(
148+
AwsXrayPropagator.getInstance(),
149+
W3CTraceContextPropagator.getInstance(),
150+
B3Propagator.injectingMultiHeaders())
151+
.map(Arguments::of);
152+
}
153+
}
154+
}

‎build.gradle.kts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -90,6 +90,7 @@ allprojects {
9090

9191
testImplementation("org.assertj:assertj-core")
9292
testImplementation("org.junit.jupiter:junit-jupiter-api")
93+
testImplementation("org.junit.jupiter:junit-jupiter-params")
9394
testRuntimeOnly("org.junit.jupiter:junit-jupiter-engine")
9495
}
9596

‎settings.gradle.kts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,7 @@ dependencyResolutionManagement {
3838
}
3939

4040
include(":awsagentprovider")
41+
include(":awspropagator")
4142
include(":dependencyManagement")
4243
include(":instrumentation:logback-1.0")
4344
include(":instrumentation:log4j-2.13.2")

0 commit comments

Comments
 (0)