Skip to content

Commit 7d004b5

Browse files
dao-junlaurit
dao-jun
and
laurit
authored
[improve][pulsar] Enhance SendCallback to address PIP-363 (#11648)
Co-authored-by: Lauri Tulmin <[email protected]> Co-authored-by: Lauri Tulmin <[email protected]>
1 parent 2b42bc2 commit 7d004b5

File tree

5 files changed

+118
-54
lines changed

5 files changed

+118
-54
lines changed

‎instrumentation/pulsar/pulsar-2.8/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/pulsar/v2_8/ProducerImplInstrumentation.java

+4-53
Original file line numberDiff line numberDiff line change
@@ -14,18 +14,14 @@
1414
import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
1515

1616
import io.opentelemetry.context.Context;
17-
import io.opentelemetry.context.Scope;
1817
import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
1918
import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
2019
import io.opentelemetry.javaagent.instrumentation.pulsar.v2_8.telemetry.PulsarRequest;
21-
import java.util.concurrent.CompletableFuture;
2220
import net.bytebuddy.asm.Advice;
2321
import net.bytebuddy.description.type.TypeDescription;
2422
import net.bytebuddy.matcher.ElementMatcher;
2523
import org.apache.pulsar.client.api.Message;
26-
import org.apache.pulsar.client.api.MessageId;
2724
import org.apache.pulsar.client.api.PulsarClient;
28-
import org.apache.pulsar.client.impl.MessageImpl;
2925
import org.apache.pulsar.client.impl.ProducerImpl;
3026
import org.apache.pulsar.client.impl.PulsarClientImpl;
3127
import org.apache.pulsar.client.impl.SendCallback;
@@ -73,7 +69,7 @@ public static class ProducerSendAsyncMethodAdvice {
7369
public static void before(
7470
@Advice.This ProducerImpl<?> producer,
7571
@Advice.Argument(value = 0) Message<?> message,
76-
@Advice.Argument(value = 1, readOnly = false) SendCallback callback) {
72+
@Advice.Argument(value = 1) SendCallback callback) {
7773
Context parent = Context.current();
7874
PulsarRequest request = PulsarRequest.create(message, VirtualFieldStore.extract(producer));
7975

@@ -82,54 +78,9 @@ public static void before(
8278
}
8379

8480
Context context = producerInstrumenter().start(parent, request);
85-
callback = new SendCallbackWrapper(context, request, callback);
86-
}
87-
}
88-
89-
public static class SendCallbackWrapper implements SendCallback {
90-
91-
private final Context context;
92-
private final PulsarRequest request;
93-
private final SendCallback delegate;
94-
95-
public SendCallbackWrapper(Context context, PulsarRequest request, SendCallback callback) {
96-
this.context = context;
97-
this.request = request;
98-
this.delegate = callback;
99-
}
100-
101-
@Override
102-
public void sendComplete(Exception e) {
103-
if (context == null) {
104-
this.delegate.sendComplete(e);
105-
return;
106-
}
107-
108-
try (Scope ignore = context.makeCurrent()) {
109-
this.delegate.sendComplete(e);
110-
} finally {
111-
producerInstrumenter().end(context, request, null, e);
112-
}
113-
}
114-
115-
@Override
116-
public void addCallback(MessageImpl<?> msg, SendCallback scb) {
117-
this.delegate.addCallback(msg, scb);
118-
}
119-
120-
@Override
121-
public SendCallback getNextSendCallback() {
122-
return this.delegate.getNextSendCallback();
123-
}
124-
125-
@Override
126-
public MessageImpl<?> getNextMessage() {
127-
return this.delegate.getNextMessage();
128-
}
129-
130-
@Override
131-
public CompletableFuture<MessageId> getFuture() {
132-
return this.delegate.getFuture();
81+
// Inject the context/request into the SendCallback. This will be extracted and used when the
82+
// message is sent and the callback is invoked. see `SendCallbackInstrumentation`.
83+
VirtualFieldStore.inject(callback, context, request);
13384
}
13485
}
13586
}

‎instrumentation/pulsar/pulsar-2.8/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/pulsar/v2_8/PulsarInstrumentationModule.java

+2-1
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@ public List<TypeInstrumentation> typeInstrumentations() {
2424
new ConsumerImplInstrumentation(),
2525
new ProducerImplInstrumentation(),
2626
new MessageInstrumentation(),
27-
new MessageListenerInstrumentation());
27+
new MessageListenerInstrumentation(),
28+
new SendCallbackInstrumentation());
2829
}
2930
}

‎instrumentation/pulsar/pulsar-2.8/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/pulsar/v2_8/SendCallbackData.java

+23
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,23 @@
1+
/*
2+
* Copyright The OpenTelemetry Authors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package io.opentelemetry.javaagent.instrumentation.pulsar.v2_8;
7+
8+
import io.opentelemetry.context.Context;
9+
import io.opentelemetry.javaagent.instrumentation.pulsar.v2_8.telemetry.PulsarRequest;
10+
11+
public final class SendCallbackData {
12+
public final Context context;
13+
public final PulsarRequest request;
14+
15+
private SendCallbackData(Context context, PulsarRequest request) {
16+
this.context = context;
17+
this.request = request;
18+
}
19+
20+
public static SendCallbackData create(Context context, PulsarRequest request) {
21+
return new SendCallbackData(context, request);
22+
}
23+
}

‎instrumentation/pulsar/pulsar-2.8/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/pulsar/v2_8/SendCallbackInstrumentation.java

+75
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,75 @@
1+
/*
2+
* Copyright The OpenTelemetry Authors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package io.opentelemetry.javaagent.instrumentation.pulsar.v2_8;
7+
8+
import static io.opentelemetry.javaagent.extension.matcher.AgentElementMatchers.hasClassesNamed;
9+
import static io.opentelemetry.javaagent.extension.matcher.AgentElementMatchers.hasSuperType;
10+
import static io.opentelemetry.javaagent.instrumentation.pulsar.v2_8.telemetry.PulsarSingletons.producerInstrumenter;
11+
import static net.bytebuddy.matcher.ElementMatchers.isMethod;
12+
import static net.bytebuddy.matcher.ElementMatchers.named;
13+
14+
import io.opentelemetry.context.Context;
15+
import io.opentelemetry.context.Scope;
16+
import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
17+
import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
18+
import io.opentelemetry.javaagent.instrumentation.pulsar.v2_8.telemetry.PulsarRequest;
19+
import net.bytebuddy.asm.Advice;
20+
import net.bytebuddy.description.type.TypeDescription;
21+
import net.bytebuddy.matcher.ElementMatcher;
22+
import org.apache.pulsar.client.impl.SendCallback;
23+
24+
public class SendCallbackInstrumentation implements TypeInstrumentation {
25+
26+
@Override
27+
public ElementMatcher<ClassLoader> classLoaderOptimization() {
28+
return hasClassesNamed("org.apache.pulsar.client.impl.SendCallback");
29+
}
30+
31+
@Override
32+
public ElementMatcher<TypeDescription> typeMatcher() {
33+
return hasSuperType(named("org.apache.pulsar.client.impl.SendCallback"));
34+
}
35+
36+
@Override
37+
public void transform(TypeTransformer transformer) {
38+
transformer.applyAdviceToMethod(
39+
isMethod().and(named("sendComplete")),
40+
SendCallbackInstrumentation.class.getName() + "$SendCallbackSendCompleteAdvice");
41+
}
42+
43+
@SuppressWarnings("unused")
44+
public static class SendCallbackSendCompleteAdvice {
45+
46+
@Advice.OnMethodEnter(suppress = Throwable.class)
47+
public static void onEnter(
48+
@Advice.This SendCallback callback,
49+
@Advice.Local("otelContext") Context otelContext,
50+
@Advice.Local("otelScope") Scope otelScope,
51+
@Advice.Local("otelRequest") PulsarRequest request) {
52+
// Extract the Context and PulsarRequest from the SendCallback instance.
53+
SendCallbackData callBackData = VirtualFieldStore.extract(callback);
54+
if (callBackData != null) {
55+
// If the extraction was successful, store the Context and PulsarRequest in local variables.
56+
otelContext = callBackData.context;
57+
request = callBackData.request;
58+
otelScope = otelContext.makeCurrent();
59+
}
60+
}
61+
62+
@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
63+
public static void onExit(
64+
@Advice.Argument(0) Throwable t,
65+
@Advice.Local("otelContext") Context otelContext,
66+
@Advice.Local("otelScope") Scope otelScope,
67+
@Advice.Local("otelRequest") PulsarRequest request) {
68+
if (otelScope != null) {
69+
// Close the Scope and end the span.
70+
otelScope.close();
71+
producerInstrumenter().end(otelContext, request, null, t);
72+
}
73+
}
74+
}
75+
}

‎instrumentation/pulsar/pulsar-2.8/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/pulsar/v2_8/VirtualFieldStore.java

+14
Original file line numberDiff line numberDiff line change
@@ -7,9 +7,11 @@
77

88
import io.opentelemetry.context.Context;
99
import io.opentelemetry.instrumentation.api.util.VirtualField;
10+
import io.opentelemetry.javaagent.instrumentation.pulsar.v2_8.telemetry.PulsarRequest;
1011
import org.apache.pulsar.client.api.Consumer;
1112
import org.apache.pulsar.client.api.Message;
1213
import org.apache.pulsar.client.api.Producer;
14+
import org.apache.pulsar.client.impl.SendCallback;
1315
import org.apache.pulsar.client.impl.TopicMessageImpl;
1416

1517
public class VirtualFieldStore {
@@ -19,6 +21,8 @@ public class VirtualFieldStore {
1921
VirtualField.find(Producer.class, ProducerData.class);
2022
private static final VirtualField<Consumer<?>, String> CONSUMER_FIELD =
2123
VirtualField.find(Consumer.class, String.class);
24+
private static final VirtualField<SendCallback, SendCallbackData> CALLBACK_FIELD =
25+
VirtualField.find(SendCallback.class, SendCallbackData.class);
2226

2327
private VirtualFieldStore() {}
2428

@@ -40,6 +44,12 @@ public static void inject(Consumer<?> instance, String serviceUrl) {
4044
CONSUMER_FIELD.set(instance, serviceUrl);
4145
}
4246

47+
public static void inject(SendCallback instance, Context context, PulsarRequest request) {
48+
if (instance != null) {
49+
CALLBACK_FIELD.set(instance, SendCallbackData.create(context, request));
50+
}
51+
}
52+
4353
public static Context extract(Message<?> instance) {
4454
if (instance instanceof TopicMessageImpl<?>) {
4555
TopicMessageImpl<?> topicMessage = (TopicMessageImpl<?>) instance;
@@ -59,4 +69,8 @@ public static ProducerData extract(Producer<?> instance) {
5969
public static String extract(Consumer<?> instance) {
6070
return CONSUMER_FIELD.get(instance);
6171
}
72+
73+
public static SendCallbackData extract(SendCallback instance) {
74+
return CALLBACK_FIELD.get(instance);
75+
}
6276
}

0 commit comments

Comments
 (0)