WebSocket推送端和接收端实现

package org.mf.socket; import com.alibaba.fastjson.JSONObject; import lombok.AllArgsConstructor; import lombok.Data; import lombok.Getter; import lombok.NoArgsConstructor; import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; import org.java_websocket.WebSocket; import org.java_websocket.framing.CloseFrame; import org.java_websocket.handshake.ClientHandshake; import org.java_websocket.server.WebSocketServer; import java.net.InetSocketAddress; import java.util.Map; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.SynchronousQueue; /** * 类名/接口名 * * @author mengfei * @date 2023年09月11日 13:36 */ public class WebSocketUtil { private static final Logger LOGGER = LogManager.getLogger(WebSocketUtil.class); /** * 自定义server容器，用来记录连接的webSocket */ private static CustomWebServer CUSTOM_WEB_SERVER = null; /** * 启动websocket服务 * @param port 端口号 * @throws Exception */ public static void startServer(Integer port) { //设置jvm退出时的清理线程 setClearWebSocketTask(); //创建同步阻塞队列，用来获取webSocket启动结果 SynchronousQueue<WebSocketStartResult> synchronousQueue = new SynchronousQueue<>(); CUSTOM_WEB_SERVER = new CustomWebServer(port, synchronousQueue); CUSTOM_WEB_SERVER.start(); //线程一直阻塞，直到拿到启动结果 WebSocketStartResult startResult = null; try { startResult = synchronousQueue.take(); } catch (InterruptedException e) { throw new RuntimeException(e); } if (startResult.getStartStatus() == 2) { //启动失败，抛出异常 Exception startException = (Exception) startResult.getStartResult(); LOGGER.error("websocket server启动失败", startException); throw new RuntimeException(startException); } LOGGER.info("websocket server启动成功，端口号：{}", port); } private static void setClearWebSocketTask() { Thread thread = new Thread(new Runnable() { @Override public void run() { stopServer(); } }); thread.setName("websocket关闭线程"); thread.setDaemon(true); Runtime.getRuntime() .addShutdownHook(thread); } public static void stopServer() { if (CUSTOM_WEB_SERVER != null) { try { CUSTOM_WEB_SERVER.stop(); } catch (InterruptedException e) { LOGGER.error("websocket服务关闭失败", e); throw new RuntimeException(e); } } } /** * 向websocket连接中推送消息，方法加锁，防止多个线程同时推送消息造成错误 * @param messageBody * @return */ public static synchronized void writeMessage(String messageBody) { if (CUSTOM_WEB_SERVER == null) { return; } ConcurrentHashMap<WebSocket, Integer> connAndStatus = CUSTOM_WEB_SERVER.getConns(); for (Map.Entry<WebSocket, Integer> webSocketAnsStatusEntry : connAndStatus.entrySet()) { WebSocket conn = webSocketAnsStatusEntry.getKey(); try { //发送消息 JSONObject jsonObject = new JSONObject(); jsonObject.put("messageBody", messageBody); conn.send(jsonObject.toJSONString()); } catch (Exception e) { LOGGER.error("websocket消息推送失败", e); } } }; /** * websocket启动结果 */ @NoArgsConstructor @AllArgsConstructor @Data private static class WebSocketStartResult { /** * socket server启动状态 */ private Integer startStatus; /** * socket server启动结果信息 * 如果status为1，result中存放启动成功字符串 * 如果status为2，result中存放启动的异常 */ private Object startResult; } /** * 自定义websocket服务 */ private static class CustomWebServer extends WebSocketServer { /** * 用于通知websocket server启动结果 */ private SynchronousQueue<WebSocketStartResult> synchronousQueue; /** * 记录webSocket的启动状态，0-已创建未启动，1-启动成功，2-启动失败 */ private Integer startState = 0; /** * 存储所有socket连接(使用线程安全的map集合，防止在往连接中推送消息时发生不必要的错误，比如连接已经关闭但还是向其推送消息) */ @Getter private ConcurrentHashMap<WebSocket, Integer> conns = new ConcurrentHashMap<>(); public CustomWebServer(int port, SynchronousQueue<WebSocketStartResult> synchronousQueue) { super(new InetSocketAddress(port)); this.synchronousQueue = synchronousQueue; } /** * websocket conn打开的回调 * @param conn The <tt>WebSocket</tt> instance this event is occurring on. * @param handshake The handshake of the websocket instance */ @Override public void onOpen(WebSocket conn, ClientHandshake handshake) { LOGGER.info("websocket 连接创建，连接：{}", conn); //将连接加入map conns.put(conn, 0); } /** * websocket conn收到消息的回调，目前仅处理认证消息，其他消息返回错误 * @param conn The <tt>WebSocket</tt> instance this event is occurring on. * @param message The UTF-8 decoded message that was received. */ @Override public void onMessage(WebSocket conn, String message) { LOGGER.info("websocket 接收到消息，连接：{}，消息：{}", conn, message); //在此处处理接收的message } /** * websocket conn被关闭的回调 * @param conn The <tt>WebSocket</tt> instance this event is occurring on. * @param code The codes can be looked up here: {@link CloseFrame} * @param reason Additional information string * @param remote Returns whether or not the closing of the connection was initiated by the remote * host. */ @Override public void onClose(WebSocket conn, int code, String reason, boolean remote) { LOGGER.info("websocket 连接关闭，连接：{}，状态：{}，关闭原因：{}，是否是远端关闭：{}", conn, code, reason, remote); //从连接map中移除 conns.remove(conn); } /** * websocketServer发生异常的回调 * @param conn Can be null if there error does not belong to one specific websocket. For example * if the servers port could not be bound. * @param ex The exception causing this error */ @Override public void onError(WebSocket conn, Exception ex) { LOGGER.error("websocket 异常，连接：{}，exception：{}", conn, ex); // 在这里，您可以处理任何错误情况 if (startState == 0) { //修改启动状态 startState = 2; //将启动结果的异常放入请求队列中 WebSocketStartResult webSocketStartResult = new WebSocketStartResult(2, ex); try { synchronousQueue.put(webSocketStartResult); } catch (Interrupt