Propagate context to armeria callbacks
ADD
instrumentation/armeria-1.3/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/armeria/v1_3/AbstractStreamMessageSubscriptionInstrumentation.java
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.armeria.v1_3;
+
+import static net.bytebuddy.matcher.ElementMatchers.isConstructor;
+import static net.bytebuddy.matcher.ElementMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
+
+import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
+import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+import org.reactivestreams.Subscriber;
+
+public class AbstractStreamMessageSubscriptionInstrumentation implements TypeInstrumentation {
+ @Override
+ public ElementMatcher<TypeDescription> typeMatcher() {
+ return named("com.linecorp.armeria.common.stream.AbstractStreamMessage$SubscriptionImpl");
+ }
+
+ @Override
+ public void transform(TypeTransformer transformer) {
+ transformer.applyAdviceToMethod(
+ isConstructor()
+ .and(
+ takesArgument(0, named("com.linecorp.armeria.common.stream.AbstractStreamMessage")))
+ .and(takesArgument(1, named("org.reactivestreams.Subscriber"))),
+ AbstractStreamMessageSubscriptionInstrumentation.class.getName() + "$WrapSubscriberAdvice");
+ }
+
+ public static class WrapSubscriberAdvice {
+ @Advice.OnMethodEnter(suppress = Throwable.class)
+ public static void attachContext(
+ @Advice.Argument(value = 1, readOnly = false) Subscriber subscriber) {
+ subscriber = SubscriberWrapper.wrap(subscriber);
+ }
+ }
+}

MODIFY
instrumentation/armeria-1.3/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/armeria/v1_3/ArmeriaInstrumentationModule.java
instrumentation/armeria-1.3/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/armeria/v1_3/ArmeriaInstrumentationModule.java
@@ -29,6 +29,8 @@ public class ArmeriaInstrumentationModule extends InstrumentationModule {
@Override
public List<TypeInstrumentation> typeInstrumentations() {
return asList(
- new ArmeriaWebClientBuilderInstrumentation(), new ArmeriaServerBuilderInstrumentation());
+ new ArmeriaWebClientBuilderInstrumentation(),
+ new ArmeriaServerBuilderInstrumentation(),
+ new AbstractStreamMessageSubscriptionInstrumentation());
}
}

ADD
instrumentation/armeria-1.3/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/armeria/v1_3/SubscriberWrapper.java
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.armeria.v1_3;
+
+import io.opentelemetry.context.Context;
+import io.opentelemetry.context.Scope;
+import org.reactivestreams.Subscriber;
+import org.reactivestreams.Subscription;
+
+public class SubscriberWrapper implements Subscriber<Object> {
+ private final Subscriber<Object> delegate;
+ private final Context context;
+
+ private SubscriberWrapper(Subscriber<Object> delegate, Context context) {
+ this.delegate = delegate;
+ this.context = context;
+ }
+
+ public static Subscriber<Object> wrap(Subscriber<Object> delegate) {
+ Context context = Context.current();
+ if (context != Context.root()) {
+ return new SubscriberWrapper(delegate, context);
+ }
+ return delegate;
+ }
+
+ @Override
+ public void onSubscribe(Subscription subscription) {
+ try (Scope ignore = context.makeCurrent()) {
+ delegate.onSubscribe(subscription);
+ }
+ }
+
+ @Override
+ public void onNext(Object o) {
+ try (Scope ignore = context.makeCurrent()) {
+ delegate.onNext(o);
+ }
+ }
+
+ @Override
+ public void onError(Throwable throwable) {
+ try (Scope ignore = context.makeCurrent()) {
+ delegate.onError(throwable);
+ }
+ }
+
+ @Override
+ public void onComplete() {
+ try (Scope ignore = context.makeCurrent()) {
+ delegate.onComplete();
+ }
+ }
+}

MODIFY
instrumentation/armeria-1.3/library/src/test/groovy/io/opentelemetry/instrumentation/armeria/v1_3/ArmeriaHttpClientTest.groovy
instrumentation/armeria-1.3/library/src/test/groovy/io/opentelemetry/instrumentation/armeria/v1_3/ArmeriaHttpClientTest.groovy
@@ -26,4 +26,9 @@ class ArmeriaHttpClientTest extends AbstractArmeriaHttpClientTest implements Lib
boolean testCallbackWithParent() {
false
}
+
+ @Override
+ boolean testErrorWithCallback() {
+ return false
+ }
}

MODIFY
instrumentation/armeria-1.3/testing/src/main/groovy/io/opentelemetry/instrumentation/armeria/v1_3/AbstractArmeriaHttpClientTest.groovy
instrumentation/armeria-1.3/testing/src/main/groovy/io/opentelemetry/instrumentation/armeria/v1_3/AbstractArmeriaHttpClientTest.groovy
@@ -63,12 +63,6 @@ abstract class AbstractArmeriaHttpClientTest extends HttpClientTest<HttpRequest>
false
}
- // TODO: context not propagated to callback
- @Override
- boolean testErrorWithCallback() {
- return false
- }
-
@Override
List<AttributeKey<?>> extraAttributes() {
[

Propagate context to armeria callbacks (#3108)