Skip to content

Commit a4d0421

Browse files
ygreeygree
committed
Properly close scope, finish span, and handle errors in HttpResponse wrappers and instrumentation.
1 parent 8718f94 commit a4d0421

12 files changed

+193
-259
lines changed

‎dd-java-agent/instrumentation/openai-java/openai-java-3.0/src/main/java/datadog/trace/instrumentation/openai_java/ChatCompletionServiceAsyncInstrumentation.java‎

Lines changed: 16 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -62,21 +62,15 @@ public static void exit(
6262
@Advice.Enter final AgentScope scope,
6363
@Advice.Return(readOnly = false) CompletableFuture<HttpResponseFor<ChatCompletion>> future,
6464
@Advice.Thrown final Throwable err) {
65-
final AgentSpan span = scope.span();
66-
try {
67-
if (err != null) {
68-
DECORATE.onError(span, err);
69-
}
70-
if (future != null) {
71-
future =
72-
HttpResponseWrapper.wrapFuture(
73-
future, span, ChatCompletionDecorator.DECORATE::withChatCompletion);
74-
} else {
75-
span.finish();
76-
}
77-
} finally {
78-
scope.close();
65+
AgentSpan span = scope.span();
66+
if (err != null || future == null) {
67+
DECORATE.finishSpan(span, err);
68+
} else {
69+
future =
70+
HttpResponseWrapper.wrapFuture(
71+
future, span, ChatCompletionDecorator.DECORATE::withChatCompletion);
7972
}
73+
scope.close();
8074
}
8175
}
8276

@@ -96,22 +90,15 @@ public static void exit(
9690
@Advice.Return(readOnly = false)
9791
CompletableFuture<HttpResponseFor<StreamResponse<ChatCompletionChunk>>> future,
9892
@Advice.Thrown final Throwable err) {
99-
final AgentSpan span = scope.span();
100-
try {
101-
if (err != null) {
102-
DECORATE.onError(span, err);
103-
}
104-
if (future != null) {
105-
future =
106-
HttpStreamResponseWrapper.wrapFuture(
107-
future, span, ChatCompletionDecorator.DECORATE::withChatCompletionChunks);
108-
} else {
109-
span.finish();
110-
}
111-
DECORATE.beforeFinish(span);
112-
} finally {
113-
scope.close();
93+
AgentSpan span = scope.span();
94+
if (err != null || future == null) {
95+
DECORATE.finishSpan(span, err);
96+
} else {
97+
future =
98+
HttpStreamResponseWrapper.wrapFuture(
99+
future, span, ChatCompletionDecorator.DECORATE::withChatCompletionChunks);
114100
}
101+
scope.close();
115102
}
116103
}
117104
}

‎dd-java-agent/instrumentation/openai-java/openai-java-3.0/src/main/java/datadog/trace/instrumentation/openai_java/ChatCompletionServiceInstrumentation.java‎

Lines changed: 16 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -61,21 +61,15 @@ public static void exit(
6161
@Advice.Enter final AgentScope scope,
6262
@Advice.Return(readOnly = false) HttpResponseFor<ChatCompletion> response,
6363
@Advice.Thrown final Throwable err) {
64-
final AgentSpan span = scope.span();
65-
try {
66-
if (err != null) {
67-
DECORATE.onError(span, err);
68-
}
69-
if (response != null) {
70-
response =
71-
HttpResponseWrapper.wrap(
72-
response, span, ChatCompletionDecorator.DECORATE::withChatCompletion);
73-
}
74-
DECORATE.beforeFinish(span);
75-
} finally {
76-
scope.close();
77-
span.finish();
64+
AgentSpan span = scope.span();
65+
if (err != null || response == null) {
66+
DECORATE.finishSpan(span, err);
67+
} else {
68+
response =
69+
HttpResponseWrapper.wrap(
70+
response, span, ChatCompletionDecorator.DECORATE::withChatCompletion);
7871
}
72+
scope.close();
7973
}
8074
}
8175

@@ -96,22 +90,15 @@ public static void exit(
9690
@Advice.Return(readOnly = false)
9791
HttpResponseFor<StreamResponse<ChatCompletionChunk>> response,
9892
@Advice.Thrown final Throwable err) {
99-
final AgentSpan span = scope.span();
100-
try {
101-
if (err != null) {
102-
DECORATE.onError(span, err);
103-
}
104-
if (response != null) {
105-
response =
106-
HttpStreamResponseWrapper.wrap(
107-
response, span, ChatCompletionDecorator.DECORATE::withChatCompletionChunks);
108-
} else {
109-
DECORATE.beforeFinish(span);
110-
span.finish();
111-
}
112-
} finally {
113-
scope.close();
93+
AgentSpan span = scope.span();
94+
if (err != null || response == null) {
95+
DECORATE.finishSpan(span, err);
96+
} else {
97+
response =
98+
HttpStreamResponseWrapper.wrap(
99+
response, span, ChatCompletionDecorator.DECORATE::withChatCompletionChunks);
114100
}
101+
scope.close();
115102
}
116103
}
117104
}

‎dd-java-agent/instrumentation/openai-java/openai-java-3.0/src/main/java/datadog/trace/instrumentation/openai_java/CompletionServiceAsyncInstrumentation.java‎

Lines changed: 16 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -57,21 +57,15 @@ public static void exit(
5757
@Advice.Enter final AgentScope scope,
5858
@Advice.Return(readOnly = false) CompletableFuture<HttpResponseFor<Completion>> future,
5959
@Advice.Thrown final Throwable err) {
60-
final AgentSpan span = scope.span();
61-
try {
62-
if (err != null) {
63-
DECORATE.onError(span, err);
64-
}
65-
if (future != null) {
66-
future =
67-
HttpResponseWrapper.wrapFuture(
68-
future, span, CompletionDecorator.DECORATE::withCompletion);
69-
} else {
70-
span.finish();
71-
}
72-
} finally {
73-
scope.close();
60+
AgentSpan span = scope.span();
61+
if (err != null || future == null) {
62+
DECORATE.finishSpan(span, err);
63+
} else {
64+
future =
65+
HttpResponseWrapper.wrapFuture(
66+
future, span, CompletionDecorator.DECORATE::withCompletion);
7467
}
68+
scope.close();
7569
}
7670
}
7771

@@ -92,22 +86,15 @@ public static void exit(
9286
@Advice.Return(readOnly = false)
9387
CompletableFuture<HttpResponseFor<StreamResponse<Completion>>> future,
9488
@Advice.Thrown final Throwable err) {
95-
final AgentSpan span = scope.span();
96-
try {
97-
if (err != null) {
98-
DECORATE.onError(span, err);
99-
}
100-
if (future != null) {
101-
future =
102-
HttpStreamResponseWrapper.wrapFuture(
103-
future, span, CompletionDecorator.DECORATE::withCompletions);
104-
} else {
105-
span.finish();
106-
}
107-
DECORATE.beforeFinish(span);
108-
} finally {
109-
scope.close();
89+
AgentSpan span = scope.span();
90+
if (err != null || future == null) {
91+
DECORATE.finishSpan(span, err);
92+
} else {
93+
future =
94+
HttpStreamResponseWrapper.wrapFuture(
95+
future, span, CompletionDecorator.DECORATE::withCompletions);
11096
}
97+
scope.close();
11198
}
11299
}
113100
}

‎dd-java-agent/instrumentation/openai-java/openai-java-3.0/src/main/java/datadog/trace/instrumentation/openai_java/CompletionServiceInstrumentation.java‎

Lines changed: 18 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,6 @@
22

33
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
44
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
5-
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
65
import static datadog.trace.instrumentation.openai_java.OpenAiDecorator.DECORATE;
76
import static net.bytebuddy.matcher.ElementMatchers.isMethod;
87
import static net.bytebuddy.matcher.ElementMatchers.returns;
@@ -13,9 +12,7 @@
1312
import com.openai.core.http.StreamResponse;
1413
import com.openai.models.completions.Completion;
1514
import com.openai.models.completions.CompletionCreateParams;
16-
import datadog.context.ContextScope;
1715
import datadog.trace.agent.tooling.Instrumenter;
18-
import datadog.trace.api.llmobs.LLMObsContext;
1916
import datadog.trace.bootstrap.instrumentation.api.AgentScope;
2017
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
2118
import net.bytebuddy.asm.Advice;
@@ -48,41 +45,25 @@ public static class CreateAdvice {
4845
@Advice.OnMethodEnter(suppress = Throwable.class)
4946
public static AgentScope enter(
5047
@Advice.Argument(0) final CompletionCreateParams params,
51-
@Advice.FieldValue("clientOptions") ClientOptions clientOptions,
52-
@Advice.Local("llmScope") ContextScope llmScope) {
53-
AgentSpan span = startSpan(OpenAiDecorator.INSTRUMENTATION_NAME, OpenAiDecorator.SPAN_NAME);
54-
// TODO get span from context?
55-
DECORATE.afterStart(span);
56-
DECORATE.withClientOptions(span, clientOptions);
48+
@Advice.FieldValue("clientOptions") ClientOptions clientOptions) {
49+
AgentSpan span = DECORATE.startSpan(clientOptions);
5750
CompletionDecorator.DECORATE.withCompletionCreateParams(span, params);
58-
59-
llmScope = LLMObsContext.attach(span.context());
60-
// TODO should the agent span be activated via the context api or keep separate?
6151
return activateSpan(span);
6252
}
6353

6454
@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
6555
public static void exit(
6656
@Advice.Enter final AgentScope scope,
6757
@Advice.Return(readOnly = false) HttpResponseFor<Completion> response,
68-
@Advice.Thrown final Throwable err,
69-
@Advice.Local("llmScope") ContextScope llmScope) {
70-
final AgentSpan span = scope.span();
71-
try {
72-
if (err != null) {
73-
DECORATE.onError(span, err);
74-
}
75-
if (response != null) {
76-
response =
77-
HttpResponseWrapper.wrap(
78-
response, span, CompletionDecorator.DECORATE::withCompletion);
79-
}
80-
DECORATE.beforeFinish(span);
81-
} finally {
82-
scope.close();
83-
llmScope.close();
84-
span.finish();
58+
@Advice.Thrown final Throwable err) {
59+
AgentSpan span = scope.span();
60+
if (err != null || response == null) {
61+
DECORATE.finishSpan(span, err);
62+
} else {
63+
response =
64+
HttpResponseWrapper.wrap(response, span, CompletionDecorator.DECORATE::withCompletion);
8565
}
66+
scope.close();
8667
}
8768
}
8869

@@ -102,22 +83,15 @@ public static void exit(
10283
@Advice.Enter final AgentScope scope,
10384
@Advice.Return(readOnly = false) HttpResponseFor<StreamResponse<Completion>> response,
10485
@Advice.Thrown final Throwable err) {
105-
final AgentSpan span = scope.span();
106-
try {
107-
if (err != null) {
108-
DECORATE.onError(span, err);
109-
}
110-
if (response != null) {
111-
response =
112-
HttpStreamResponseWrapper.wrap(
113-
response, span, CompletionDecorator.DECORATE::withCompletions);
114-
} else {
115-
span.finish();
116-
}
117-
DECORATE.beforeFinish(span);
118-
} finally {
119-
scope.close();
86+
AgentSpan span = scope.span();
87+
if (err != null || response == null) {
88+
DECORATE.finishSpan(span, err);
89+
} else {
90+
response =
91+
HttpStreamResponseWrapper.wrap(
92+
response, span, CompletionDecorator.DECORATE::withCompletions);
12093
}
94+
scope.close();
12195
}
12296
}
12397
}

‎dd-java-agent/instrumentation/openai-java/openai-java-3.0/src/main/java/datadog/trace/instrumentation/openai_java/EmbeddingServiceInstrumentation.java‎

Lines changed: 8 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -48,21 +48,15 @@ public static void exit(
4848
@Advice.Enter final AgentScope scope,
4949
@Advice.Return(readOnly = false) HttpResponseFor<CreateEmbeddingResponse> response,
5050
@Advice.Thrown final Throwable err) {
51-
final AgentSpan span = scope.span();
52-
try {
53-
if (err != null) {
54-
DECORATE.onError(span, err);
55-
}
56-
if (response != null) {
57-
response =
58-
HttpResponseWrapper.wrap(
59-
response, span, EmbeddingDecorator.DECORATE::withCreateEmbeddingResponse);
60-
}
61-
DECORATE.beforeFinish(span);
62-
} finally {
63-
scope.close();
64-
span.finish();
51+
AgentSpan span = scope.span();
52+
if (err != null || response == null) {
53+
DECORATE.finishSpan(span, err);
54+
} else {
55+
response =
56+
HttpResponseWrapper.wrap(
57+
response, span, EmbeddingDecorator.DECORATE::withCreateEmbeddingResponse);
6558
}
59+
scope.close();
6660
}
6761
}
6862
}

‎dd-java-agent/instrumentation/openai-java/openai-java-3.0/src/main/java/datadog/trace/instrumentation/openai_java/HttpResponseWrapper.java‎

Lines changed: 21 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -9,21 +9,24 @@
99
import java.util.concurrent.CompletableFuture;
1010
import java.util.function.BiConsumer;
1111
import org.jetbrains.annotations.NotNull;
12+
import org.slf4j.Logger;
13+
import org.slf4j.LoggerFactory;
1214

1315
public final class HttpResponseWrapper<T> implements HttpResponseFor<T> {
16+
private static final Logger log = LoggerFactory.getLogger(HttpResponseWrapper.class);
1417

1518
public static <T> HttpResponseFor<T> wrap(
16-
HttpResponseFor<T> response, AgentSpan span, BiConsumer<AgentSpan, T> afterParse) {
19+
HttpResponseFor<T> response, AgentSpan span, BiConsumer<AgentSpan, T> decorate) {
1720
DECORATE.withHttpResponse(span, response.headers());
18-
return new HttpResponseWrapper<>(response, span, afterParse);
21+
return new HttpResponseWrapper<>(response, span, decorate);
1922
}
2023

2124
public static <T> CompletableFuture<HttpResponseFor<T>> wrapFuture(
2225
CompletableFuture<HttpResponseFor<T>> future,
2326
AgentSpan span,
24-
BiConsumer<AgentSpan, T> afterParse) {
27+
BiConsumer<AgentSpan, T> decorate) {
2528
return future
26-
.thenApply(response -> wrap(response, span, afterParse))
29+
.thenApply(response -> wrap(response, span, decorate))
2730
.whenComplete(
2831
(r, t) -> {
2932
try {
@@ -39,25 +42,32 @@ public static <T> CompletableFuture<HttpResponseFor<T>> wrapFuture(
3942

4043
private final HttpResponseFor<T> delegate;
4144
private final AgentSpan span;
42-
private final BiConsumer<AgentSpan, T> afterParse;
45+
private final BiConsumer<AgentSpan, T> decorate;
4346

4447
private HttpResponseWrapper(
45-
HttpResponseFor<T> delegate, AgentSpan span, BiConsumer<AgentSpan, T> afterParse) {
48+
HttpResponseFor<T> delegate, AgentSpan span, BiConsumer<AgentSpan, T> decorate) {
4649
this.delegate = delegate;
4750
this.span = span;
48-
this.afterParse = afterParse;
51+
this.decorate = decorate;
4952
}
5053

5154
@Override
5255
public T parse() {
56+
T parsed;
5357
try {
54-
T parsed = delegate.parse();
55-
afterParse.accept(span, parsed);
56-
return parsed;
58+
parsed = delegate.parse();
5759
} catch (Throwable err) {
58-
DECORATE.onError(span, err);
60+
DECORATE.finishSpan(span, err);
5961
throw err;
6062
}
63+
try {
64+
decorate.accept(span, parsed);
65+
} catch (Throwable t) {
66+
log.debug("Span decorator failed", t);
67+
} finally {
68+
DECORATE.finishSpan(span, null);
69+
}
70+
return parsed;
6171
}
6272

6373
@Override
@@ -79,7 +89,6 @@ public InputStream body() {
7989

8090
@Override
8191
public void close() {
82-
// span finished after the response is available
8392
delegate.close();
8493
}
8594
}

0 commit comments

Comments
 (0)