simulation-backend/src/main/java/com/hivekion/ws/WsServer.java

package com.hivekion.ws;

import com.hivekion.Global;
import com.hivekion.common.entity.RequestCmdInfo;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import javax.websocket.OnClose;
import javax.websocket.OnError;
import javax.websocket.OnMessage;
import javax.websocket.OnOpen;
import javax.websocket.Session;
import javax.websocket.server.PathParam;
import javax.websocket.server.ServerEndpoint;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;

/**
 * [类的简要说明]
 * <p>
 * [详细描述，可选]
 * <p>
 *
 * @author LiDongYU
 * @since 2025/7/22
 */
@ServerEndpoint("/ws/{scenarioId}/{room}")
@Component
@Slf4j
public class WsServer {

  private static final ConcurrentHashMap<Integer, Object> lockMap = new ConcurrentHashMap<>();
  // key -> 当前key下所有会话
  private static final Map<String, Map<String, Map<String, Session>>> SESSION_MAP = new ConcurrentHashMap<>();

  private final Object lock = new Object();

  @OnOpen
  public void onOpen(Session session,
      @PathParam("room") String room, @PathParam("scenarioId") String scenarioId) {
    log.info("onOpen::scenarioId: {}, room: {}", room, room);

    synchronized (lock) {
      if (SESSION_MAP.get(scenarioId) == null) {
        SESSION_MAP.put(scenarioId, new ConcurrentHashMap<>());
        Map<String, Session> sessionMap = new ConcurrentHashMap<>();
        sessionMap.put(session.getId(), session);
        SESSION_MAP.get(scenarioId).put(room, sessionMap);
      } else {

        if (SESSION_MAP.get(scenarioId).get(room) == null) {

          Map<String, Session> sessionMap = new ConcurrentHashMap<>();
          sessionMap.put(session.getId(), session);
          SESSION_MAP.get(scenarioId).put(room, sessionMap);
        }


      }
    }
  }

  @OnClose
  public void onClose(Session session, @PathParam("room") String room,
      @PathParam("scenarioId") String scenarioId) {
    log.info("onClose::room: {}, scenarioId: {}", room, scenarioId);
    synchronized (lock) {
      //获取session信息
      if (SESSION_MAP.get(scenarioId) != null) {
        if (SESSION_MAP.get(scenarioId).get(room) != null) {
          SESSION_MAP.get(scenarioId).get(room).remove(session.getId());
        }
      }

    }

  }

  @OnMessage
  public void onMessage(String message, Session session, @PathParam("room") String room,
      @PathParam("scenarioId") String scenarioId) {
    try {
      log.info("onMessage::room: {}, message: {}", room, message);
      RequestCmdInfo requestCmdInfo = new RequestCmdInfo();
      requestCmdInfo.setScenarioId(Integer.parseInt(scenarioId));
      requestCmdInfo.setRoom(room);
      requestCmdInfo.setMessage(message);
      Global.receiveCmdInfoQueue.add(requestCmdInfo);
    } catch (Exception e) {
      log.error("onMessage::room: {}, message: {},error::", room, message, e);
    }


  }

  @OnError
  public void onError(Session session, Throwable error, @PathParam("room") String room,
      @PathParam("scenarioId") String scenarioId) {
    log.info("onError::room: {},sessionId::{} error: ", room, session.getId(), error);
    onClose(session, room, scenarioId);
  }

  public static void sendMessage(Integer scenarioId, String message) {

  }

  public static void sendMessage(Integer scenarioId, String room, String message) {
    Object lock = lockMap.computeIfAbsent(scenarioId, k -> new Object());
    synchronized (lock) {
      Map<String, Map<String, Session>> roomMap = SESSION_MAP.get(String.valueOf(scenarioId));
      if (roomMap.containsKey(room)) {
        Map<String, Session> singleRoomMap = roomMap.get(room);
        singleRoomMap.forEach((sessionId, session) -> {
          session.getAsyncRemote().sendText(message);
        });
      }
    }
  }

}