调整

Author: MentosL

arex-instrumentation/mq/arex-integration-mqtt/src/main/java/io/arex/inst/mqtt/adapter/MessageAdapterImpl.java

@@ -10,7 +10,7 @@
 import java.nio.charset.StandardCharsets;
 
 /**
- * MessageAdapterImpl
+ * MessageImpl
  */
 public class MessageAdapterImpl implements MessageAdapter<MessageChannel, Message> {
 

arex-instrumentation/mq/arex-integration-mqtt/src/main/java/io/arex/inst/mqtt/inst/MQGenericInstrumentationV3.java

@@ -1,5 +1,6 @@
 package io.arex.inst.mqtt.inst;
 
+import io.arex.agent.bootstrap.internal.Pair;
 import io.arex.inst.extension.MethodInstrumentation;
 import io.arex.inst.extension.TypeInstrumentation;
 import io.arex.inst.mqtt.MQTTAdapterHelper;
@@ -40,13 +41,23 @@ public static class ArrivedAdvice {
 
         @Advice.OnMethodEnter(suppress = Throwable.class)
         public static void onEnter(@Advice.Argument(value = 0, readOnly = false) MessageChannel messageChannel,
-                                   @Advice.Argument(value = 1, readOnly = false) Message<?> message) {
-            MQTTAdapterHelper.onServiceEnter(MessageAdapterImpl.getInstance(), messageChannel, message);
+                                   @Advice.Argument(value = 1, readOnly = false) Message<?> message,
+                                   @Advice.Local("channelMsgPair") Pair<MessageChannel, Message> pair) {
+            Pair<MessageChannel, Message> messageChannelMessagePair =
+                    MQTTAdapterHelper.onServiceEnter(MessageAdapterImpl.getInstance(), messageChannel, message);
+            if (messageChannelMessagePair == null){
+                return;
+            }
+            pair = messageChannelMessagePair;
         }
 
         @Advice.OnMethodExit(suppress = Throwable.class)
         public static void onExit(@Advice.Argument(value = 0, readOnly = false) MessageChannel messageChannel,
-                                  @Advice.Argument(value = 1, readOnly = false) Message<?> message) {
+                                  @Advice.Argument(value = 1, readOnly = false) Message<?> message,
+                                  @Advice.Local("channelMsgPair") Pair<MessageChannel, Message> pair) {
+            if (pair != null){
+                message = pair.getSecond();
+            }
             MQTTAdapterHelper.onServiceExit(MessageAdapterImpl.getInstance(), messageChannel, message);
         }
 

arex-instrumentation/mq/arex-integration-mqtt/src/main/java/io/arex/inst/mqtt/warp/GenericMessageWarp.java

@@ -12,23 +12,18 @@
 public class GenericMessageWarp extends GenericMessage {
 
     private MessageHeaderWarp messageHeaderWarp;
-    private Object payload;
 
     public GenericMessageWarp(Object payload) {
         super(payload);
-        this.payload = payload;
-        this.messageHeaderWarp = null;
     }
 
     public GenericMessageWarp(Object payload, Map headers) {
         super(payload, headers);
-        this.payload = payload;
         this.messageHeaderWarp = new MessageHeaderWarp(headers);
     }
 
     public GenericMessageWarp(Object payload, MessageHeaders headers) {
         super(payload, headers);
-        this.payload = payload;
         this.messageHeaderWarp = new MessageHeaderWarp(headers);
     }
 

arex-instrumentation/mq/arex-integration-mqtt/src/main/java/io/arex/inst/mqtt/warp/MessageHeaderWarp.java

@@ -2,8 +2,6 @@
 
 import org.springframework.messaging.MessageHeaders;
 
-import java.util.Collections;
-import java.util.HashMap;
 import java.util.Map;
 import java.util.UUID;
 
@@ -13,45 +11,26 @@
  */
 public class MessageHeaderWarp extends MessageHeaders {
 
-    private final Map<String, Object> headers;
-    private final UUID id;
-    private final Long timestamp;
 
     public MessageHeaderWarp(MessageHeaders messageHeaders) {
         super(messageHeaders);
-        this.headers = new HashMap<>();
-        for (Map.Entry<String, Object> entry : messageHeaders.entrySet()) {
-            this.headers.put(entry.getKey(), entry.getValue());
-        }
-        this.id = messageHeaders.getId();
-        this.timestamp = messageHeaders.getTimestamp();
     }
 
     public MessageHeaderWarp(Map<String, Object> headers) {
         super(headers);
-        this.headers = headers;
-        this.id = getIdGenerator().generateId();
-        this.timestamp = System.currentTimeMillis();
     }
 
     public MessageHeaderWarp(Map<String, Object> headers, UUID id, Long timestamp) {
         super(headers, id, timestamp);
-        this.headers = headers;
-        this.id = id;
-        this.timestamp = timestamp;
     }
 
-
     public Object put(String key, Object value){
-        this.headers.put(key,value);
+        super.getRawHeaders().put(key,value);
         return value;
     }
 
     public void remove(String key){
-        this.headers.remove(key);
+        super.getRawHeaders().remove(key);
     }
 
-    public Map<String, Object> getHeaders() {
-        return Collections.unmodifiableMap(headers);
-    }
 }