Skip to content

Commit 4ad5752

Browse files
lukeina2zlukeina2z
committed
xfix:1
1 parent 754d1d4 commit 4ad5752

File tree

5 files changed

+316
-0
lines changed

5 files changed

+316
-0
lines changed

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

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -204,6 +204,10 @@ private SdkTracerProviderBuilder customizeTracerProviderBuilder(
204204
tracerProviderBuilder.addSpanProcessor(
205205
AttributePropagatingSpanProcessorBuilder.create().build());
206206

207+
if (isLambdaEnvironment()) {
208+
tracerProviderBuilder.addSpanProcessor(new AwsLambdaSpanProcessor());
209+
}
210+
207211
// If running on Lambda, we just need to export 100% spans and skip generating any Application
208212
// Signals metrics.
209213
if (isLambdaEnvironment()

‎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
@@ -97,4 +97,7 @@ private AwsAttributeKeys() {}
9797
AttributeKey.stringKey("aws.bedrock.guardrail.id");
9898
static final AttributeKey<String> AWS_GUARDRAIL_ARN =
9999
AttributeKey.stringKey("aws.bedrock.guardrail.arn");
100+
101+
static final AttributeKey<Boolean> AWS_TRACE_LAMBDA_MULTIPLE_SERVER =
102+
AttributeKey.booleanKey("aws.trace.lambda.multiple-server");
100103
}

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

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,72 @@
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.api.trace.Span;
19+
import io.opentelemetry.context.Context;
20+
import io.opentelemetry.sdk.trace.ReadWriteSpan;
21+
import io.opentelemetry.sdk.trace.ReadableSpan;
22+
import io.opentelemetry.sdk.trace.SpanProcessor;
23+
import java.util.logging.Logger;
24+
import javax.annotation.concurrent.Immutable;
25+
26+
@Immutable
27+
public final class AwsLambdaSpanProcessor implements SpanProcessor {
28+
private static final Logger logger =
29+
Logger.getLogger(AwsApplicationSignalsCustomizerProvider.class.getName());
30+
private ReadWriteSpan lambdaServerSpan = null;
31+
32+
public AwsLambdaSpanProcessor() {}
33+
34+
@Override
35+
public void onStart(Context parentContext, ReadWriteSpan span) {
36+
if (AwsSpanProcessingUtil.isLambdaServerSpan(span)) {
37+
if (lambdaServerSpan != null) {
38+
logger.warning("More than one Lambda server span detected.");
39+
}
40+
lambdaServerSpan = span;
41+
logger.info("now lambda server span recorded!!!.");
42+
}
43+
44+
if (AwsSpanProcessingUtil.isServletServerSpan(span)) {
45+
if (lambdaServerSpan == null) {
46+
logger.warning("lambda server span is null when servlet span is detected.");
47+
}
48+
49+
Span parentSpan = Span.fromContextOrNull(parentContext);
50+
if (parentSpan != null && (parentSpan instanceof ReadableSpan)) {
51+
ReadableSpan parentReadableSpan = (ReadableSpan) parentSpan;
52+
if (parentReadableSpan.getSpanContext().getSpanId()
53+
== lambdaServerSpan.getSpanContext().getSpanId()) {
54+
lambdaServerSpan.setAttribute(AwsAttributeKeys.AWS_TRACE_LAMBDA_MULTIPLE_SERVER, true);
55+
}
56+
}
57+
}
58+
}
59+
60+
@Override
61+
public boolean isStartRequired() {
62+
return true;
63+
}
64+
65+
@Override
66+
public void onEnd(ReadableSpan span) {}
67+
68+
@Override
69+
public boolean isEndRequired() {
70+
return false;
71+
}
72+
}

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

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@
3737
import io.opentelemetry.api.trace.SpanContext;
3838
import io.opentelemetry.api.trace.SpanKind;
3939
import io.opentelemetry.sdk.common.InstrumentationScopeInfo;
40+
import io.opentelemetry.sdk.trace.ReadableSpan;
4041
import io.opentelemetry.sdk.trace.data.SpanData;
4142
import java.io.IOException;
4243
import java.io.InputStream;
@@ -68,6 +69,10 @@ final class AwsSpanProcessingUtil {
6869

6970
private static final String SQL_DIALECT_KEYWORDS_JSON = "configuration/sql_dialect_keywords.json";
7071

72+
static final AttributeKey<String> OTEL_SCOPE_NAME = AttributeKey.stringKey("otel.scope.name");
73+
static final String LAMBDA_SCOPE_PREFIX = "io.opentelemetry.aws-lambda-";
74+
static final String SERVLET_SCOPE_PREFIX = "io.opentelemetry.servlet-";
75+
7176
static List<String> getDialectKeywords() {
7277
try (InputStream jsonFile =
7378
AwsSpanProcessingUtil.class
@@ -91,6 +96,10 @@ static List<String> getDialectKeywords() {
9196
*/
9297
static String getIngressOperation(SpanData span) {
9398
if (isLambdaEnvironment()) {
99+
String op = generateIngressOperation(span);
100+
if (op != UNKNOWN_OPERATION) {
101+
return System.getenv(AWS_LAMBDA_FUNCTION_NAME_CONFIG) + " " + op;
102+
}
94103
return System.getenv(AWS_LAMBDA_FUNCTION_NAME_CONFIG) + "/FunctionHandler";
95104
}
96105
String operation = span.getName();
@@ -248,4 +257,18 @@ static boolean isDBSpan(SpanData span) {
248257
|| isKeyPresent(span, DB_OPERATION)
249258
|| isKeyPresent(span, DB_STATEMENT);
250259
}
260+
261+
static boolean isLambdaServerSpan(ReadableSpan span) {
262+
String scopeName = span.toSpanData().getInstrumentationScopeInfo().getName();
263+
return scopeName != null
264+
&& scopeName.startsWith(LAMBDA_SCOPE_PREFIX)
265+
&& SpanKind.SERVER == span.getKind();
266+
}
267+
268+
static boolean isServletServerSpan(ReadableSpan span) {
269+
String scopeName = span.toSpanData().getInstrumentationScopeInfo().getName();
270+
return scopeName != null
271+
&& scopeName.startsWith(SERVLET_SCOPE_PREFIX)
272+
&& SpanKind.SERVER == span.getKind();
273+
}
251274
}

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

Lines changed: 214 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,214 @@
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.junit.jupiter.api.Assertions.assertNotNull;
19+
import static org.junit.jupiter.api.Assertions.assertTrue;
20+
import static org.mockito.ArgumentMatchers.any;
21+
import static org.mockito.Mockito.*;
22+
import static org.mockito.Mockito.mock;
23+
import static org.mockito.Mockito.when;
24+
25+
import io.opentelemetry.api.common.AttributeKey;
26+
import io.opentelemetry.api.trace.Span;
27+
import io.opentelemetry.api.trace.SpanContext;
28+
import io.opentelemetry.api.trace.SpanKind;
29+
import io.opentelemetry.api.trace.Tracer;
30+
import io.opentelemetry.context.Context;
31+
import io.opentelemetry.sdk.OpenTelemetrySdk;
32+
import io.opentelemetry.sdk.common.InstrumentationScopeInfo;
33+
import io.opentelemetry.sdk.trace.ReadWriteSpan;
34+
import io.opentelemetry.sdk.trace.SdkTracerProvider;
35+
import io.opentelemetry.sdk.trace.data.SpanData;
36+
import java.util.HashMap;
37+
import java.util.Map;
38+
import org.junit.jupiter.api.BeforeEach;
39+
import org.junit.jupiter.api.Test;
40+
41+
class AwsLambdaSpanProcessorTest {
42+
43+
private AwsLambdaSpanProcessor processor;
44+
private ReadWriteSpan mockLambdaServerSpan;
45+
private SpanData mockLambdaSpanData;
46+
private InstrumentationScopeInfo mockLambdaScopeInfo;
47+
private Map<AttributeKey<?>, Object> attributeMapForLambdaSpan;
48+
private SpanContext mockSpanContext;
49+
50+
private ReadWriteSpan mockServletServerSpan;
51+
private SpanData mockServletSpanData;
52+
private InstrumentationScopeInfo mockServletScopeInfo;
53+
54+
@BeforeEach
55+
void setupMock() {
56+
processor = new AwsLambdaSpanProcessor();
57+
58+
mockLambdaServerSpan = mock(ReadWriteSpan.class);
59+
mockLambdaSpanData = mock(SpanData.class);
60+
mockLambdaScopeInfo = mock(InstrumentationScopeInfo.class);
61+
when(mockLambdaSpanData.getInstrumentationScopeInfo()).thenReturn(mockLambdaScopeInfo);
62+
when(mockLambdaScopeInfo.getName()).thenReturn(AwsSpanProcessingUtil.LAMBDA_SCOPE_PREFIX);
63+
when(mockLambdaServerSpan.toSpanData()).thenReturn(mockLambdaSpanData);
64+
when(mockLambdaServerSpan.getKind()).thenReturn(SpanKind.SERVER);
65+
66+
mockSpanContext = mock(SpanContext.class);
67+
when(mockSpanContext.getSpanId()).thenReturn("mock-lambda-span-id");
68+
when(mockLambdaServerSpan.getSpanContext()).thenReturn(mockSpanContext);
69+
70+
attributeMapForLambdaSpan = new HashMap<>();
71+
72+
when(mockLambdaServerSpan.setAttribute(any(AttributeKey.class), any()))
73+
.thenAnswer(
74+
invocation -> {
75+
AttributeKey<?> key = invocation.getArgument(0);
76+
Object value = invocation.getArgument(1);
77+
attributeMapForLambdaSpan.put(key, value);
78+
System.out.println("xxxxxx " + key.toString());
79+
return mockLambdaServerSpan;
80+
});
81+
82+
when(mockLambdaServerSpan.getAttribute(any(AttributeKey.class)))
83+
.thenAnswer(
84+
invocation -> {
85+
AttributeKey<?> key = invocation.getArgument(0);
86+
System.out.println("xxxxxx2 " + key.toString());
87+
return attributeMapForLambdaSpan.get(key);
88+
});
89+
90+
assertNotNull(mockLambdaServerSpan);
91+
92+
mockServletServerSpan = mock(ReadWriteSpan.class);
93+
mockServletSpanData = mock(SpanData.class);
94+
mockServletScopeInfo = mock(InstrumentationScopeInfo.class);
95+
when(mockServletSpanData.getInstrumentationScopeInfo()).thenReturn(mockServletScopeInfo);
96+
when(mockServletScopeInfo.getName()).thenReturn(AwsSpanProcessingUtil.SERVLET_SCOPE_PREFIX);
97+
when(mockServletServerSpan.toSpanData()).thenReturn(mockServletSpanData);
98+
when(mockServletServerSpan.getKind()).thenReturn(SpanKind.SERVER);
99+
}
100+
101+
// @Test
102+
// void testCreateSpansWithoutStarting() {
103+
// Tracer lambdaTracer =
104+
// GlobalOpenTelemetry.getTracer("io.opentelemetry.aws-lambda-core-1.0");
105+
// Tracer servletTracer =
106+
// GlobalOpenTelemetry.getTracer("io.opentelemetry.servlet-5.0");
107+
108+
// // Create the Lambda span builder
109+
// SpanBuilder lambdaSpanBuilder =
110+
// lambdaTracer.spanBuilder("AwsLambdaServerSpan")
111+
// .setSpanKind(SpanKind.SERVER);
112+
113+
// // Create the Servlet span builder with Lambda span as parent
114+
// SpanBuilder servletSpanBuilder =
115+
// servletTracer.spanBuilder("ServletServerSpan")
116+
// .setSpanKind(SpanKind.SERVER);
117+
118+
// // Verify that spans are created but not started
119+
// assert lambdaSpanBuilder != null;
120+
// assert servletSpanBuilder != null;
121+
// }
122+
123+
// @Test
124+
// void testOnStart_lambdaServerSpan() {
125+
//
126+
// when(mockSpan.toSpanData().getInstrumentationScopeInfo().getName()).thenReturn("AWS.Lambda");
127+
// when(mockSpan.getKind()).thenReturn(SpanKind.SERVER);
128+
129+
// processor.onStart(mockContext, mockSpan);
130+
131+
// // Verify that lambdaServerSpan is set
132+
// verify(mockSpan, never()).setAttribute(any(), any());
133+
// }
134+
135+
// @Test
136+
// void testOnStart_multipleLambdaServerSpans() {
137+
//
138+
// when(mockSpan.toSpanData().getInstrumentationScopeInfo().getName()).thenReturn("AWS.Lambda");
139+
// when(mockSpan.getKind()).thenReturn(SpanKind.SERVER);
140+
141+
// processor.onStart(mockContext, mockSpan);
142+
// processor.onStart(mockContext, mockSpan);
143+
144+
// // Logger should warn about multiple Lambda server spans (this part is hard
145+
// to verify in
146+
// unit tests, would need a logger spy)
147+
// }
148+
149+
// @Test
150+
// void testOnStart_servletServerSpan_withoutLambdaSpan() {
151+
//
152+
// when(mockSpan.toSpanData().getInstrumentationScopeInfo().getName()).thenReturn("AWS.Servlet");
153+
// when(mockSpan.getKind()).thenReturn(SpanKind.SERVER);
154+
155+
// processor.onStart(mockContext, mockSpan);
156+
157+
// // Verify that warning is logged due to missing lambda server span
158+
// }
159+
160+
@Test
161+
void testOnStart_servletServerSpan_withLambdaSpan() {
162+
163+
System.out.println("Running test...");
164+
165+
SdkTracerProvider tracerProvider = SdkTracerProvider.builder().build();
166+
OpenTelemetrySdk openTelemetrySdk =
167+
OpenTelemetrySdk.builder().setTracerProvider(tracerProvider).build();
168+
169+
// Tracer tracerLambda =
170+
// GlobalOpenTelemetry.getTracer(AwsSpanProcessingUtil.LAMBDA_SCOPE_PREFIX);
171+
Tracer tracerLambda =
172+
openTelemetrySdk.getTracerProvider().get(AwsSpanProcessingUtil.LAMBDA_SCOPE_PREFIX);
173+
174+
Span lambdaSpan =
175+
tracerLambda.spanBuilder("Lambda-Function").setSpanKind(SpanKind.SERVER).startSpan();
176+
177+
processor.onStart(Context.current(), (ReadWriteSpan) lambdaSpan);
178+
lambdaSpan.makeCurrent();
179+
180+
Tracer tracerServlet =
181+
openTelemetrySdk.getTracerProvider().get(AwsSpanProcessingUtil.SERVLET_SCOPE_PREFIX);
182+
Span servletServerSpan =
183+
tracerServlet.spanBuilder("Servlet-Function").setSpanKind(SpanKind.SERVER).startSpan();
184+
185+
processor.onStart(Context.current(), (ReadWriteSpan) servletServerSpan);
186+
187+
boolean retVal =
188+
((ReadWriteSpan) lambdaSpan)
189+
.getAttribute(AwsAttributeKeys.AWS_TRACE_LAMBDA_MULTIPLE_SERVER);
190+
191+
openTelemetrySdk.shutdown();
192+
openTelemetrySdk.close();
193+
assertTrue(retVal);
194+
195+
// ReadableSpan parentSpan = mock(ReadableSpan.class);
196+
// SpanContext parentSpanContext = mock(SpanContext.class);
197+
198+
// when(mockSpan.toSpanData().getInstrumentationScopeInfo().getName()).thenReturn("AWS.Lambda");
199+
// when(mockSpan.getKind()).thenReturn(SpanKind.SERVER);
200+
// when(parentSpan.getSpanContext()).thenReturn(parentSpanContext);
201+
202+
// processor.onStart(mockContext, mockSpan);
203+
// when(mockSpan.getSpanContext()).thenReturn(parentSpanContext);
204+
205+
// ReadWriteSpan servletSpan = mock(ReadWriteSpan.class);
206+
// when(servletSpan.toSpanData().getInstrumentationScopeInfo().getName()).thenReturn("AWS.Servlet");
207+
// when(servletSpan.getKind()).thenReturn(SpanKind.SERVER);
208+
209+
// processor.onStart(mockContext, servletSpan);
210+
211+
// // Verify that attribute is set on lambdaServerSpan
212+
// verify(mockSpan, atLeastOnce()).setAttribute(any(), any());
213+
}
214+
}

0 commit comments

Comments
 (0)