WebSocket 实时进度推送实操方法（AI 大模型评测平台的批量任务通知）

25000 个子任务的进度要实时刷在页面上，这不是加个定时器能解决的。平台跑一次批量评测就是 5000 条用例 × 5 个模型，用户盯着页面等结果，进度条要实时增长，还要能看清哪些失败、什么时间完成。我先试了 HTTP 轮询，发现 5 秒一次频率拉高了浪费、拉低了滞后；又考察了 SSE，发现它只能服务端单向推送，用户想取消或暂停任务根本没法做。最后选了 WebSocket，把整套实时推送链路搭了起来。下面把选型、实现和踩过的坑完整写一遍。

一、为什么必须 WebSocket

先把方案选型讲清楚。实时进度推送有三个候选方案，各有适用场景：

方案	实时性	双向	复杂度	适用
HTTP 轮询	差	客户端主动	低	简单状态
SSE	好	服务端→客户端	中	单向流（AI 输出）
WebSocket	好	全双工	中	实时进度、协作

批量任务进度是”服务端主动推+客户端偶尔取消/暂停”，选 WebSocket 最合适。如果只是展示模型输出的文字流，SSE 反而更轻；但批量任务的交互是双向的，WebSocket 的全双工优势没法替代。

二、协议选型

定了 WebSocket 之后，还要在”原生实现”和”上层协议”之间做取舍。我们列了一张表对比：

协议	优点	缺点	平台选择
原生 WebSocket	简单、零依赖	要自己实现 STOMP/重连	备选
STOMP over WebSocket	主题订阅、消息路由	协议复杂	❌
Socket.IO	自动重连、房间	不是纯 WS、Java 端弱	❌

平台选 Spring 原生 WebSocket（spring-boot-starter-websocket），轻量、易控。我们只有”按 batchId 推进度”一个订阅维度，用不上 STOMP 那套主题路由，原生 API 反而更直接。

三、握手与鉴权

WebSocket 握手就是 HTTP Upgrade，既然是 HTTP，鉴权这一步必须在握手阶段完成，否则等于把实时通道裸奔出去。注册处理器时挂上鉴权拦截器：

@Configuration
@EnableWebSocket
public class WebSocketConfig implements WebSocketConfigurer {
    @Override
    public void registerWebSocketHandlers(WebSocketHandlerRegistry registry) {
        registry.addHandler(batchProgressHandler, "/ws/batch/{batchId}")
                .addInterceptors(new AuthHandshakeInterceptor())
                .setAllowedOrigins("*");
    }
}

AuthHandshakeInterceptor 握手时校验 token：

public class AuthHandshakeInterceptor implements HandshakeInterceptor {
    @Override
    public boolean beforeHandshake(ServerHttpRequest req, 
                                    ServerHttpResponse resp,
                                    WebSocketHandler wsHandler,
                                    Map<String, Object> attrs) {
        String token = getToken(req);  // 从 query 拿 ?token=xxx
        if (token == null) {
            resp.setStatusCode(HttpStatus.UNAUTHORIZED);
            return false;
        }
        Long userId = JwtUtil.parse(token);
        attrs.put("userId", userId);
        return true;
    }
}

前端连接：new WebSocket("wss://eval.example.com/ws/batch/123?token=...")。token 放在 query 上是 WebSocket 的惯例做法，Cookie 在部分浏览器和跨域场景下不靠谱。

四、消息格式

消息格式是前后端约定的核心，我们用 JSON 文本帧，按 type 区分事件。三种主要消息：

// 服务端 → 客户端
{
  "type": "PROGRESS",
  "batchId": 123,
  "percent": 45,
  "done": 11250,
  "total": 25000,
  "failed": 23,
  "ts": 1692518400000
}

// 失败详情
{
  "type": "TASK_FAILED",
  "batchId": 123,
  "taskId": 456,
  "subtaskId": 789,
  "modelCode": "gpt-4o",
  "error": "RATE_LIMIT",
  "ts": 1692518400000
}

// 整批完成
{
  "type": "BATCH_DONE",
  "batchId": 123,
  "summary": {"total": 25000, "success": 24800, "failed": 200},
  "ts": 1692518400000
}

所有事件都带 batchId 和 ts，前端收到消息后可以按 batchId 过滤、按 ts 判断是否过期，避免多条连接的消息串台。

五、连接管理

服务端 Session 表

服务端要维护”batchId → 连接”的映射，才能把进度推给对的人。这个表用 ConcurrentHashMap 实现：

public class BatchProgressHandler extends TextWebSocketHandler {
    // batchId -> Set<WebSocketSession>
    private final Map<Long, Set<WebSocketSession>> sessions = new ConcurrentHashMap<>();
    
    @Override
    public void afterConnectionEstablished(WebSocketSession s) {
        Long batchId = getBatchId(s);
        Long userId = (Long) s.getAttributes().get("userId");
        
        // 鉴权：只能订阅自己 batch 的进度
        Batch batch = batchService.getById(batchId);
        if (!batch.getUserId().equals(userId)) {
            s.close(CloseStatus.NOT_ACCEPTABLE);
            return;
        }
        
        sessions.computeIfAbsent(batchId, k -> ConcurrentHashMap.newKeySet()).add(s);
        // 立即推一次当前快照（断线重连场景）
        sendSnapshot(batchId);
    }
    
    @Override
    public void afterConnectionClosed(WebSocketSession s, CloseStatus status) {
        Long batchId = getBatchId(s);
        Set<WebSocketSession> set = sessions.get(batchId);
        if (set != null) set.remove(s);
    }
}

建立连接后先推一次快照，这条在断线重连场景下特别关键，后面单独讲。连接的增删都在这两个回调里完成，不会出现”连接断了还在推”的脏数据。

推送服务

有了 session 表，推送就简单了。推送服务负责查进度、组消息、广播：

@Service
public class WebSocketPushService {
    @Autowired private BatchProgressHandler handler;
    
    public void pushProgress(long batchId) {
        BatchProgress p = batchService.getProgress(batchId);
        String msg = JsonUtil.toJson(new ProgressEvent("PROGRESS", batchId, p));
        broadcast(batchId, msg);
    }
    
    private void broadcast(long batchId, String msg) {
        Set<WebSocketSession> set = handler.getSessions(batchId);
        for (WebSocketSession s : set) {
            try {
                if (s.isOpen()) s.sendMessage(new TextMessage(msg));
            } catch (Exception e) {
                log.warn("WS 推失败", e);
            }
        }
    }
}

单个连接发送失败只记日志，不影响其他连接，保证一次批量任务的推送不会因为个别连接异常而中断。

六、与 MQ 消费者解耦

子任务完成事件来自 worker 进程，如果让 worker 直接调 WebSocket 推送，两个模块就耦合死了。我们用 MQ 在中间解耦，数据要经过三道关卡：

worker 把完成事件发到 MQ；
消费者写库并聚合进度；
推送服务广播到对应的 WebSocket 连接。

链路如下：

子任务完成 → MQ(eval.subtask.complete) 
            → ProgressConsumer 消费
                → 写 DB（累加进度）
                → 推 WebSocket

@RabbitListener(queues = "eval.subtask.complete")
public void onComplete(SubtaskCompleteMessage msg) {
    // 1) 写库
    batchService.incDone(msg.getBatchId());
    
    // 2) 批量聚合推送（避免推送风暴）
    progressAggregator.maybePush(msg.getBatchId());
}

progressAggregator 攒 500ms 或 50 条再推一次。25000 个子任务如果完成一个推一次，消息风暴直接打挂推送通道，聚合后推送频率降了两个量级，前端进度条依旧平滑。

七、心跳与重连

WebSocket 连接很容易被 NAT、代理超时静默断开，两端都不知情。解决办法是心跳：

// 服务端每 30s 发 ping
@Override
protected void handlePing(WebSocketSession session, PingMessage message) {
    session.sendMessage(new PongMessage());
}

// 前端 30s 没收到消息主动 reconnect
socket.onclose = () => {
  setTimeout(() => connect(), 3000);
};

服务端发 ping、客户端回 pong，双向心跳能及时发现死连接。前端在 onclose 回调里 3 秒后自动重连，用户基本无感知。

八、断线重连的快照补偿

用户断网 5 分钟重连，中间的进度全丢了。重连后服务端立刻推一份当前快照，前端用快照覆盖本地状态：

// 重连时立即推一次当前快照
private void sendSnapshot(long batchId) {
    BatchProgress p = batchService.getProgress(batchId);
    sendToBatch(batchId, JsonUtil.toJson(
        new SnapshotEvent("SNAPSHOT", batchId, p)));
}

前端在收到 SNAPSHOT 后用最新数据覆盖本地状态，避免重复累加。这里有个约定：SNAPSHOT 是”全量覆盖”，PROGRESS 是”增量推进”，两种消息语义不同，前端处理逻辑要分开。

九、踩过的坑

实时链路的问题往往在压测和线上才暴露，把踩过的坑都记下来：

WebSocket 路径冲突：/ws/batch/{id} 中的 {id} 必须是数字，否则 Spring 路由匹配不到。用正则 @MessageMapping("/batch/{id:\\d+}")。
大消息体：单条消息 < 1MB，单帧不超过 16KB（WebSocket 协议）。超长用 chunked transfer。
多实例部署：用户连到实例 A，MQ 消费者在实例 B，跨实例推送要走 Redis Pub/Sub。
鉴权信息泄漏：握手时把 userId 存进 session attribute，afterConnectionEstablished 里再校验一次（防止有人篡改握手数据）。
连接数无限制：Nginx 默认 worker_connections 1024，要调到 65535，并加 IP 维度的连接数限制。
同 batch 多次打开：用户多 tab 打开同一 batch，progressAggregator 要按 batch 维度聚合推送，不能重复推。

十、监控与告警

实时通道最怕”悄悄挂掉”。我们给 WebSocket 加了关键指标，推送链路异常时能第一时间发现：

metrics.gauge("ws.connections.active", sessions.size());
metrics.counter("ws.message.sent", "type", "PROGRESS").increment();
metrics.counter("ws.message.drop", "reason", "session_closed").increment();

监控项：在线连接数、推送频率、断线率、消息大小 P99。上线第一天就靠连接数指标发现了一个没释放连接的 bug，session 表只增不减，这个指标救了一次线上事故。

到这里，WebSocket 实时进度推送从选型、握手、消息格式、连接管理到心跳重连、监控告警的完整链路就齐了。

常见问题（FAQ）

Q1：为什么不用 SSE？

SSE 只能服务端推，客户端无法发”取消/暂停”指令。WebSocket 全双工更灵活。

Q2：怎么解决跨实例推送？

Redis Pub/Sub。生产者 publish 到 ws:batch:{id} 频道，所有实例订阅，匹配自己的 session 就推。

Q3：消息顺序怎么保证？

WebSocket 单连接内消息天然有序。批量场景下用 batchId 维度串行化处理即可。