当前位置：首页 > article >正文

WebSocket 实现消息推送

article 2024/9/21 1:20:41

使用 WebSocket 实现消息推送是一个非常流行的方案，它允许服务器和客户端之间建立持久的双向通信通道。以下是如何在一个典型的 Java Spring Boot 项目中使用 WebSocket 实现消息推送的详细步骤。

1. 添加依赖

在 pom.xml 文件中添加 WebSocket 和 Spring Boot 相关的依赖：

<dependencies>
    <!-- Spring Boot Starter WebSocket -->
    <dependency>
        <groupId>org.springframework.boot</groupId>
        <artifactId>spring-boot-starter-websocket</artifactId>
    </dependency>
</dependencies>

2. 配置 WebSocket

创建一个 WebSocket 配置类，配置 WebSocket 的端点和处理器：

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 WebSocketConfig implements WebSocketConfigurer {

    @Override
    public void registerWebSocketHandlers(WebSocketHandlerRegistry registry) {
        registry.addHandler(new MyWebSocketHandler(), "/ws/my-endpoint")
                .setAllowedOrigins("*"); // 配置允许跨域访问
    }
}

3. 实现 WebSocket 处理器

创建一个 WebSocket 处理器来处理消息的发送和接收：

import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketMessage;
import java.io.IOException;

public class MyWebSocketHandler extends TextWebSocketHandler {

    @Override
    public void handleTextMessage(WebSocketSession session, TextMessage message) throws IOException {
        String payload = message.getPayload();
        // 处理接收到的消息（例如：记录日志、处理业务逻辑等）
        System.out.println("Received message: " + payload);

        // 发送消息到客户端
        session.sendMessage(new TextMessage("Message received: " + payload));
    }

    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        System.out.println("WebSocket connection established: " + session.getId());
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
        System.out.println("WebSocket connection closed: " + session.getId());
    }
}

4. 创建 WebSocket 客户端

在前端页面中使用 JavaScript 创建 WebSocket 客户端，并与服务器进行通信。以下是一个简单的 HTML 和 JavaScript 示例：

<!DOCTYPE html>
<html>
<head>
    <title>WebSocket Example</title>
</head>
<body>
    <h1>WebSocket Example</h1>
    <textarea id="messages" cols="100" rows="10" readonly></textarea><br>
    <input type="text" id="messageInput" placeholder="Enter your message" />
    <button onclick="sendMessage()">Send</button>

    <script>
        var socket = new WebSocket('ws://localhost:8080/ws/my-endpoint');

        socket.onopen = function(event) {
            console.log('WebSocket connection established');
        };

        socket.onmessage = function(event) {
            var messages = document.getElementById('messages');
            messages.value += 'Server: ' + event.data + '\n';
        };

        socket.onclose = function(event) {
            console.log('WebSocket connection closed');
        };

        function sendMessage() {
            var messageInput = document.getElementById('messageInput');
            var message = messageInput.value;
            socket.send(message);
            messageInput.value = '';
        }
    </script>
</body>
</html>

5. 推送消息到客户端

在服务器端推送消息到客户端的一个例子是使用 @Scheduled 注解定期发送消息：

import org.springframework.scheduling.annotation.Scheduled;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import java.io.IOException;
import java.util.Collection;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

@Component
public class WebSocketPushService {

    private final Map<String, WebSocketSession> sessions = new ConcurrentHashMap<>();

    public void addSession(WebSocketSession session) {
        sessions.put(session.getId(), session);
    }

    public void removeSession(WebSocketSession session) {
        sessions.remove(session.getId());
    }

    @Scheduled(fixedRate = 5000) // 每5秒推送一次
    public void pushMessage() {
        String message = "Current Time: " + System.currentTimeMillis();
        Collection<WebSocketSession> sessionCollection = sessions.values();
        for (WebSocketSession session : sessionCollection) {
            try {
                session.sendMessage(new TextMessage(message));
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
    }
}

6. 更新 WebSocket 处理器

将 WebSocket 处理器更新为能够处理连接和断开连接事件，并将客户端会话保存到 WebSocketPushService：

import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;
import org.springframework.web.socket.TextMessage;

public class MyWebSocketHandler extends TextWebSocketHandler {

    @Autowired
    private WebSocketPushService pushService;

    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        System.out.println("WebSocket connection established: " + session.getId());
        pushService.addSession(session);
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
        System.out.println("WebSocket connection closed: " + session.getId());
        pushService.removeSession(session);
    }

    @Override
    public void handleTextMessage(WebSocketSession session, TextMessage message) throws IOException {
        String payload = message.getPayload();
        System.out.println("Received message: " + payload);
        session.sendMessage(new TextMessage("Message received: " + payload));
    }
}