feat: telnet invoke command with context argument. (#9974)

apache · Jun 27, 2022 · f1ca2b4 · f1ca2b4
1 parent 2a758e2
commit f1ca2b4
Showing 1 changed file with 37 additions and 13 deletions.
diff --git a/dubbo-plugin/dubbo-qos/src/main/java/org/apache/dubbo/qos/legacy/InvokeTelnetHandler.java b/dubbo-plugin/dubbo-qos/src/main/java/org/apache/dubbo/qos/legacy/InvokeTelnetHandler.java
@@ -16,6 +16,7 @@
  */
 package org.apache.dubbo.qos.legacy;
 
+import com.alibaba.fastjson.JSON;
 import org.apache.dubbo.common.extension.Activate;
 import org.apache.dubbo.common.utils.CollectionUtils;
 import org.apache.dubbo.common.utils.ReflectUtils;
@@ -24,15 +25,15 @@
 import org.apache.dubbo.remoting.telnet.TelnetHandler;
 import org.apache.dubbo.remoting.telnet.support.Help;
 import org.apache.dubbo.rpc.AppResponse;
+import org.apache.dubbo.rpc.RpcContext;
 import org.apache.dubbo.rpc.model.ApplicationModel;
 import org.apache.dubbo.rpc.model.MethodDescriptor;
 import org.apache.dubbo.rpc.model.ProviderModel;
 
-import com.alibaba.fastjson.JSON;
-
 import java.lang.reflect.Method;
 import java.util.ArrayList;
 import java.util.Collection;
+import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
@@ -43,14 +44,16 @@
  * InvokeTelnetHandler
  */
 @Activate
-@Help(parameter = "[service.]method(args) ", summary = "Invoke the service method.",
+@Help(parameter = "[service.]method(args) [{'attachment':{'a':'b'}}]", summary = "Invoke the service method.",
         detail = "Invoke the service method.")
 public class InvokeTelnetHandler implements TelnetHandler {
 
     public static final String INVOKE_MESSAGE_KEY = "telnet.invoke.method.message";
     public static final String INVOKE_METHOD_LIST_KEY = "telnet.invoke.method.list";
     public static final String INVOKE_METHOD_PROVIDER_KEY = "telnet.invoke.method.provider";
 
+    private static final String ATTACHMENT_KEY  = "attachment";
+
     @Override
     @SuppressWarnings("unchecked")
     public String telnet(Channel channel, String message) {
@@ -62,18 +65,19 @@ public String telnet(Channel channel, String message) {
 
         String service = (String) channel.getAttribute(ChangeTelnetHandler.SERVICE_KEY);
 
-        int i = message.indexOf("(");
+        int leftBracket = message.indexOf("(");
+        int rightBracket = message.indexOf(")");
 
-        if (i < 0 || !message.endsWith(")")) {
+        if (leftBracket < 0 || rightBracket < 0) {
             return "Invalid parameters, format: service.method(args)";
         }
 
-        String method = message.substring(0, i).trim();
-        String args = message.substring(i + 1, message.length() - 1).trim();
-        i = method.lastIndexOf(".");
-        if (i >= 0) {
-            service = method.substring(0, i).trim();
-            method = method.substring(i + 1).trim();
+        String method = message.substring(0, leftBracket).trim();
+        String args = message.substring(leftBracket + 1, rightBracket).trim();
+        leftBracket = method.lastIndexOf(".");
+        if (leftBracket >= 0) {
+            service = method.substring(0, leftBracket).trim();
+            method = method.substring(leftBracket + 1).trim();
         }
 
         List<Object> list;
@@ -82,6 +86,17 @@ public String telnet(Channel channel, String message) {
         } catch (Throwable t) {
             return "Invalid json argument, cause: " + t.getMessage();
         }
+
+        Map<Object, Object> context = Collections.emptyMap();
+        String contextJson = message.substring(Math.min(message.length(), rightBracket + 1)).trim();
+        if (StringUtils.isNotEmpty(contextJson)) {
+            try {
+                context = JSON.parseObject(contextJson, Map.class);
+            } catch (Throwable t) {
+                // parse context json string.
+            }
+        }
+
         StringBuilder buf = new StringBuilder();
         Method invokeMethod = null;
         ProviderModel selectedProvider = null;
@@ -137,13 +152,22 @@ public String telnet(Channel channel, String message) {
             Object[] array = realize(list.toArray(), invokeMethod.getParameterTypes(),
                     invokeMethod.getGenericParameterTypes());
             long start = System.currentTimeMillis();
+
+            // before invoke
+            if (context.get(ATTACHMENT_KEY) != null) {
+                RpcContext.getContext().setAttachments((Map<String, String>) context.get(ATTACHMENT_KEY));
+            }
+
             AppResponse result = new AppResponse();
             try {
-                Object o = invokeMethod.invoke(selectedProvider.getServiceInstance(), array);
-                result.setValue(o);
+                result.setValue(invokeMethod.invoke(selectedProvider.getServiceInstance(), array));
             } catch (Throwable t) {
                 result.setException(t);
+            } finally {
+                // after invoke
+                RpcContext.getContext().clearAttachments();
             }
+
             long end = System.currentTimeMillis();
             buf.append("\r\nresult: ");
             buf.append(JSON.toJSONString(result.recreate()));