Java - WebSocket配置及使用

引入依赖

Spring Boot 默认支持 WebSocket，但需要引入 spring-boot-starter-websocket 依赖，然后重新构建项目

<dependency>
    <groupId>org.springframework.bootgroupId>
    <artifactId>spring-boot-starter-websocketartifactId>
dependency>

搭建 WebSocket 服务

创建 WebSocketServer.java

package com.project.module.webSocket;

import com.fasterxml.jackson.databind.ObjectMapper;
import jakarta.websocket.*;
import jakarta.websocket.server.PathParam;
import jakarta.websocket.server.ServerEndpoint;
import org.springframework.stereotype.Service;

import java.io.IOException;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

@Service
@ServerEndpoint("/ws/{userId}") // WebSocket 连接端点
public class WebSocketServer {

    // 存储所有用户的 WebSocket 连接
    private static final Map<String, WebSocketServer> clients = new ConcurrentHashMap<>();
    private Session session;
    private String userId; // 当前连接的用户ID

    @OnOpen
    public void onOpen(Session session, @PathParam("userId") String userId) {
        this.session = session;
        this.userId = userId;
        clients.put(userId, this);
        System.out.println("新客户端连接，当前在线用户：" + clients);
    }

    @OnMessage
    public void onMessage(String message) {
        System.out.println("收到消息：" + message);

        // 假设消息格式是 "userId:message"
        String[] parts = message.split(":", 2);
        if (parts.length == 2) {
            String targetUserId = parts[0];
            String msg = parts[1];

            sendToUser(targetUserId, "私信：" + msg);
        } else {
            System.out.println("无效的消息格式：" + message);
        }
    }

//    @OnMessage
//    public void onMessage(String message, @PathParam("userId") String userId) {
//        System.out.println("收到消息：" + message);
//        sendToUser(userId, "私信：" + message);
//    }

    // 给指定用户发送消息
    public void sendToUser(String userId, Object messageObject) {
        WebSocketServer client = clients.get(userId);
        System.out.println("当前用户：" + userId);
        if (client != null) {
            try {
                // 使用 Jackson 转换对象为 JSON
                ObjectMapper objectMapper = new ObjectMapper();
                String jsonMessage = objectMapper.writeValueAsString(messageObject);
                client.session.getBasicRemote().sendText(jsonMessage);
            } catch (IOException e) {
                e.printStackTrace();
            }
        } else {
            System.out.println("用户 " + userId + " 不在线，消息未发送");
        }
    }

    @OnClose
    public void onClose() {
        for (Map.Entry<String, WebSocketServer> entry : clients.entrySet()) {
            if (entry.getValue().session.getId().equals(session.getId())) {
                clients.remove(entry.getKey());
                System.out.println("客户端断开连接，当前在线人数：" + clients.size());
                break;
            }
        }
    }

    @OnError
    public void onError(Session session, Throwable error) {
        error.printStackTrace();
    }

    public void sendToAll(Object messageObject) {
        for (Map.Entry<String, WebSocketServer> entry : clients.entrySet()) {
            WebSocketServer client = entry.getValue();
            try {
                // 使用 Jackson 转换对象为 JSON
                ObjectMapper objectMapper = new ObjectMapper();
                String jsonMessage = objectMapper.writeValueAsString(messageObject);
                client.session.getBasicRemote().sendText(jsonMessage);
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
    }

    public Map<String, WebSocketServer> getClients() {
        return clients;
    }
}

手动注册 WebSocket

@ServerEndpoint 不能直接用 @Autowired 注入 Spring Bean，所以需要 手动注册
在 WebSocketConfig 里添加 Bean
创建 WebSocketConfig.java：

package com.project.module.webSocket;

import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import javax.websocket.server.ServerEndpointConfig;
import javax.websocket.server.ServerEndpointExporter;

@Configuration
public class WebSocketConfig {
    
    // 启动 WebSocket 服务器
    @Bean
    public ServerEndpointExporter serverEndpointExporter() {
        return new ServerEndpointExporter();
    }
}

发送Websocket服务

package com.project.module.webSocket;

import org.springframework.scheduling.annotation.Scheduled;
import org.springframework.stereotype.Component;

@Component
public class WebSocketTask {

    private final WebSocketServer webSocketServer;

    public WebSocketTask(WebSocketServer webSocketServer) {
        this.webSocketServer = webSocketServer;
    }

    // 每一分钟推送一次消息
    @Scheduled(fixedRate = 1000*60)
    public void pushMessage() {
        webSocketServer.sendToUser("userId","定时任务推送消息：" + System.currentTimeMillis());
    }
}

前端链接 WebSocket

创建 webSocket.js
注：此处的 import {host} from './base.js'为引入的全局服务地址，例如 host 为 localhost ，开发者可自行更换为自己的 WebSocket 地址

import {host} from './base.js'
const WebSocketService = {
    ws: null,
    connect(userId, callback) {
        this.ws = new WebSocket(`ws://${host}:8980/ws/` + userId);

        this.ws.onopen = () => {
            console.log("WebSocket 连接成功");
        };

        this.ws.onmessage = (event) => {
            callback(JSON.parse(event.data));
        };

        this.ws.onclose = () => {
            console.log("WebSocket 连接已关闭");
        };
    },
    sendMessage(message) {
        if (this.ws && this.ws.readyState === WebSocket.OPEN) {
            this.ws.send(message);
        }
    },
    disconnect() {
        if (this.ws) {
            this.ws.close();
        }
    },
};

export default WebSocketService;

在Vue组将中使用

<template>
  <div class="container">
  div>
template>

<script setup>
import { onMounted, ref } from "vue";
import WebSocketService from "@/utils/websocket";

const CURRENT_CHART_NAME = 'safety'
const safetyRef = ref()
let safetyChart = null

const { setChart } = useSetChart()

onMounted(() => {
  websocketConn()
})

function websocketConn() {
  const userId = `userId`
  WebSocketService.connect(userId, (msg) => {
    console.log(msg)
  });
}
script>

注：此处的userId对应后端 webSocketServer.sendToUser("userId ","定时任务推送消息：" + System.currentTimeMillis()); 中的 userId,代表前端订阅WebSocket消息的用户以及WebSocket要发送的用户