优化 Spring WebSocket 大型 Switch 语句：命令/策略模式

2025-01-01 16:50:05

优化 Java Spring WebSocket 处理中的大型 Switch 语句

在处理 WebSocket 连接时，经常需要根据客户端发送的不同消息执行相应的操作。一种常见做法是使用 switch 语句来区分不同的消息类型。但随着业务逻辑的增长，这种方法容易导致代码变得难以维护和扩展，特别是在一个大型的 switch 语句中，所有逻辑都耦合在一起。本文将介绍一些方法来解决这个问题，并通过代码示例进行演示。

问题所在

大型的 switch 语句会引发以下问题：

代码可读性差： 嵌套的 case 语句使代码难以理解。
维护性降低： 修改和添加新的消息类型需要更改整个 switch 语句。
可扩展性不足： 新增操作将增加代码复杂度，容易出错。
代码耦合度高： 所有逻辑都集中在一个地方，难以复用。

解决方案

针对以上问题，这里提供一些改进的方案，可以按需采用：

1. 使用命令模式

命令模式是一种行为设计模式，将请求封装成对象，从而使不同的请求具有参数化、队列化、可撤销的能力。此模式的核心在于将请求的发送者和接收者解耦，提升灵活性。

原理：

定义一个命令接口，包含 execute() 方法。
每个 case 对应一个具体的命令实现，处理各自的消息。
使用 Map<String, Command> 来存储消息类型和命令实现之间的映射关系。
在 WebSocket handler 中根据 action 查找对应的命令，执行。

代码示例：

首先定义 Command 接口和 Message 基类。

interface Command {
    void execute(WebSocketSession session, Message<?> message) throws IOException;
}
@Data
@AllArgsConstructor
@NoArgsConstructor
class Message<T> {
   private String action;
   private T data;

    public Message(String action){
        this.action=action;
    }
}

然后创建具体的命令，如CreateFileCommand。

@Component
@RequiredArgsConstructor
public class CreateFileCommand implements Command {

    private final FileService fileService;
    private final MessageParser messageParser;
    private final ProjectService projectService;

    @Override
    public void execute(WebSocketSession session, Message<?> message) throws IOException {
          Integer userId = (Integer) session.getAttributes().get("userId");
          Integer projectId = (Integer) session.getAttributes().get("projectId");
          Message<CreateFileDTO> createFileDTOMessage = messageParser.parseMessage(message.toString(), CreateFileDTO.class);
          fileService.createFileByUser(createFileDTOMessage.getData(), userId);
         new WebSocketUtils().broadcastMessageToProjectMembers(projectSessions,projectId, new Message("get-project-structure", projectService.getFolderTree(projectId)),messageParser);
    }

}

WebSocketUtils


class WebSocketUtils{
    private final ConcurrentHashMap<Integer, ConcurrentLinkedQueue<WebSocketSession>> projectSessions = new ConcurrentHashMap<>();

    public void broadcastMessageToProjectMembers( ConcurrentHashMap<Integer, ConcurrentLinkedQueue<WebSocketSession>> projectSessions,Integer projectId, Message messageContent,MessageParser messageParser) {
        ConcurrentLinkedQueue<WebSocketSession> sessions = projectSessions.get(projectId);
         broadcastMessageToMembers(sessions, messageContent, messageParser);
    }

    public void broadcastMessageToMembers(ConcurrentLinkedQueue<WebSocketSession> sessions, Message messageContent,MessageParser messageParser) {
        if (sessions != null) {
            try {
                TextMessage message = new TextMessage(messageParser.parseMessage(messageContent));
                for (WebSocketSession session : sessions) {
                    if (session.isOpen()) {
                        session.sendMessage(message);
                    }
                }
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
    }

}

最后更新 WebSocketHandler


@RequiredArgsConstructor
@Component
public class EditorOperationsWebSocketHandler extends TextWebSocketHandler {

    private final MessageParser messageParser;

    private final ApplicationContext applicationContext;

    private final ConcurrentHashMap<Integer, ConcurrentLinkedQueue<WebSocketSession>> projectSessions = new ConcurrentHashMap<>();

    private final ConcurrentHashMap<Integer, ConcurrentLinkedQueue<WebSocketSession>> discussionSessions = new ConcurrentHashMap<>();
      private  final Map<String, Command> commandMap = new HashMap<>();

      @PostConstruct
    public void init(){
           commandMap.put("create-file",applicationContext.getBean(CreateFileCommand.class));
        //注册其他的命令
       }

    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        String username = (String) session.getAttributes().get("sub");
        Integer projectId = (Integer) session.getAttributes().get("projectId");
        if (username != null && projectId != null) {
            projectSessions.computeIfAbsent(projectId, k -> new ConcurrentLinkedQueue<>()).add(session);
        } else {
            session.close(CloseStatus.NOT_ACCEPTABLE);
        }
    }

    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) {
        String payload = message.getPayload();
       handleMessage(session, payload);
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) {
       // cleanup(session);
    }

    public void handleMessage(WebSocketSession session, String jsonMessage) {
        try {
           Message<?> message = messageParser.parseMessage(jsonMessage);
            Command command = commandMap.get(message.getAction());
            if(command!=null)
            {
              command.execute(session,message);
            }

        } catch (IOException e) {
           e.printStackTrace();
        }
    }



}

操作步骤：

定义 Command 接口和 Message 基类。
为每个消息类型创建实现了 Command 接口的类。
将这些命令注入 commandMap。
在 handleMessage() 方法中，通过 action 找到对应的命令，执行。

2. 使用策略模式

策略模式定义了一组算法，并将每个算法封装起来，使它们可以互相替换。策略模式让算法独立于使用它的客户端而变化。

原理:

定义一个策略接口，包含 handle() 方法。
每个 case 对应一个策略实现类。
使用 Map<String, HandlerStrategy> 来管理 action 到对应策略的映射。
在 handleMessage() 方法中，获取相应的策略并执行。

代码示例:

类似于命令模式，我们定义一个 HandlerStrategy 接口：

interface HandlerStrategy {
  void handle(WebSocketSession session, Message<?> message) throws IOException;
}

然后创建具体策略如FileHandlerStrategy:


@Component
@RequiredArgsConstructor
public class FileHandlerStrategy implements HandlerStrategy{

   private final  FileService fileService;
    private final  MessageParser messageParser;
    private final  ProjectService projectService;

   @Override
    public void handle(WebSocketSession session, Message<?> message) throws IOException {
      String action = message.getAction();
      Integer userId = (Integer) session.getAttributes().get("userId");
        Integer projectId = (Integer) session.getAttributes().get("projectId");

        switch (action){
          case "create-file":
              Message<CreateFileDTO> createFileDTOMessage = messageParser.parseMessage(message.toString(), CreateFileDTO.class);
              var file = fileService.createFileByUser(createFileDTOMessage.getData(), userId);
             new WebSocketUtils().broadcastMessageToProjectMembers(projectSessions,projectId, new Message("get-project-structure", projectService.getFolderTree(projectId)),messageParser);
            break;
          //其他的case
        }
    }
}

更新WebSocketHandler:


@RequiredArgsConstructor
@Component
public class EditorOperationsWebSocketHandler extends TextWebSocketHandler {

    private final MessageParser messageParser;

    private final ApplicationContext applicationContext;

    private final ConcurrentHashMap<Integer, ConcurrentLinkedQueue<WebSocketSession>> projectSessions = new ConcurrentHashMap<>();
    private final ConcurrentHashMap<Integer, ConcurrentLinkedQueue<WebSocketSession>> discussionSessions = new ConcurrentHashMap<>();

    private final Map<String,HandlerStrategy> strategyMap = new HashMap<>();


      @PostConstruct
    public void init(){
         strategyMap.put("file",applicationContext.getBean(FileHandlerStrategy.class));
          //其他的策略
       }


    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        String username = (String) session.getAttributes().get("sub");
        Integer projectId = (Integer) session.getAttributes().get("projectId");
        if (username != null && projectId != null) {
            projectSessions.computeIfAbsent(projectId, k -> new ConcurrentLinkedQueue<>()).add(session);

        } else {
            session.close(CloseStatus.NOT_ACCEPTABLE);
        }
    }

    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) {
        String payload = message.getPayload();

        handleMessage(session, payload);
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) {
       // cleanup(session);
    }

    public void handleMessage(WebSocketSession session, String jsonMessage) {
        try {
           Message<?> message = messageParser.parseMessage(jsonMessage);
           String action = message.getAction();
           String type =  action.split("-")[0];

            HandlerStrategy handlerStrategy = strategyMap.get(type);
            if(handlerStrategy!=null){
              handlerStrategy.handle(session,message);
            }

        } catch (IOException e) {
          e.printStackTrace();
        }
    }


}