使用Websocket+ Runtime 将shell 脚本执行信息输出到前端

需求

公司要做一个类似jenkins 的发布平台，在执行shell 脚本的时候，需要并同步将shell的执行信息在前端页面展示，所以我用websocket + RuntimeUtil 实现了此功能。

jar包依赖

        <dependency>
            <groupId>org.springframework.boot</groupId>
            <artifactId>spring-boot-starter-websocket</artifactId>
            <version>2.3.1.RELEASE</version>
        </dependency>
        <dependency>
            <groupId>cn.hutool</groupId>
            <artifactId>hutool-all</artifactId>
            <version>5.3.8</version>
        </dependency>

代码实现

配置websocket

• WebSocketAutoConfig.java

@Configuration
@EnableWebSocket
public class WebSocketAutoConfig implements WebSocketConfigurer {

    @Override
    public void registerWebSocketHandlers(WebSocketHandlerRegistry registry) {
        // webSocket通道
        // 指定处理器和路径
        registry.addHandler(new WebSocketHandler(), "/websocket")
                // 指定自定义拦截器
                .addInterceptors(new WebSocketInterceptor())
                // 允许跨域
                .setAllowedOrigins("*");
        // sockJs通道
        registry.addHandler(new WebSocketHandler(), "/sock-js")
                .addInterceptors(new WebSocketInterceptor())
                .setAllowedOrigins("*")
                // 开启sockJs支持
                .withSockJS();
    }
}

• ControllerExceptionHandler.java

@Slf4j
@ControllerAdvice
public class ControllerExceptionHandler {

    @ResponseBody
    @ExceptionHandler(value = Exception.class)
    public Result<Boolean> exceptionHandler(Exception exception) {
        log.error("error: ", exception);
        if (exception instanceof MethodArgumentNotValidException) {
            MethodArgumentNotValidException methodArgumentNotValidException = (MethodArgumentNotValidException) exception;
            FieldError fieldError = methodArgumentNotValidException.getBindingResult().getFieldError();
            String defaultMessage = fieldError == null ? "参数错误" : fieldError.getDefaultMessage();
            return Result.fail(String.format("%s:%s", methodArgumentNotValidException.getClass().getName(), defaultMessage));
        }
        if (exception instanceof IllegalArgumentException) {
            IllegalArgumentException illegalArgumentException = (IllegalArgumentException) exception;
            return Result.fail(String.format("%s:%s", illegalArgumentException.getClass().getName(), illegalArgumentException.getMessage()));
        } else {
            return Result.fail(String.format("%s:%s", exception.getClass().getName(), exception.getMessage()));
        }
    }
}

• WebSocketHandler.java

@Component
public class WebSocketHandler extends AbstractWebSocketHandler {

    /**
     * 存储sessionId和webSocketSession
     * 需要注意的是，webSocketSession没有提供无参构造，不能进行序列化，也就不能通过redis存储
     * 在分布式系统中，要想别的办法实现webSocketSession共享
     */
    private static Map<String, WebSocketSession> sessionMap = new ConcurrentHashMap<>();
    private static Map<String, String> userMap = new ConcurrentHashMap<>();

    /**
     * 获取sessionId
     */
    private String getSessionId(WebSocketSession session) {
        if (session instanceof WebSocketServerSockJsSession) {
            // sock js 连接
            try {
                return ((WebSocketSession) FieldUtils.readField(session, "webSocketSession", true)).getId();
            } catch (IllegalAccessException e) {
                throw new RuntimeException("get sessionId error");
            }
        }
        return session.getId();
    }

    /**
     * webSocket连接创建后调用
     */
    @Override
    public void afterConnectionEstablished(WebSocketSession session) {
        // 获取参数
        String user = String.valueOf(session.getAttributes().get("user"));
        if (StringUtils.isBlank(user)) {
            return;
        }
        String sessionId = getSessionId(session);
        userMap.put(user, getSessionId(session));
        sessionMap.put(sessionId, session);
    }

    /**
     * 接收到消息会调用
     */
    @Override
    public void handleMessage(WebSocketSession session, WebSocketMessage<?> message) throws Exception {
        if (message instanceof TextMessage) {

        } else if (message instanceof BinaryMessage) {

        } else if (message instanceof PongMessage) {

        } else {
            System.out.println("Unexpected WebSocket message type: " + message);
        }
    }

    /**
     * 连接出错会调用
     */
    @Override
    public void handleTransportError(WebSocketSession session, Throwable exception) {
        sessionMap.remove(getSessionId(session));
    }

    /**
     * 连接关闭会调用
     */
    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) {
        sessionMap.remove(getSessionId(session));
    }

    @Override
    public boolean supportsPartialMessages() {
        return false;
    }

    /**
     * 后端发送消息
     */
    public void sendMessage(String user, String message) {
        try {
            String sessionId = userMap.get(user);
            if (StringUtils.isBlank(sessionId)) {
                return;
            }
            WebSocketSession session = sessionMap.get(sessionId);
            if (session == null) {
                return;
            }
            session.sendMessage(new TextMessage(message));
        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    public WebSocketSession getSession(String username) {
        String sessionId = userMap.get(username);
        return sessionMap.get(sessionId);
    }

}

配置Service

• MessageSendRepository.java

@Slf4j
@Component
public class MessageSendRepository {

    @Autowired
    private WebSocketHandler webSocketHandler;
    @Getter
    ThreadLocal<String> threadLocal = new ThreadLocal<>();

    public void sendWebSocketMessage(String message) {
        webSocketHandler.sendMessage(threadLocal.get(), message);
    }

    public String sendWebSocketMessage(Process process) {
        StringBuilder builder = new StringBuilder();
        String websocketToken = threadLocal.get();
        try (
                //执行信息
                BufferedReader br = new BufferedReader(new InputStreamReader(process.getInputStream()));
                //错误信息
                BufferedReader errorBr = new BufferedReader(new InputStreamReader(process.getErrorStream()));
        ) {
            //输出执行信息
            String inline;
            while ((inline = br.readLine()) != null) {
                //记录信息
                builder.append(inline);
                if (StringUtils.isBlank(inline)) {
                    webSocketHandler.sendMessage(websocketToken, "\n");
                    continue;
                }
                StringUtils.replace(inline, "<", "&lt;");
                StringUtils.replace(inline, ">", "&gt;");
                //发送websocket信息
                webSocketHandler.sendMessage(websocketToken, inline);
            }

            //输出异常信息
            String errorInline;
            while ((errorInline = errorBr.readLine()) != null) {
                //记录信息
                builder.append(errorInline);
                if (StringUtils.isBlank(errorInline)) {
                    webSocketHandler.sendMessage(websocketToken, "<font color='red'>" + errorInline + "</font>");
                    continue;
                }
                webSocketHandler.sendMessage(websocketToken, "\n");
            }
        } catch (Exception e) {
            log.error(e.getMessage(), e);
            webSocketHandler.sendMessage(websocketToken, "<font color='red'>" + e.getMessage() + "</font>");
        } finally {
            webSocketHandler.sendMessage(websocketToken, "End.");
        }
        return builder.toString();
    }
}

• VersionBuildServiceImpl.java

@Slf4j
@Component
public class VersionBuildServiceImpl {
  /**
     * 项目编译打包
     */
    private Boolean projectBuild(ProjectDevelopmentVO project, String releaseVersionNum, String websocketToken) {
        try {

            //调用脚本，并传入对应参数
            String shShellCmd = String.format(
                    "sh %s %s %s %s ",
                    project.getProjectName(),
                    project.getGitRepositoryPath(),
                    project.getGitBranchName(),
                    project.getJarPath());
            log.info("项目编译打包脚本: " + shShellCmd);
            String[] shShellCmdAry = new String[]{ShellConstants.CMD_BIN_SH, "-c", shShellCmd};
            
            //获取脚本执行的process
            Process process = RuntimeUtil.exec(shShellCmdAry);
            
            //发送执行信息给前端
            String shShellCmdRet = messageSendRepository.sendWebSocketMessage(process);

            log.info("shShellCmdRet:{} ", shShellCmdRet);
            return StringUtils.contains(shShellCmdRet, "===脚本执行完成，项目编译打包完成===");
        } catch (Exception e) {
            e.printStackTrace();
        }
        return Boolean.FALSE;
    }
}

js 测试脚本

这里是在console 里面输出的信息

<script>
    var websocket = null;
    if('WebSocket' in window){
        websocket = new WebSocket('ws://localhost:8100/websocket?user=admin');
    }else{
        alert('当前浏览器不支持websocket消息通知');
    }

    //连接成功建立的回调方法
    websocket.onopen = function (event) {
        console.log("ws建立连接成功");
    }

    //连接关闭的回调方法
    websocket.onclose = function (event) {
        console.log("ws连接关闭");
    }

    //接收到消息的回调方法
    websocket.onmessage = function (event) {
        /*setMessageInnerHTML(event.data);*/
       // alert("ws接收返回消息："+event.data);
        console.log("服务器返回消息： " + event.data);
        //弹窗提醒(要用到JQuary,所以要先引入JQuary)   播放音乐
        // $('#mymodal').modal('show')
    }

    //连接发生错误的回调方法
    websocket.onerror = function(event){
        alert('websocket通信发生错误！')
    }

    //监听窗口关闭事件，当窗口关闭时，主动去关闭websocket连接，防止连接还没断开就关闭窗口，server端会抛异常。
    window.onbeforeunload = function() {
        websocket.close();
    }
</script>

运行效果