diff --git a/instrumentation/rocketmq/rocketmq-client/rocketmq-client-java-5.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/rocketmqclientjava/v5_0/MapSetter.java b/instrumentation/rocketmq/rocketmq-client/rocketmq-client-java-5.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/rocketmqclientjava/v5_0/MapSetter.java index 2658e7318a94..21e73ce6efb5 100644 --- a/instrumentation/rocketmq/rocketmq-client/rocketmq-client-java-5.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/rocketmqclientjava/v5_0/MapSetter.java +++ b/instrumentation/rocketmq/rocketmq-client/rocketmq-client-java-5.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/rocketmqclientjava/v5_0/MapSetter.java @@ -6,6 +6,9 @@ package io.opentelemetry.javaagent.instrumentation.rocketmqclientjava.v5_0; import io.opentelemetry.context.propagation.TextMapSetter; +import io.opentelemetry.instrumentation.api.util.VirtualField; +import java.util.HashMap; +import java.util.Map; import javax.annotation.Nullable; import org.apache.rocketmq.client.java.message.PublishingMessageImpl; @@ -13,10 +16,17 @@ enum MapSetter implements TextMapSetter { INSTANCE; @Override - public void set(@Nullable PublishingMessageImpl carrier, String key, String value) { - if (carrier == null) { + public void set(@Nullable PublishingMessageImpl message, String key, String value) { + if (message == null) { return; } - carrier.setTraceContext(value); + VirtualField> virtualField = + VirtualField.find(PublishingMessageImpl.class, Map.class); + Map extraProperties = virtualField.get(message); + if (extraProperties == null) { + extraProperties = new HashMap<>(); + virtualField.set(message, extraProperties); + } + extraProperties.put(key, value); } } diff --git a/instrumentation/rocketmq/rocketmq-client/rocketmq-client-java-5.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/rocketmqclientjava/v5_0/RocketMqProducerInstrumentation.java b/instrumentation/rocketmq/rocketmq-client/rocketmq-client-java-5.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/rocketmqclientjava/v5_0/RocketMqProducerInstrumentation.java index c9dadb563f1f..3f2e6c0f15a2 100644 --- a/instrumentation/rocketmq/rocketmq-client/rocketmq-client-java-5.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/rocketmqclientjava/v5_0/RocketMqProducerInstrumentation.java +++ b/instrumentation/rocketmq/rocketmq-client/rocketmq-client-java-5.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/rocketmqclientjava/v5_0/RocketMqProducerInstrumentation.java @@ -6,6 +6,7 @@ package io.opentelemetry.javaagent.instrumentation.rocketmqclientjava.v5_0; import static net.bytebuddy.matcher.ElementMatchers.isMethod; +import static net.bytebuddy.matcher.ElementMatchers.isPrivate; import static net.bytebuddy.matcher.ElementMatchers.named; import static net.bytebuddy.matcher.ElementMatchers.takesArgument; import static net.bytebuddy.matcher.ElementMatchers.takesArguments; @@ -41,6 +42,7 @@ public void transform(TypeTransformer transformer) { transformer.applyAdviceToMethod( isMethod() .and(named("send0")) + .and(isPrivate()) .and(takesArguments(6)) .and( takesArgument( diff --git a/instrumentation/rocketmq/rocketmq-client/rocketmq-client-java-5.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/rocketmqclientjava/v5_0/RocketMqPublishingMessageImplInstrumentation.java b/instrumentation/rocketmq/rocketmq-client/rocketmq-client-java-5.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/rocketmqclientjava/v5_0/RocketMqPublishingMessageImplInstrumentation.java index ec398e598f5e..9cb30adbd8ea 100644 --- a/instrumentation/rocketmq/rocketmq-client/rocketmq-client-java-5.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/rocketmqclientjava/v5_0/RocketMqPublishingMessageImplInstrumentation.java +++ b/instrumentation/rocketmq/rocketmq-client/rocketmq-client-java-5.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/rocketmqclientjava/v5_0/RocketMqPublishingMessageImplInstrumentation.java @@ -6,38 +6,49 @@ package io.opentelemetry.javaagent.instrumentation.rocketmqclientjava.v5_0; import static net.bytebuddy.matcher.ElementMatchers.isConstructor; +import static net.bytebuddy.matcher.ElementMatchers.isMethod; +import static net.bytebuddy.matcher.ElementMatchers.isPublic; import static net.bytebuddy.matcher.ElementMatchers.named; +import static net.bytebuddy.matcher.ElementMatchers.namedOneOf; import static net.bytebuddy.matcher.ElementMatchers.takesArgument; import io.opentelemetry.context.Context; import io.opentelemetry.instrumentation.api.util.VirtualField; import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation; import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer; +import java.util.Map; import net.bytebuddy.asm.Advice; import net.bytebuddy.description.type.TypeDescription; import net.bytebuddy.matcher.ElementMatcher; import org.apache.rocketmq.client.apis.message.Message; import org.apache.rocketmq.client.apis.producer.Producer; import org.apache.rocketmq.client.apis.producer.Transaction; +import org.apache.rocketmq.client.java.message.MessageImpl; import org.apache.rocketmq.client.java.message.PublishingMessageImpl; final class RocketMqPublishingMessageImplInstrumentation implements TypeInstrumentation { @Override public ElementMatcher typeMatcher() { - return named("org.apache.rocketmq.client.java.message.PublishingMessageImpl"); + return namedOneOf( + "org.apache.rocketmq.client.java.message.PublishingMessageImpl", + "org.apache.rocketmq.client.java.message.MessageImpl"); } @Override public void transform(TypeTransformer transformer) { transformer.applyAdviceToMethod( isConstructor() + .and(isPublic()) .and(takesArgument(0, named("org.apache.rocketmq.client.apis.message.Message"))) .and( takesArgument( 1, named("org.apache.rocketmq.client.java.impl.producer.PublishingSettings"))) .and(takesArgument(2, boolean.class)), RocketMqPublishingMessageImplInstrumentation.class.getName() + "$ConstructorAdvice"); + transformer.applyAdviceToMethod( + isMethod().and(named("getProperties")).and(isPublic()), + RocketMqPublishingMessageImplInstrumentation.class.getName() + "$GetPropertiesAdvice"); } @SuppressWarnings("unused") @@ -59,4 +70,24 @@ public static void onExit(@Advice.This PublishingMessageImpl message) { } } } + + @SuppressWarnings("unused") + public static class GetPropertiesAdvice { + /** Update the message properties to propagate context recorded by {@link MapSetter}. */ + @Advice.OnMethodExit(suppress = Throwable.class) + public static void onExit( + @Advice.This MessageImpl messageImpl, + @Advice.Return(readOnly = false) Map properties) { + if (!(messageImpl instanceof PublishingMessageImpl)) { + return; + } + PublishingMessageImpl message = (PublishingMessageImpl) messageImpl; + VirtualField> virtualField = + VirtualField.find(PublishingMessageImpl.class, Map.class); + Map extraProperties = virtualField.get(message); + if (extraProperties != null) { + properties.putAll(extraProperties); + } + } + } }