springboot整合SSE(Server Send Event)推送流信息

Spring boot整合SSE实现服务器实时推送流信息代码实现

本文只描述了sse的代码使用

服务端

1、控制层的链接接口必须返回SseEmitter，前端才能获取到该链接

2、业务层写业务逻辑即可，sse相关逻辑可以放在通用类中

工具类

1、sse发送消息是通过阻塞队列进行触发的

package com.zbnsec.simulator.framework.sse;

import com.zbnsec.simulator.project.topology.controller.dto.TopologyChangeDTO;
import jakarta.annotation.PostConstruct;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.util.Map;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * sse工具类
 *
 * @author user
 */
@Slf4j
@Component
public class SSEServer {

    /**
     * 当前连接数
     */
    private static AtomicInteger count = new AtomicInteger(0);

    private static Map<String, SseEmitter> cache = new ConcurrentHashMap<>();

    private static final BlockingQueue<TopologyChangeDTO> QUEUE = new LinkedBlockingQueue<>();

    /**
     * 创建sse链接
     *
     * @param clientId 客户端id
     * @return
     */
    public static SseEmitter connect(String clientId, TopologyChangeDTO topologyChangeDTO) {
        //设置超时时间，0表示不过期，默认是30秒，超过时间未完成会抛出异常
        SseEmitter sseEmitter = new SseEmitter(0L);
        pushMessage(sseEmitter, topologyChangeDTO);
        // 注册回调
        // 通信完成
        sseEmitter.onCompletion(() -> cache.remove(clientId));
        sseEmitter.onError(e -> log.error("SSEServer.connect error:", e));
        /*sseEmitter.onTimeout(timeOutCallBack(userId));*/
        cache.put(clientId, sseEmitter);
        //数量+1
        count.getAndIncrement();
        log.info("SSEServer create new sse connect ,current clientId:{}, 连接数={}", clientId, count);
        return sseEmitter;
    }

    /**
     * 关闭sse链接
     *
     * @param clientId 客户端id
     */
    public static void close(String clientId) {
        SseEmitter sseEmitter = cache.get(clientId);
        if (sseEmitter != null) {
            sseEmitter.complete();
            cache.remove(clientId);
        }
    }

    private static Thread SEND_THREAD;

    @PostConstruct
    public void pushMessage() {
        log.info("SSEServer start simulator SSE push thread.");
        SEND_THREAD = new Thread(new Runnable() {
            @Override
            public void run() {
                while (true) {
                    TopologyChangeDTO message = pull();
                    if (message != null && !message.getChange().equals(message.getLastChange())) {
                        for (Map.Entry<String, SseEmitter> entry : cache.entrySet()) {
                            pushMessage(entry.getValue(), message);
                        }
                    }
                }
            }
        }, "Simulator-SSE-Push");
        SEND_THREAD.setDaemon(true);
        SEND_THREAD.start();
        log.info("SSEServer start simulator SSE push thread completed.");
    }

    private static void pushMessage(SseEmitter sseEmitter, TopologyChangeDTO message) {
        try {
            sseEmitter.send(
                    SseEmitter
                            .event()
                            .data(message)
                            .reconnectTime(3000)
            );
        } catch (Exception e) {
            log.error("SSEServer pushMessage error:", e);
        }
    }


    public static void pushQueue(TopologyChangeDTO message) {
        if (null == message) {
            return;
        }
        try {
            QUEUE.put(message);
        } catch (InterruptedException exception) {
            log.error("SSEServer pushQueue：", exception);
        }
    }

    public static TopologyChangeDTO pull() {
        try {
            return QUEUE.take();
        } catch (InterruptedException exception) {
            log.error("SSEServer io thread interrupted", exception);
            return null;
        }
    }
}