replace netty Promise with Mono

2024-04-21 10:21:39 +00:00 · 2020-09-11 17:47:43 +08:00
parent b9ce500a97
commit da8c26b671
7 changed files with 236 additions and 315 deletions
@@ -2,7 +2,6 @@ package com.alibaba.arthas.channel.server;

 import org.springframework.boot.SpringApplication;
 import org.springframework.boot.autoconfigure.SpringBootApplication;
-import org.springframework.scheduling.annotation.EnableScheduling;

@SpringBootApplication
 public class ArthasChannelApplication {
@@ -7,7 +7,7 @@ package com.alibaba.arthas.channel.server.api;
 public class ApiResponse<T> {
    private String agentId;
    private String requestId;
-    private ApiState state;
+    private ApiStatus status;
    private String message;
    private String sessionId;
    // private String consumerId;
@@ -31,12 +31,12 @@ public class ApiResponse<T> {
        return this;
    }

-    public ApiState getState() {
-        return state;
+    public ApiStatus getStatus() {
+        return status;
    }

-    public ApiResponse<T> setState(ApiState state) {
-        this.state = state;
+    public ApiResponse<T> setStatus(ApiStatus status) {
+        this.status = status;
        return this;
    }

@@ -5,7 +5,7 @@ package com.alibaba.arthas.channel.server.api;
 *
 * @author gongdewei 2020-03-19
 */
-public enum ApiState {
+public enum ApiStatus {
    /**
     * Response is CONTINUOUS, receiving streaming data of async exec job
     */
@@ -1,9 +1,7 @@
 package com.alibaba.arthas.channel.server.service;

+import com.alibaba.arthas.channel.proto.ActionRequest;
 import com.alibaba.arthas.channel.proto.ActionResponse;
-import com.alibaba.arthas.channel.server.api.ApiRequest;
-import com.alibaba.arthas.channel.server.api.ApiResponse;
-import io.netty.util.concurrent.Promise;
 import reactor.core.publisher.Mono;

 /**
@@ -11,24 +9,22 @@ import reactor.core.publisher.Mono;
 */
 public interface ApiActionDelegateService {

-    Promise<ApiResponse> initSession(String agentId) throws Exception;
+    Mono<ActionResponse> initSession(String agentId) throws Exception;

-//    Promise<ApiResponse> joinSession(String agentId, String sessionId) throws Exception;
+    Mono<ActionResponse> closeSession(String agentId, String sessionId) throws Exception;

-    Promise<ApiResponse> closeSession(String agentId, String sessionId) throws Exception;
+    Mono<ActionResponse> interruptJob(String agentId, String sessionId) throws Exception;

-    Promise<ApiResponse> interruptJob(String agentId, String sessionId) throws Exception;
+    Mono<ActionResponse> execCommand(String agentId, ActionRequest request) throws Exception;

-    Promise<ApiResponse> execCommand(String agentId, ApiRequest request) throws Exception;
-
-    ApiResponse asyncExecCommand(String agentId, ApiRequest request) throws Exception;
+    Mono<ActionResponse> asyncExecCommand(String agentId, ActionRequest request) throws Exception;


    /**
     * Open WebConsole and create new session
     * @return
     */
-    Promise<ActionResponse> openConsole(String agentId, int timeout) throws Exception;
+    Mono<ActionResponse> openConsole(String agentId, int timeout) throws Exception;

    /**
     * proxy pass WebConsole input
@@ -40,7 +36,7 @@ public interface ApiActionDelegateService {
     */
    void closeConsole(String agentId, String consoleId) throws Exception;

-    Mono<ApiResponse> pullResults(String agentId, String requestId, int timeout);
+    Mono<ActionResponse> pullResults(String agentId, String requestId, int timeout);

    void subscribeResults(String agentId, String requestId, int timeout, ResponseListener responseListener) throws Exception;

@@ -3,15 +3,8 @@ package com.alibaba.arthas.channel.server.service.impl;
 import com.alibaba.arthas.channel.proto.ActionRequest;
 import com.alibaba.arthas.channel.proto.ActionResponse;
 import com.alibaba.arthas.channel.proto.ConsoleParams;
-import com.alibaba.arthas.channel.proto.ExecuteParams;
-import com.alibaba.arthas.channel.proto.ExecuteResult;
 import com.alibaba.arthas.channel.proto.RequestAction;
 import com.alibaba.arthas.channel.proto.ResponseStatus;
-import com.alibaba.arthas.channel.proto.ResultFormat;
-import com.alibaba.arthas.channel.server.api.ApiAction;
-import com.alibaba.arthas.channel.server.api.ApiRequest;
-import com.alibaba.arthas.channel.server.api.ApiResponse;
-import com.alibaba.arthas.channel.server.api.ApiState;
 import com.alibaba.arthas.channel.server.message.MessageExchangeException;
 import com.alibaba.arthas.channel.server.message.MessageExchangeService;
 import com.alibaba.arthas.channel.server.message.topic.ActionRequestTopic;
@@ -19,8 +12,6 @@ import com.alibaba.arthas.channel.server.message.topic.ActionResponseTopic;
 import com.alibaba.arthas.channel.server.service.AgentManageService;
 import com.alibaba.arthas.channel.server.service.ApiActionDelegateService;
 import com.google.protobuf.StringValue;
-import io.netty.util.concurrent.GlobalEventExecutor;
-import io.netty.util.concurrent.Promise;
 import org.apache.commons.lang3.RandomStringUtils;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -44,34 +35,23 @@ public class ApiActionDelegateServiceImpl implements ApiActionDelegateService {


    @Override
-    public Promise<ApiResponse> initSession(String agentId) throws Exception {
-        ApiRequest apiRequest = new ApiRequest();
-        apiRequest.setAction(ApiAction.INIT_SESSION.name());
-        return sendRequestAndSubscribe(agentId, apiRequest);
-    }
-
-//    @Override
-//    public Promise<ApiResponse> joinSession(String agentId, String sessionId) throws Exception {
-//        ApiRequest apiRequest = new ApiRequest();
-//        apiRequest.setAction(ApiAction.JOIN_SESSION.name());
-//        apiRequest.setSessionId(sessionId);
-//        return sendRequestAndSubscribe(agentId, apiRequest);
-//    }
-
-    @Override
-    public Promise<ApiResponse> closeSession(String agentId, String sessionId) throws Exception {
-        ApiRequest apiRequest = new ApiRequest();
-        apiRequest.setAction(ApiAction.CLOSE_SESSION.name());
-        apiRequest.setSessionId(sessionId);
-        return sendRequestAndSubscribe(agentId, apiRequest);
+    public Mono<ActionResponse> initSession(String agentId) throws Exception {
+        return sendRequestAndSubscribe(agentId, ActionRequest.newBuilder()
+                .setAction(RequestAction.INIT_SESSION));
    }

    @Override
-    public Promise<ApiResponse> interruptJob(String agentId, String sessionId) throws Exception {
-        ApiRequest apiRequest = new ApiRequest();
-        apiRequest.setAction(ApiAction.INTERRUPT_JOB.name());
-        apiRequest.setSessionId(sessionId);
-        return sendRequestAndSubscribe(agentId, apiRequest);
+    public Mono<ActionResponse> closeSession(String agentId, String sessionId) throws Exception {
+        return sendRequestAndSubscribe(agentId, ActionRequest.newBuilder()
+                .setAction(RequestAction.CLOSE_SESSION)
+                .setSessionId(StringValue.of(sessionId)));
+    }
+
+    @Override
+    public Mono<ActionResponse> interruptJob(String agentId, String sessionId) throws Exception {
+        return sendRequestAndSubscribe(agentId, ActionRequest.newBuilder()
+                .setAction(RequestAction.INTERRUPT_JOB)
+                .setSessionId(StringValue.of(sessionId)));
    }

    /**
@@ -83,8 +63,8 @@ public class ApiActionDelegateServiceImpl implements ApiActionDelegateService {
     * @return
     */
    @Override
-    public Promise<ApiResponse> execCommand(String agentId, ApiRequest request) throws Exception {
-        return sendRequestAndSubscribe(agentId, request);
+    public Mono<ActionResponse> execCommand(String agentId, ActionRequest request) throws Exception {
+        return sendRequestAndSubscribe(agentId, request.toBuilder());
    }

    /**
@@ -95,27 +75,24 @@ public class ApiActionDelegateServiceImpl implements ApiActionDelegateService {
     * @return
     */
    @Override
-    public ApiResponse asyncExecCommand(final String agentId, ApiRequest request) throws Exception {
+    public Mono<ActionResponse> asyncExecCommand(final String agentId, ActionRequest request) throws Exception {
        //send request
        String requestId = generateRandomRequestId();
-        sendRequest(agentId, requestId, request);
+        sendRequest(agentId, requestId, request.toBuilder());

-        //TODO 获取JobId？
-        return new ApiResponse()
-                .setState(ApiState.CONTINUOUS)
-                .setRequestId(requestId);
+        // 获取JobId？
+        return Mono.just(ActionResponse.newBuilder()
+                .setStatus(ResponseStatus.CONTINUOUS)
+                .setRequestId(requestId)
+                .build());
    }

    @Override
-    public Promise<ActionResponse> openConsole(String agentId, int timeout) throws Exception {
+    public Mono<ActionResponse> openConsole(String agentId, int timeout) throws Exception {
        //send request
        String requestId = generateRandomRequestId();
-        ActionRequest actionRequest = ActionRequest.newBuilder()
-                .setAgentId(agentId)
-                .setRequestId(requestId)
-                .setAction(RequestAction.OPEN_CONSOLE)
-                .build();
-        messageExchangeService.pushMessage(new ActionRequestTopic(agentId), actionRequest.toByteArray());
+        sendRequest(agentId, requestId, ActionRequest.newBuilder()
+                .setAction(RequestAction.OPEN_CONSOLE));

        //subscribe response
        return subscribeResponse(agentId, requestId, timeout);
@@ -124,62 +101,55 @@ public class ApiActionDelegateServiceImpl implements ApiActionDelegateService {
    @Override
    public void consoleInput(String agentId, String consoleId, String inputData) throws Exception {
        //send request
-        ActionRequest actionRequest = ActionRequest.newBuilder()
-                .setAgentId(agentId)
-                .setRequestId(consoleId)
+        sendRequest(agentId, consoleId, ActionRequest.newBuilder()
                .setAction(RequestAction.CONSOLE_INPUT)
                .setConsoleParams(ConsoleParams.newBuilder()
                        .setConsoleId(consoleId)
-                        .setInputData(inputData))
-                .build();
-
-        messageExchangeService.pushMessage(new ActionRequestTopic(agentId), actionRequest.toByteArray());
+                        .setInputData(inputData)));
    }

    @Override
    public void closeConsole(String agentId, String consoleId) throws Exception {
        //send request
-        ActionRequest actionRequest = ActionRequest.newBuilder()
-                .setAgentId(agentId)
-                .setRequestId(consoleId)
+        sendRequest(agentId, consoleId, ActionRequest.newBuilder()
                .setAction(RequestAction.CLOSE_CONSOLE)
                .setConsoleParams(ConsoleParams.newBuilder()
-                        .setConsoleId(consoleId))
-                .build();
-        messageExchangeService.pushMessage(new ActionRequestTopic(agentId), actionRequest.toByteArray());
+                        .setConsoleId(consoleId)));
    }

    @Override
-    public Mono<ApiResponse> pullResults(final String agentId, String requestId, int timeout) {
+    public Mono<ActionResponse> pullResults(final String agentId, String requestId, int timeout) {
        //subscribe response
        ActionResponseTopic topic = new ActionResponseTopic(agentId, requestId);
-        Mono<ApiResponse> responseMono = messageExchangeService.pollMessage(topic, timeout)
-                .flatMap((Function<byte[], Mono<ApiResponse>>) messageBytes -> {
+        Mono<ActionResponse> responseMono = messageExchangeService.pollMessage(topic, timeout)
+                .flatMap((Function<byte[], Mono<ActionResponse>>) messageBytes -> {
                    try {
                        ActionResponse actionResponse = ActionResponse.parseFrom(messageBytes);
-                        ApiResponse apiResponse = convertApiResponse(actionResponse);
-                        return Mono.just(apiResponse);
+                        return Mono.just(actionResponse);
                    } catch (Throwable e) {
                        logger.error("process action response message failure", e);
-                        ApiResponse apiResponse = new ApiResponse()
+                        ActionResponse actionResponse = ActionResponse.newBuilder()
                                .setAgentId(agentId)
                                .setRequestId(requestId)
-                                .setState(ApiState.FAILED)
-                                .setMessage("process action response message failure");
-                        return Mono.just(apiResponse);
+                                .setStatus(ResponseStatus.FAILED)
+                                .setMessage(StringValue.of("process action response message failure"))
+                                .build();
+                        return Mono.just(actionResponse);
                    }
-                }).switchIfEmpty(Mono.just(new ApiResponse()
+                }).switchIfEmpty(Mono.just(ActionResponse.newBuilder()
                        .setAgentId(agentId)
                        .setRequestId(requestId)
-                        .setState(ApiState.FAILED)
-                        .setMessage("Timeout")))
+                        .setStatus(ResponseStatus.FAILED)
+                        .setMessage(StringValue.of("Timeout"))
+                        .build()))
                .onErrorResume(throwable -> {
                    logger.error("pull results error", throwable);
-                    return Mono.just(new ApiResponse()
+                    return Mono.just(ActionResponse.newBuilder()
                            .setAgentId(agentId)
                            .setRequestId(requestId)
-                            .setState(ApiState.FAILED)
-                            .setMessage(throwable.getMessage()));
+                            .setStatus(ResponseStatus.FAILED)
+                            .setMessage(StringValue.of(throwable.getMessage()))
+                            .build());
                });

        return responseMono;
@@ -221,175 +191,78 @@ public class ApiActionDelegateServiceImpl implements ApiActionDelegateService {
     * Send one-time request and subscribe it's response.
     * NOTE: This method do not support streaming results.
     * @param agentId
-     * @param request
+     * @param requestBuilder
     * @return
     */
-    private Promise<ApiResponse> sendRequestAndSubscribe(String agentId, ApiRequest request) throws Exception {
+    private Mono<ActionResponse> sendRequestAndSubscribe(String agentId, ActionRequest.Builder requestBuilder) throws Exception {
        //send request
        String requestId = generateRandomRequestId();
-        sendRequest(agentId, requestId, request);
+        ActionRequest actionRequest = sendRequest(agentId, requestId, requestBuilder);
+
+        int execTimeout = 30000;
+        if (actionRequest.hasExecuteParams()) {
+            int timeout = actionRequest.getExecuteParams().getExecTimeout();
+            if (timeout > 0) {
+                execTimeout = timeout;
+            }
+        }

        //subscribe response
-        return subscribeResponse1(agentId, requestId, request.getExecTimeout());
+        return subscribeResponse(agentId, requestId, execTimeout);
    }

-    private Promise<ApiResponse> subscribeResponse1(String agentId, String requestId, Integer timeout) throws MessageExchangeException {
-        final Promise<ApiResponse> promise = GlobalEventExecutor.INSTANCE.newPromise();
-        int execTimeout = 30000;
-        if (timeout != null && timeout > 0) {
-            execTimeout = timeout;
-        }
+    private Mono<ActionResponse> subscribeResponse(String agentId, String requestId, int timeout) throws MessageExchangeException {
+
        final ActionResponseTopic responseTopic = new ActionResponseTopic(agentId, requestId);
-        messageExchangeService.subscribe(responseTopic, execTimeout, new MessageExchangeService.MessageHandler() {
-            @Override
-            public boolean onMessage(byte[] messageBytes) {
-                try {
-                    ActionResponse actionResponse = ActionResponse.parseFrom(messageBytes);
-                    ApiResponse apiResponse = convertApiResponse(actionResponse);
+        return Mono.create(monoSink -> {
+            try {
+                messageExchangeService.subscribe(responseTopic, timeout, new MessageExchangeService.MessageHandler() {
+                    @Override
+                    public boolean onMessage(byte[] messageBytes) {
+                        try {
+                            ActionResponse actionResponse = ActionResponse.parseFrom(messageBytes);

-                    promise.setSuccess(apiResponse);
-                } catch (Throwable e) {
-                    logger.error("process response message failure: "+e.getMessage(), e);
-                    promise.setSuccess(new ApiResponse()
-                            .setState(ApiState.FAILED)
-                            .setMessage("process response message failure: "+e.getMessage()));
-                }
+                            monoSink.success(actionResponse);
+                        } catch (Throwable e) {
+                            logger.error("process response message failure: "+e.getMessage(), e);
+                            monoSink.success(ActionResponse.newBuilder()
+                                    .setStatus(ResponseStatus.FAILED)
+                                    .setMessage(StringValue.of("process response message failure: "+e.getMessage()))
+                                    .build());
+                        }

-                //promise is one-time subscribe, just remove it after received message
-                try {
-                    messageExchangeService.removeTopic(responseTopic);
-                } catch (Throwable e) {
-                    logger.error("remove topic failure", e);
-                }
-                return false;
+                        //promise is one-time subscribe, just remove it after received message
+                        try {
+                            messageExchangeService.removeTopic(responseTopic);
+                        } catch (Throwable e) {
+                            logger.error("remove topic failure", e);
+                        }
+                        return false;
+                    }
+
+                    @Override
+                    public boolean onTimeout() {
+                        monoSink.success(ActionResponse.newBuilder()
+                                .setStatus(ResponseStatus.FAILED)
+                                .setMessage(StringValue.of("timeout"))
+                                .build());
+                        return false;
+                    }
+                });
+            } catch (MessageExchangeException e) {
+                monoSink.error(e);
            }

-            @Override
-            public boolean onTimeout() {
-                promise.setSuccess(new ApiResponse()
-                        .setState(ApiState.FAILED)
-                        .setMessage("Timeout"));
-                return false;
-            }
        });
-        return promise;
    }

-    private Promise<ActionResponse> subscribeResponse(String agentId, String requestId, Integer timeout) throws MessageExchangeException {
-        final Promise<ActionResponse> promise = GlobalEventExecutor.INSTANCE.newPromise();
-        int execTimeout = 30000;
-        if (timeout != null && timeout > 0) {
-            execTimeout = timeout;
-        }
-        final ActionResponseTopic responseTopic = new ActionResponseTopic(agentId, requestId);
-        messageExchangeService.subscribe(responseTopic, execTimeout, new MessageExchangeService.MessageHandler() {
-            @Override
-            public boolean onMessage(byte[] messageBytes) {
-                try {
-                    ActionResponse actionResponse = ActionResponse.parseFrom(messageBytes);
-
-                    promise.setSuccess(actionResponse);
-                } catch (Throwable e) {
-                    logger.error("process response message failure: "+e.getMessage(), e);
-                    promise.setSuccess(ActionResponse.newBuilder()
-                            .setStatus(ResponseStatus.FAILED)
-                            .setMessage(StringValue.of("process response message failure: "+e.getMessage()))
-                            .build());
-                }
-
-                //promise is one-time subscribe, just remove it after received message
-                try {
-                    messageExchangeService.removeTopic(responseTopic);
-                } catch (Throwable e) {
-                    logger.error("remove topic failure", e);
-                }
-                return false;
-            }
-
-            @Override
-            public boolean onTimeout() {
-                promise.setSuccess(ActionResponse.newBuilder()
-                        .setStatus(ResponseStatus.FAILED)
-                        .setMessage(StringValue.of("timeout"))
-                        .build());
-                return false;
-            }
-        });
-        return promise;
-    }
-
-    private void sendRequest(String agentId, String requestId, ApiRequest request) throws Exception {
-        final RequestAction action = getAction(request.getAction());
-
-        ActionRequest.Builder actionRequestBuilder = ActionRequest.newBuilder()
+    private ActionRequest sendRequest(String agentId, String requestId, ActionRequest.Builder actionRequestBuilder) throws MessageExchangeException {
+        ActionRequest actionRequest = actionRequestBuilder
                .setAgentId(agentId)
                .setRequestId(requestId)
-                .setAction(action);
-
-        if (request.getSessionId() != null) {
-            actionRequestBuilder = actionRequestBuilder.setSessionId(StringValue.of(request.getSessionId()));
-        }
-//        if (request.getConsumerId() != null) {
-//            actionRequestBuilder = actionRequestBuilder.setConsumerId(StringValue.of(request.getConsumerId()));
-//        }
-        if (request.getCommand() != null) {
-            int execTimeout = request.getExecTimeout() !=null && request.getExecTimeout() > 0 ? request.getExecTimeout() : 30000;
-            actionRequestBuilder = actionRequestBuilder.setExecuteParams(ExecuteParams.newBuilder()
-                    .setResultFormat(ResultFormat.JSON)
-                    .setCommandLine(request.getCommand())
-                    .setExecTimeout(execTimeout)
-                    .build());
-        }
-        ActionRequest actionRequest = actionRequestBuilder.build();
+                .build();
        messageExchangeService.pushMessage(new ActionRequestTopic(agentId), actionRequest.toByteArray());
-    }
-
-    private ApiResponse convertApiResponse(ActionResponse actionResponse) {
-        ApiResponse apiResponse = new ApiResponse()
-                .setState(getState(actionResponse.getStatus()))
-                .setAgentId(actionResponse.getAgentId())
-                .setRequestId(actionResponse.getRequestId());
-
-        if (actionResponse.hasSessionId()) {
-            apiResponse.setSessionId(actionResponse.getSessionId().getValue());
-        }
-//        if (actionResponse.hasConsumerId()) {
-//            apiResponse.setConsumerId(actionResponse.getConsumerId().getValue());
-//        }
-        if (actionResponse.hasMessage()) {
-            apiResponse.setMessage(actionResponse.getMessage().getValue());
-        }
-        if (actionResponse.hasExecuteResult()) {
-            ExecuteResult executeResult = actionResponse.getExecuteResult();
-            if (executeResult.hasResultsJson()) {
-                apiResponse.setResult(executeResult.getResultsJson().getValue());
-            }
-        }
-        return apiResponse;
-    }
-
-    private RequestAction getAction(String action) {
-        ApiAction apiAction = ApiAction.valueOf(action.trim().toUpperCase());
-
-        switch (apiAction) {
-            case EXEC:
-                return RequestAction.EXECUTE;
-            case ASYNC_EXEC:
-                return RequestAction.ASYNC_EXECUTE;
-//            case JOIN_SESSION:
-//                return RequestAction.JOIN_SESSION;
-            case INIT_SESSION:
-                return RequestAction.INIT_SESSION;
-            case CLOSE_SESSION:
-                return RequestAction.CLOSE_SESSION;
-            case INTERRUPT_JOB:
-                return RequestAction.INTERRUPT_JOB;
-            case OPEN_CONSOLE:
-                return RequestAction.OPEN_CONSOLE;
-            case CONSOLE_INPUT:
-                return RequestAction.CONSOLE_INPUT;
-        }
-        throw new IllegalArgumentException("Unsupported request action: " + action);
+        return actionRequest;
    }

    private String generateRandomRequestId() {
@@ -397,23 +270,4 @@ public class ApiActionDelegateServiceImpl implements ApiActionDelegateService {
        return RandomStringUtils.random(12, true, true);
    }

-    private ApiState getState(ResponseStatus status) {
-        switch (status) {
-            case SUCCEEDED:
-                return ApiState.SUCCEEDED;
-            case REFUSED:
-                return ApiState.REFUSED;
-            case CONTINUOUS:
-                return ApiState.CONTINUOUS;
-            case INTERRUPTED:
-                return ApiState.INTERRUPTED;
-
-            case FAILED:
-            case UNRECOGNIZED:
-                return ApiState.FAILED;
-        }
-        return ApiState.FAILED;
-    }
-
-
 }
@@ -1,18 +1,23 @@
 package com.alibaba.arthas.channel.server.web;

+import com.alibaba.arthas.channel.proto.ActionRequest;
 import com.alibaba.arthas.channel.proto.ActionResponse;
+import com.alibaba.arthas.channel.proto.ExecuteParams;
 import com.alibaba.arthas.channel.proto.ExecuteResult;
+import com.alibaba.arthas.channel.proto.RequestAction;
 import com.alibaba.arthas.channel.proto.ResponseStatus;
+import com.alibaba.arthas.channel.proto.ResultFormat;
+import com.alibaba.arthas.channel.server.api.ApiAction;
 import com.alibaba.arthas.channel.server.api.ApiException;
 import com.alibaba.arthas.channel.server.conf.ScheduledExecutorConfig;
 import com.alibaba.arthas.channel.server.model.AgentVO;
 import com.alibaba.arthas.channel.server.api.ApiRequest;
 import com.alibaba.arthas.channel.server.api.ApiResponse;
-import com.alibaba.arthas.channel.server.api.ApiState;
+import com.alibaba.arthas.channel.server.api.ApiStatus;
 import com.alibaba.arthas.channel.server.service.ApiActionDelegateService;
 import com.alibaba.arthas.channel.server.service.AgentManageService;
 import com.alibaba.fastjson.JSON;
-import io.netty.util.concurrent.Promise;
+import com.google.protobuf.StringValue;
 import org.apache.commons.lang3.StringUtils;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -30,9 +35,9 @@ import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
 import reactor.core.publisher.Mono;

 import java.io.IOException;
+import java.time.Duration;
 import java.util.List;
 import java.util.Optional;
-import java.util.concurrent.TimeUnit;

 /**
 * @author gongdewei 2020/8/10
@@ -64,55 +69,63 @@ public class AgentController {
    }

    @RequestMapping("/agent/{agentId}/init_session")
-    public ApiResponse initSession(@PathVariable String agentId) {
+    public Mono<ApiResponse> initSession(@PathVariable String agentId) throws Exception {
        checkAgentExists(agentId);
        try {
-            Promise<ApiResponse> responsePromise = apiActionDelegateService.initSession(agentId);
-            ApiResponse apiResponse = responsePromise.get(30, TimeUnit.SECONDS);
-            return apiResponse;
+            return apiActionDelegateService.initSession(agentId)
+                    .timeout(Duration.ofMillis(30000))
+                    .map(AgentController::convertApiResponse);
        } catch (Throwable e) {
            logger.error("create session failure: " + e.toString(), e);
-            return new ApiResponse()
+            return Mono.just(ActionResponse.newBuilder()
                    .setAgentId(agentId)
-                    .setState(ApiState.FAILED)
-                    .setMessage("create session failure: " + e.toString());
+                    .setStatus(ResponseStatus.FAILED)
+                    .setMessage(StringValue.of("create session failure: " + e.toString()))
+                    .build())
+                    .map(AgentController::convertApiResponse);
        }
    }

    @RequestMapping("/agent/{agentId}/close_session/{sessionId}")
-    public ApiResponse closeSession(@PathVariable String agentId, @PathVariable String sessionId) {
+    public Mono<ApiResponse> closeSession(@PathVariable String agentId, @PathVariable String sessionId) {
        checkAgentExists(agentId);
        try {
-            Promise<ApiResponse> responsePromise = apiActionDelegateService.closeSession(agentId, sessionId);
-            return responsePromise.get(30, TimeUnit.SECONDS);
+            return apiActionDelegateService.closeSession(agentId, sessionId)
+                    .timeout(Duration.ofMillis(30000))
+                    .map(AgentController::convertApiResponse);
        } catch (Throwable e) {
            logger.error("close session failure: " + e.toString(), e);
-            return new ApiResponse()
+            return Mono.just(ActionResponse.newBuilder()
                    .setAgentId(agentId)
-                    .setSessionId(sessionId)
-                    .setState(ApiState.FAILED)
-                    .setMessage("close session failure: " + e.toString());
+                    .setSessionId(StringValue.of(sessionId))
+                    .setStatus(ResponseStatus.FAILED)
+                    .setMessage(StringValue.of("close session failure: " + e.toString()))
+                    .build())
+                    .map(AgentController::convertApiResponse);
        }
    }

    @RequestMapping("/agent/{agentId}/interrupt_job/{sessionId}")
-    public ApiResponse interruptJob(@PathVariable String agentId, @PathVariable String sessionId) {
+    public Mono<ApiResponse> interruptJob(@PathVariable String agentId, @PathVariable String sessionId) {
        checkAgentExists(agentId);
        try {
-            Promise<ApiResponse> responsePromise = apiActionDelegateService.interruptJob(agentId, sessionId);
-            return responsePromise.get(30, TimeUnit.SECONDS);
+            return apiActionDelegateService.interruptJob(agentId, sessionId)
+                    .timeout(Duration.ofMillis(30000))
+                    .map(AgentController::convertApiResponse);
        } catch (Throwable e) {
            logger.error("interrupt job failure: " + e.toString(), e);
-            return new ApiResponse()
+            return Mono.just(ActionResponse.newBuilder()
                    .setAgentId(agentId)
-                    .setSessionId(sessionId)
-                    .setState(ApiState.FAILED)
-                    .setMessage("interrupt job failure: " + e.toString());
+                    .setSessionId(StringValue.of(sessionId))
+                    .setStatus(ResponseStatus.FAILED)
+                    .setMessage(StringValue.of("interrupt job failure: " + e.toString()))
+                    .build())
+                    .map(AgentController::convertApiResponse);
        }
    }

    @PostMapping("/agent/{agentId}/exec")
-    public ApiResponse execCommand(@PathVariable String agentId, @RequestBody String requestBody) {
+    public Mono<ApiResponse> execCommand(@PathVariable String agentId, @RequestBody String requestBody) {
        try {
            ApiRequest apiRequest = parseRequest(requestBody);
            long execTimeout = 30000;
@@ -120,34 +133,45 @@ public class AgentController {
                execTimeout = apiRequest.getExecTimeout();
            }

-            Promise<ApiResponse> responsePromise = apiActionDelegateService.execCommand(agentId, apiRequest);
-            return responsePromise.get(execTimeout, TimeUnit.MILLISECONDS);
+            ActionRequest request = convertApiRequest(apiRequest);
+            return apiActionDelegateService.execCommand(agentId, request)
+                    .timeout(Duration.ofMillis(execTimeout))
+                    .map(AgentController::convertApiResponse);

        } catch (Throwable e) {
            logger.error("exec command failure: " + e.toString(), e);
-            return new ApiResponse()
-                    .setState(ApiState.FAILED)
-                    .setMessage("exec command failure: " + e.toString());
+            return Mono.just(ActionResponse.newBuilder()
+                    .setAgentId(agentId)
+                    .setStatus(ResponseStatus.FAILED)
+                    .setMessage(StringValue.of("exec command failure: " + e.toString()))
+                    .build())
+                    .map(AgentController::convertApiResponse);
        }
    }

    @PostMapping("/agent/{agentId}/async_exec")
-    public ApiResponse asyncExecCommand(@PathVariable String agentId, @RequestBody String requestBody) {
+    public Mono<ApiResponse> asyncExecCommand(@PathVariable String agentId, @RequestBody String requestBody) {
        try {
            ApiRequest apiRequest = parseRequest(requestBody);
-            return apiActionDelegateService.asyncExecCommand(agentId, apiRequest);
+            ActionRequest request = convertApiRequest(apiRequest);
+            return apiActionDelegateService.asyncExecCommand(agentId, request)
+                    .map(AgentController::convertApiResponse);
        } catch (Throwable e) {
            logger.error("async exec command failure: " + e.getMessage(), e);
-            return new ApiResponse()
-                    .setState(ApiState.FAILED)
-                    .setMessage("async exec command failure: " + e.getMessage());
+            return Mono.just(ActionResponse.newBuilder()
+                    .setAgentId(agentId)
+                    .setStatus(ResponseStatus.FAILED)
+                    .setMessage(StringValue.of("async exec command failure: " + e.toString()))
+                    .build())
+                    .map(AgentController::convertApiResponse);
        }
    }

    @RequestMapping("/agent/{agentId}/results/{requestId}")
    public Mono<ApiResponse> pullResults(@PathVariable String agentId, @PathVariable String requestId,
                                         @RequestParam(value = "timeout", defaultValue = "30000") final int timeout) {
-        return apiActionDelegateService.pullResults(agentId, requestId, timeout);
+        return apiActionDelegateService.pullResults(agentId, requestId, timeout)
+                .map(AgentController::convertApiResponse);
    }

    @GetMapping("/agent/{agentId}/sse_results/{requestId}")
@@ -166,13 +190,14 @@ public class AgentController {
        final SseEmitter emitter = new SseEmitter(Long.valueOf(timeout));
        try {
            ApiRequest apiRequest = parseRequest(requestBody);
-            ApiResponse apiResponse = apiActionDelegateService.asyncExecCommand(agentId, apiRequest);
-            String requestId = apiResponse.getRequestId();
+            ActionRequest request = convertApiRequest(apiRequest);
+            ActionResponse actionResponse = apiActionDelegateService.asyncExecCommand(agentId, request).block();
+            String requestId = actionResponse.getRequestId();
            subscribeResults(agentId, requestId, timeout, emitter);
        } catch (Throwable e) {
            logger.error("async exec command failure: " + e.getMessage(), e);
            ApiResponse response = new ApiResponse()
-                    .setState(ApiState.FAILED)
+                    .setStatus(ApiStatus.FAILED)
                    .setMessage("async exec command failure: " + e.getMessage());
            try {
                emitter.send(JSON.toJSONString(response));
@@ -223,9 +248,9 @@ public class AgentController {
        }
    }

-    private ApiResponse convertApiResponse(ActionResponse actionResponse) {
+    private static ApiResponse convertApiResponse(ActionResponse actionResponse) {
        ApiResponse apiResponse = new ApiResponse()
-                .setState(getState(actionResponse.getStatus()))
+                .setStatus(getStatus(actionResponse.getStatus()))
                .setAgentId(actionResponse.getAgentId())
                .setRequestId(actionResponse.getRequestId());

@@ -247,22 +272,68 @@ public class AgentController {
        return apiResponse;
    }

-    private ApiState getState(ResponseStatus status) {
+    private static ActionRequest convertApiRequest(ApiRequest request) throws Exception {
+        RequestAction action = getAction(request.getAction());
+        ActionRequest.Builder actionRequestBuilder = ActionRequest.newBuilder()
+                .setAction(action);
+
+        if (request.getSessionId() != null) {
+            actionRequestBuilder = actionRequestBuilder.setSessionId(StringValue.of(request.getSessionId()));
+        }
+//        if (request.getConsumerId() != null) {
+//            actionRequestBuilder = actionRequestBuilder.setConsumerId(StringValue.of(request.getConsumerId()));
+//        }
+        if (request.getCommand() != null) {
+            int execTimeout = request.getExecTimeout() !=null && request.getExecTimeout() > 0 ? request.getExecTimeout() : 30000;
+            actionRequestBuilder = actionRequestBuilder.setExecuteParams(ExecuteParams.newBuilder()
+                    .setResultFormat(ResultFormat.JSON)
+                    .setCommandLine(request.getCommand())
+                    .setExecTimeout(execTimeout)
+                    .build());
+        }
+        return actionRequestBuilder.build();
+    }
+
+    private static RequestAction getAction(String action) {
+        ApiAction apiAction = ApiAction.valueOf(action.trim().toUpperCase());
+
+        switch (apiAction) {
+            case EXEC:
+                return RequestAction.EXECUTE;
+            case ASYNC_EXEC:
+                return RequestAction.ASYNC_EXECUTE;
+//            case JOIN_SESSION:
+//                return RequestAction.JOIN_SESSION;
+            case INIT_SESSION:
+                return RequestAction.INIT_SESSION;
+            case CLOSE_SESSION:
+                return RequestAction.CLOSE_SESSION;
+            case INTERRUPT_JOB:
+                return RequestAction.INTERRUPT_JOB;
+            case OPEN_CONSOLE:
+                return RequestAction.OPEN_CONSOLE;
+            case CONSOLE_INPUT:
+                return RequestAction.CONSOLE_INPUT;
+        }
+        throw new IllegalArgumentException("Unsupported request action: " + action);
+    }
+
+    private static ApiStatus getStatus(ResponseStatus status) {
        switch (status) {
            case SUCCEEDED:
-                return ApiState.SUCCEEDED;
+                return ApiStatus.SUCCEEDED;
            case REFUSED:
-                return ApiState.REFUSED;
+                return ApiStatus.REFUSED;
            case CONTINUOUS:
-                return ApiState.CONTINUOUS;
+                return ApiStatus.CONTINUOUS;
            case INTERRUPTED:
-                return ApiState.INTERRUPTED;
+                return ApiStatus.INTERRUPTED;

            case FAILED:
            case UNRECOGNIZED:
-                return ApiState.FAILED;
+                return ApiStatus.FAILED;
        }
-        return ApiState.FAILED;
+        return ApiStatus.FAILED;
    }

    private ApiRequest parseRequest(String requestBody) throws ApiException {
@@ -24,6 +24,7 @@ import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.springframework.util.MultiValueMap;
 import org.springframework.web.util.UriComponentsBuilder;
+import reactor.core.publisher.Mono;

 import java.util.Collections;
 import java.util.List;
@@ -97,8 +98,8 @@ public class WebSocketFrameHandler extends SimpleChannelInboundHandler<WebSocket
            throw new IllegalArgumentException(error);
        }

-        Promise<ActionResponse> responsePromise = apiActionDelegateService.openConsole(agentId, 15000);
-        ActionResponse actionResponse = responsePromise.get();
+        Mono<ActionResponse> responseMono = apiActionDelegateService.openConsole(agentId, 15000);
+        ActionResponse actionResponse = responseMono.block();
        if (!actionResponse.getStatus().equals(ResponseStatus.SUCCEEDED)) {
            logger.error("open console failure, response: {}", actionResponse);
            throw new Exception("open console failure");