websocket实现进度条功能

WebSocket‌ 是一种为实现全双工通信而设计的协议，同样位于应用层，尽管其握手阶段使用 HTTP 协议进行协议升级，但后续通信不再依赖 HTTP ‌

效果演示

引入springboot依赖

<dependency>
    <groupId>org.springframework.boot</groupId>
    <artifactId>spring-boot-starter-websocket</artifactId>
</dependency>

java代码

package cn.websocket.config;

import cn.websocket.server.ProgressWebSocketHandler;
import org.springframework.context.annotation.Configuration;
import org.springframework.web.socket.config.annotation.EnableWebSocket;
import org.springframework.web.socket.config.annotation.WebSocketConfigurer;
import org.springframework.web.socket.config.annotation.WebSocketHandlerRegistry;

@Configuration
@EnableWebSocket
public class WebSocketConfig1 implements WebSocketConfigurer {

    @Override
    public void registerWebSocketHandlers(WebSocketHandlerRegistry registry) {
        // 注册处理器，路径为 /ws/progress
        // setAllowedOrigins("*") 允许跨域（生产环境请指定具体域名）
        registry.addHandler(new ProgressWebSocketHandler(), "/ws/progress")
                .setAllowedOrigins("*");
    }
}

package cn.websocket.server;

import lombok.Data;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;
import com.fasterxml.jackson.databind.ObjectMapper;
import java.io.IOException;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
@Data
// 定义消息 DTO (简化版，实际可用 Record 或 Lombok)
class TaskMessage {
    public String type; // "START", "CANCEL"
    public String taskId;
    // getters/setters 省略
}
@Data
class ProgressMessage {
    public String taskId;
    public int progress; // 0-100
    public String status; // "RUNNING", "DONE", "ERROR"
    public String message;
    // 构造函数省略


    public ProgressMessage() {
    }

    public ProgressMessage(String taskId, int progress, String status, String message) {
        this.taskId = taskId;
        this.progress = progress;
        this.status = status;
        this.message = message;
    }
}

public class ProgressWebSocketHandler extends TextWebSocketHandler {

    // 线程安全的 Map 存储会话：Key=taskId, Value=Session
    private static final Map<String, WebSocketSession> sessionMap = new ConcurrentHashMap<>();

    // 线程池用于执行耗时任务
    private final ExecutorService taskExecutor = Executors.newFixedThreadPool(10);
    private final ObjectMapper objectMapper = new ObjectMapper();

    /**
     * 连接建立后触发
     */
    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        System.out.println("新连接建立: " + session.getId());
        // 可以在这里发送欢迎消息
    }

    /**
     * 收到前端消息时触发 (例如：前端发送 "开始任务")
     */
    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
        String payload = message.getPayload();
        TaskMessage req = objectMapper.readValue(payload, TaskMessage.class);

        if ("START".equals(req.getType())) {
            String taskId = req.getTaskId();

            // 1. 将会话存入 Map，关联 taskId
            sessionMap.put(taskId, session);

            // 2. 异步执行耗时任务
            taskExecutor.submit(() -> runLongTask(taskId));

            // 3. 立即回复确认
            sendMessage(session, new ProgressMessage(taskId, 0, "RUNNING", "任务已启动"));
        } else if ("CANCEL".equals(req.getType())) {
            // 处理取消逻辑 (需要配合任务中的中断标志位)
            System.out.println("收到取消请求: " + req.getTaskId());
        }
    }

    /**
     * 连接关闭时触发 (清理资源)
     */
    @Override
    public void afterConnectionClosed(WebSocketSession session, org.springframework.web.socket.CloseStatus status) throws Exception {
        System.out.println("连接关闭: " + session.getId());
        // 遍历 Map 移除该 session 对应的 taskId (简化起见，实际可能需要反向索引)
        sessionMap.values().remove(session);
    }

    /**
     * 模拟耗时任务
     */
    private void runLongTask(String taskId) {
        WebSocketSession session = sessionMap.get(taskId);
        if (session == null || !session.isOpen()) return;

        try {
            for (int i = 1; i <= 100; i++) {
                // 模拟业务逻辑耗时
                Thread.sleep(200);

                // 检查连接是否还活着，避免向断开的连接发消息报错
                if (!session.isOpen()) break;

                // 构建进度消息
                ProgressMessage msg = new ProgressMessage();
                msg.taskId = taskId;
                msg.progress = i;
                msg.status = i == 100 ? "DONE" : "RUNNING";
                msg.message = i == 100 ? "任务完成！" : "处理中...";

                // 推送消息
                sendMessage(session, msg);

                if (i == 100) {
                    sessionMap.remove(taskId); // 任务完成，移除映射
                    break;
                }
            }
        } catch (Exception e) {
            e.printStackTrace();
            try {
                if (session != null && session.isOpen()) {
                    sendMessage(session, new ProgressMessage(taskId, -1, "ERROR", e.getMessage()));
                }
            } catch (IOException ex) {
                ex.printStackTrace();
            }
        }
    }

    /**
     * 通用发送方法
     */
    private void sendMessage(WebSocketSession session, ProgressMessage data) throws IOException {
        if (session != null && session.isOpen()) {
            String json = objectMapper.writeValueAsString(data);
            session.sendMessage(new TextMessage(json));
        }
    }
}

前端代码

保存在src/main/resource/static中 index.html,启动springboot项目后直接localhost:8080访问即可

<!DOCTYPE html>
<html lang="zh">
<head>
    <meta charset="UTF-8">
    <title>WebSocket 进度条</title>
    <style>
        #progress-container { width: 100%; background-color: #f3f3f3; border-radius: 5px; margin-top: 20px;}
        #progress-bar { width: 0%; height: 30px; background-color: #4caf50; text-align: center; line-height: 30px; color: white; border-radius: 5px; transition: width 0.2s;}
    </style>
</head>
<body>

<h2>实时任务进度</h2>
<button onclick="startTask()">开始任务</button>

<div id="progress-container">
    <div id="progress-bar">0%</div>
</div>
<p id="status-text">等待开始...</p>

<script>
    let ws;
    const taskId = 'task_' + Date.now(); // 生成唯一任务ID

    function connect() {
        // 1. 建立 WebSocket 连接
        ws = new WebSocket('ws://localhost:2006/ws/progress');

        ws.onopen = function() {
            console.log('WebSocket 连接成功');
            document.getElementById('status-text').innerText = '连接已建立，准备启动...';
        };

        ws.onmessage = function(event) {
            // 2. 解析后端发来的 JSON
            const data = JSON.parse(event.data);

            if (data.taskId === taskId) {
                updateUI(data.progress, data.status, data.message);

                if (data.status === 'DONE' || data.status === 'ERROR') {
                    ws.close(); // 任务结束，关闭连接
                }
            }
        };

        ws.onerror = function(error) {
            console.error('WebSocket 错误:', error);
            document.getElementById('status-text').innerText = '连接错误';
        };

        ws.onclose = function() {
            console.log('连接关闭');
        };
    }

    function startTask() {
        if (!ws || ws.readyState !== WebSocket.OPEN) {
            connect();
            // 等待一小会儿确保连接建立再发消息，或者在 onopen 中发
            setTimeout(() => sendStart(), 500);
        } else {
            sendStart();
        }
    }

    function sendStart() {
        // 3. 发送启动指令
        const msg = {
            type: 'START',
            taskId: taskId
        };
        ws.send(JSON.stringify(msg));
        document.getElementById('status-text').innerText = '任务运行中...';
    }

    function updateUI(progress, status, message) {
        const bar = document.getElementById('progress-bar');
        const text = document.getElementById('status-text');

        bar.style.width = progress + '%';
        bar.innerText = progress + '%';
        text.innerText = message;

        if (status === 'DONE') {
            bar.style.backgroundColor = '#4caf50'; // 绿色
        } else if (status === 'ERROR') {
            bar.style.backgroundColor = '#f44336'; // 红色
        }
    }
</script>
</body>
</html>