Spring实现自定义WebSocket

Spring实现自定义WebSocket

引入gradle

// https://mvnrepository.com/artifact/org.springframework.boot/spring-boot-starter-websocket
implementation group: 'org.springframework.boot', name: 'spring-boot-starter-websocket', version: '2.3.12.RELEASE'

配置类

import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.web.socket.server.standard.ServerEndpointExporter;

@Configuration
public class WebSocketConfig {

    @Bean
    public ServerEndpointExporter serverEndpointExporter() {
        return new ServerEndpointExporter();
    }
}

WebSocket实现类

import org.springframework.stereotype.Component;

import javax.websocket.*;
import javax.websocket.server.ServerEndpoint;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

@Component
@ServerEndpoint(value = "/onLine")
public class OnLineWebSocket {

    /**
     * 存放所有在线的客户端
     */
    private static final Map<String, Session> CLIENTS = new ConcurrentHashMap<>();

    /**
     * 存放所有在线账号信息与所对应的sessionId
     */
    private static final Map<String, String> SESSIONIDS = new ConcurrentHashMap<>();


    /**
     * webSocket连接
     * 连接建立成功调用的方法
     *
     * @param session webSocket的连接Session
     */
    @OnOpen
    public void onOpen(Session session) {
        String id = session.getId();
        CLIENTS.put(id, session);
    }

    /**
     * 连接关闭
     *
     * @param session seesion客户端
     */
    @OnClose
    public void onClose(Session session) {
        String sessionId = session.getId();
        // 连接关闭，移除session
        CLIENTS.remove(sessionId);
        // 移除SESSIONIDS中记录数据
        SESSIONIDS.remove(sessionId);
    }

    /**
     * 登录成功接收的消息
     *
     * @param message 前端发送的信息
     * @param session 前端的session
     */
    @OnMessage
    public void onMessage(String message, Session session) {
        // 接收到的信息返回客户端
        if (!sendMessage(message, session)) {
            System.out.println("信息发送失败");
        }
    }

    /**
     * WebSocket 连接异常
     *
     * @param session
     * @param error
     */
    @OnError
    public void onError(Session session, Throwable error) {
        // 响应前端异常
        String message = error.getMessage();
        if (!sendMessage(message, session)) {
            System.out.println("信息发送失败");
        }
    }

    /**
     * 响应前端信息
     * 外部调用，实现推送websocket消息
     *
     * @param message   响应信息
     * @param toSession 传输Session
     */
    public boolean sendMessage(String message, Session toSession) {
        try {
            toSession.getBasicRemote().sendText(message);
            return true;
        } catch (Exception e) {
            e.printStackTrace();
            return false;
        }
    }

    /**
     * 登录成功信息存入
     * 外部调用，绑定当前登录账号到Session
     * 当websocket推送到特定用户时，通过account可查用户，再关联到Session，即可实现推送websocket消息到特定用户
     *
     * @param account   登录账号
     * @param sessionId sessionId
     */
    public synchronized void addSessionIds(String account, String sessionId) {
        SESSIONIDS.put(sessionId, account);
    }

}