Skip to content

Commit 67efe98

Browse files
majanjua-amznmajanjua-amzn
authored
Add AWS X-Ray Adaptive Sampling Support (#2147)
1 parent 828f00c commit 67efe98

File tree

17 files changed

+2452
-117
lines changed

17 files changed

+2452
-117
lines changed

‎aws-xray-propagator/src/main/java/io/opentelemetry/contrib/awsxray/propagator/AwsXrayPropagator.java‎

Lines changed: 21 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010

1111
import io.opentelemetry.api.baggage.Baggage;
1212
import io.opentelemetry.api.baggage.BaggageBuilder;
13+
import io.opentelemetry.api.baggage.propagation.W3CBaggagePropagator;
1314
import io.opentelemetry.api.internal.StringUtils;
1415
import io.opentelemetry.api.trace.Span;
1516
import io.opentelemetry.api.trace.SpanContext;
@@ -80,6 +81,9 @@ public final class AwsXrayPropagator implements TextMapPropagator {
8081
private static final String INVALID_LINEAGE = "-1:11111111:0";
8182
private static final int NUM_OF_LINEAGE_DELIMITERS = 2;
8283

84+
// Copied from AwsSamplingResult in aws-xray extension
85+
private static final String AWS_XRAY_SAMPLING_RULE_TRACE_STATE_KEY = "xrsr";
86+
8387
private static final List<String> FIELDS = singletonList(TRACE_HEADER_KEY);
8488

8589
private static final AwsXrayPropagator INSTANCE = new AwsXrayPropagator();
@@ -140,6 +144,16 @@ public <C> void inject(Context context, @Nullable C carrier, TextMapSetter<C> se
140144

141145
Baggage baggage = Baggage.fromContext(context);
142146
String lineageHeader = baggage.getEntryValue(LINEAGE_KEY);
147+
// Get sampling rule from trace state and inject into baggage
148+
// This is a back up in case the next service does not have trace state propagation
149+
String ruleFromTraceState =
150+
spanContext.getTraceState().get(AWS_XRAY_SAMPLING_RULE_TRACE_STATE_KEY);
151+
if (ruleFromTraceState != null) {
152+
baggage =
153+
baggage.toBuilder()
154+
.put(AWS_XRAY_SAMPLING_RULE_TRACE_STATE_KEY, ruleFromTraceState)
155+
.build();
156+
}
143157

144158
if (lineageHeader != null) {
145159
traceHeader
@@ -152,6 +166,9 @@ public <C> void inject(Context context, @Nullable C carrier, TextMapSetter<C> se
152166
// add 256 character truncation
153167
String truncatedTraceHeader = traceHeader.substring(0, Math.min(traceHeader.length(), 256));
154168
setter.set(carrier, TRACE_HEADER_KEY, truncatedTraceHeader);
169+
170+
// Ensure baggage is propagated with any modifications
171+
W3CBaggagePropagator.getInstance().inject(context.with(baggage), carrier, setter);
155172
}
156173

157174
@Override
@@ -245,12 +262,15 @@ private static <C> Context getContextFromHeader(
245262
logger.finest("Both traceId and spanId are required to extract a valid span context. ");
246263
}
247264

265+
SpanContext upstreamSpanContext = Span.fromContext(context).getSpanContext();
248266
SpanContext spanContext =
249267
SpanContext.createFromRemoteParent(
250268
StringUtils.padLeft(traceId, TraceId.getLength()),
251269
spanId,
252270
isSampled ? TraceFlags.getSampled() : TraceFlags.getDefault(),
253-
TraceState.getDefault());
271+
upstreamSpanContext.isValid()
272+
? upstreamSpanContext.getTraceState()
273+
: TraceState.getDefault());
254274

255275
if (spanContext.isValid()) {
256276
context = context.with(Span.wrap(spanContext));

‎aws-xray/build.gradle.kts‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ dependencies {
1111
api("io.opentelemetry:opentelemetry-sdk-trace")
1212

1313
compileOnly("io.opentelemetry:opentelemetry-sdk-extension-autoconfigure")
14+
implementation("io.opentelemetry.semconv:opentelemetry-semconv:1.32.0-alpha")
1415

1516
implementation("com.squareup.okhttp3:okhttp")
1617
implementation("io.opentelemetry.semconv:opentelemetry-semconv")
@@ -25,6 +26,7 @@ dependencies {
2526

2627
implementation("com.fasterxml.jackson.core:jackson-core")
2728
implementation("com.fasterxml.jackson.core:jackson-databind")
29+
implementation("com.github.ben-manes.caffeine:caffeine:2.9.3")
2830

2931
testImplementation("com.linecorp.armeria:armeria-junit5")
3032
testImplementation("io.opentelemetry:opentelemetry-sdk-extension-autoconfigure")

‎aws-xray/src/main/java/io/opentelemetry/contrib/awsxray/AwsSamplingResult.java‎

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,56 @@
1+
/*
2+
* Copyright The OpenTelemetry Authors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package io.opentelemetry.contrib.awsxray;
7+
8+
import io.opentelemetry.api.common.Attributes;
9+
import io.opentelemetry.api.trace.TraceState;
10+
import io.opentelemetry.sdk.trace.samplers.SamplingDecision;
11+
import io.opentelemetry.sdk.trace.samplers.SamplingResult;
12+
import javax.annotation.Nullable;
13+
14+
final class AwsSamplingResult implements SamplingResult {
15+
16+
// OTel trace state is a space shared with other vendors with a 256 character limit
17+
// We keep the key and values as short as possible while still identifiable
18+
public static final String AWS_XRAY_SAMPLING_RULE_TRACE_STATE_KEY = "xrsr";
19+
20+
private final SamplingDecision decision;
21+
private final Attributes attributes;
22+
@Nullable private final String samplingRuleName;
23+
24+
private AwsSamplingResult(
25+
SamplingDecision decision, Attributes attributes, @Nullable String samplingRuleName) {
26+
this.decision = decision;
27+
this.attributes = attributes;
28+
this.samplingRuleName = samplingRuleName;
29+
}
30+
31+
static AwsSamplingResult create(
32+
SamplingDecision decision, Attributes attributes, @Nullable String samplingRuleName) {
33+
return new AwsSamplingResult(decision, attributes, samplingRuleName);
34+
}
35+
36+
@Override
37+
public SamplingDecision getDecision() {
38+
return decision;
39+
}
40+
41+
@Override
42+
public Attributes getAttributes() {
43+
return attributes;
44+
}
45+
46+
@Override
47+
public TraceState getUpdatedTraceState(TraceState parentTraceState) {
48+
if (parentTraceState.get(AWS_XRAY_SAMPLING_RULE_TRACE_STATE_KEY) == null
49+
&& this.samplingRuleName != null) {
50+
return parentTraceState.toBuilder()
51+
.put(AWS_XRAY_SAMPLING_RULE_TRACE_STATE_KEY, samplingRuleName)
52+
.build();
53+
}
54+
return parentTraceState;
55+
}
56+
}

‎aws-xray/src/main/java/io/opentelemetry/contrib/awsxray/AwsXrayAdaptiveSamplingConfig.java‎

Lines changed: 148 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,148 @@
1+
/*
2+
* Copyright The OpenTelemetry Authors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package io.opentelemetry.contrib.awsxray;
7+
8+
import com.fasterxml.jackson.annotation.JsonCreator;
9+
import com.fasterxml.jackson.annotation.JsonProperty;
10+
import com.fasterxml.jackson.annotation.JsonValue;
11+
import com.fasterxml.jackson.databind.annotation.JsonDeserialize;
12+
import com.fasterxml.jackson.databind.annotation.JsonSerialize;
13+
import com.google.auto.value.AutoValue;
14+
import java.util.List;
15+
import javax.annotation.Nullable;
16+
17+
@AutoValue
18+
@JsonSerialize(as = AwsXrayAdaptiveSamplingConfig.class)
19+
@JsonDeserialize(builder = AutoValue_AwsXrayAdaptiveSamplingConfig.Builder.class)
20+
public abstract class AwsXrayAdaptiveSamplingConfig {
21+
22+
@JsonProperty("version")
23+
public abstract double getVersion();
24+
25+
@JsonProperty("anomalyConditions")
26+
@Nullable
27+
public abstract List<AnomalyConditions> getAnomalyConditions();
28+
29+
@JsonProperty("anomalyCaptureLimit")
30+
@Nullable
31+
public abstract AnomalyCaptureLimit getAnomalyCaptureLimit();
32+
33+
public static Builder builder() {
34+
return new AutoValue_AwsXrayAdaptiveSamplingConfig.Builder();
35+
}
36+
37+
@AutoValue.Builder
38+
public abstract static class Builder {
39+
@JsonProperty("version")
40+
public abstract Builder setVersion(double value);
41+
42+
@JsonProperty("anomalyConditions")
43+
public abstract Builder setAnomalyConditions(List<AnomalyConditions> value);
44+
45+
@JsonProperty("anomalyCaptureLimit")
46+
public abstract Builder setAnomalyCaptureLimit(AnomalyCaptureLimit value);
47+
48+
public abstract AwsXrayAdaptiveSamplingConfig build();
49+
}
50+
51+
@AutoValue
52+
@JsonDeserialize(
53+
builder = AutoValue_AwsXrayAdaptiveSamplingConfig_AnomalyConditions.Builder.class)
54+
public abstract static class AnomalyConditions {
55+
@JsonProperty("errorCodeRegex")
56+
@Nullable
57+
public abstract String getErrorCodeRegex();
58+
59+
@JsonProperty("operations")
60+
@Nullable
61+
public abstract List<String> getOperations();
62+
63+
@JsonProperty("highLatencyMs")
64+
@Nullable
65+
public abstract Long getHighLatencyMs();
66+
67+
@JsonProperty("usage")
68+
@Nullable
69+
public abstract UsageType getUsage();
70+
71+
public static Builder builder() {
72+
return new AutoValue_AwsXrayAdaptiveSamplingConfig_AnomalyConditions.Builder();
73+
}
74+
75+
@AutoValue.Builder
76+
public abstract static class Builder {
77+
@JsonProperty("errorCodeRegex")
78+
public abstract Builder setErrorCodeRegex(String value);
79+
80+
@JsonProperty("operations")
81+
public abstract Builder setOperations(List<String> value);
82+
83+
@JsonProperty("highLatencyMs")
84+
public abstract Builder setHighLatencyMs(Long value);
85+
86+
@JsonProperty("usage")
87+
public abstract Builder setUsage(UsageType value);
88+
89+
public abstract AnomalyConditions build();
90+
}
91+
}
92+
93+
public enum UsageType {
94+
BOTH("both"),
95+
SAMPLING_BOOST("sampling-boost"),
96+
ANOMALY_TRACE_CAPTURE("anomaly-trace-capture"),
97+
NEITHER("neither"); // Not meant to be used by customers
98+
99+
private final String value;
100+
101+
UsageType(String value) {
102+
this.value = value;
103+
}
104+
105+
@JsonValue
106+
public String getValue() {
107+
return value;
108+
}
109+
110+
@JsonCreator
111+
public static UsageType fromValue(String value) {
112+
for (UsageType type : values()) {
113+
if (type.value.equals(value)) {
114+
return type;
115+
}
116+
}
117+
throw new IllegalArgumentException("Invalid usage value: " + value);
118+
}
119+
120+
public static boolean isUsedForBoost(UsageType usage) {
121+
return BOTH.equals(usage) || SAMPLING_BOOST.equals(usage);
122+
}
123+
124+
public static boolean isUsedForAnomalyTraceCapture(UsageType usage) {
125+
return BOTH.equals(usage) || ANOMALY_TRACE_CAPTURE.equals(usage);
126+
}
127+
}
128+
129+
@AutoValue
130+
@JsonDeserialize(
131+
builder = AutoValue_AwsXrayAdaptiveSamplingConfig_AnomalyCaptureLimit.Builder.class)
132+
public abstract static class AnomalyCaptureLimit {
133+
@JsonProperty("anomalyTracesPerSecond")
134+
public abstract int getAnomalyTracesPerSecond();
135+
136+
public static Builder builder() {
137+
return new AutoValue_AwsXrayAdaptiveSamplingConfig_AnomalyCaptureLimit.Builder();
138+
}
139+
140+
@AutoValue.Builder
141+
public abstract static class Builder {
142+
@JsonProperty("anomalyTracesPerSecond")
143+
public abstract Builder setAnomalyTracesPerSecond(int value);
144+
145+
public abstract AnomalyCaptureLimit build();
146+
}
147+
}
148+
}

0 commit comments

Comments
 (0)