tomcat websocket 并发问题解决（二）

现在问题回到最开始了，上一篇文章说过，因为 tomcat 的 session.sendMessage() 方法在并发环境下会抛出异常，我们为了保证程序的正确性，在 session 上加了同步限制。然而这种做法会因对 session 的竞争导致所有的消息事实上是同步进行的，系统中的全部消息形成了一个超长队列，造成了性能上的问题，下边我就讲讲我的解决思路和方法。

解决思路

首先，我想到的有两点：

1. session.sendMessage() 方法仅提供了最小的线程安全保证，同时发送多条消息时虽然不会给客户端发送错误信息，但是会抛出异常，这是不符合我们预期的。
2. 一条消息发给多个用户即多个 session 时，不同 session 之间不存在竞争，不需要同步进行。

那么对于 1 ，可以将对同一个 session 发送的消息丢到一个队列里，同步地进行发送，这样可以避免对 session 的争用，也可以避免 session.sendMessage() 的并发执行问题。

对于 2 ，可以使用多个队列来存储对不同 session 的消息，启用多个线程来消费队列中的信息，提高程序性能。

这里使用的这个队列将消息先缓存起来的方法，就是我理解的 issue 中提到的 tricky buffering，下面来讲讲我的具体实现。

具体实现

写到这里，突然有些兴趣索然，这就是一个常见的生产者-消费者模型，感觉上边两段已经说得很明白了，不过既然坑已经挖了，还是好好填上吧。

消息

发送消息时，消费者最少需要知道两个要素：发送什么和发送给谁，我们就把这两个要素简单封装下，代码如下：

@Slf4j
@AllArgsConstructor
public class TextMessageJob {

    private WebSocketSession session;

    private String messageBody;

    public void send() throws IOException {
        if(session.isOpen()) {
            session.sendMessage(new TextMessage(messageBody));
            log.debug("Message {} send success", messageBody);
        } else {
            log.warn("Message {}, session : {} sent failed cause the session was closed", messageBody, session.getId());
        }
    }
    
}

生产者和消费者

接口都非常简单，在这里略过，我们直接就上实现的代码了。

@Slf4j
public abstract class AbstractJobHandler implements Runnable, MessageHandleErrorCallback, JobHandler {

    private BlockingQueue<MessageJob> jobContainer;

    public AbstractJobHandler() {
        this.jobContainer = new LinkedTransferQueue<>();
    }

    @Override
    public void run() {
        log.info("{} init", Thread.currentThread().getName());
        while (true) {
            MessageJob messageJob = jobContainer.take();
            internalSendMessage(messageJob);
        }
    }

    @Override
    public void pushMessage(MessageJob messageJob) {
        try {
            this.jobContainer.put(messageJob);
        } catch (InterruptedException e) {
            log.warn("Message send failed , messageJob : {}, exception : {}", messageJob, e.getMessage());
            onMessageHandleError(messageJob);
        }
    }

    public void internalSendMessage(MessageJob messageJob){
        try {
            try {
                log.info("Message {} pop from queue", messageJob);
                messageJob.send();
            } catch (IOException e) {
                log.warn("Message send failed , messageJob : {}, exception : {}", messageJob, e.getMessage());
                onMessageHandleError(messageJob);
            }
        } catch (Throwable e) {
            e.printStackTrace();
        }
    }
    
}

因为对消息量不是太有把握，所以这里直接使用了一个无界阻塞队列来存放消息，有新消息时丢到队列里。

JobHandler 自己跑起来的时候，会不停地从这个队列里 take 新的消息，拿去处理。

消息分配

开始我们说到，我们会使用多个队列，开启多条线程来处理消息，同一个 session 的消息会丢到同一个队列里来避免对 session 的争用，这些工作是在 JobCenter 中完成的，具体代码如下：

@Slf4j
public class JobCenter {

    private SimpleJobHandler[] jobHandlers;

    private int capacity;

    private volatile int status;

    private static final int STATUS_NEW = 0;

    private static final int STATUS_ON_EXECUTE = 1;

    public JobCenter(int capacity) {
        this.capacity = capacity;
        status = STATUS_NEW;
        jobHandlers = new SimpleJobHandler[capacity];
    }

    private void initJobHandler(int index) {
        jobHandlers[index] = new SimpleJobHandler(index, "SimpleJobHandler-" + index);
        Thread thread = new Thread(jobHandlers[index]);
        thread.setName(jobHandlers[index].getName());
        thread.start();
    }

    public synchronized void start() {
        if (status == STATUS_NEW) {
            IntStream.range(0, capacity).forEach(this::initJobHandler);
            log.info("job center started");
            status = STATUS_ON_EXECUTE;
        } else {
            log.error("JobCenter already on running!");
        }
    }

    public void pushMessage(MessageJob message) {
        int index = (Math.abs(message.getSession().getId().hashCode()) % capacity);
        jobHandlers[index].pushMessage(message);
    }

}
    private static final int DEFAULT_WORK_COUNT = 10;

...
        JobCenter jobCenter = new JobCenter(DEFAULT_WORK_COUNT);
        jobCenter.start();
...

这里的代码就非常简单了，SimpleJobHandler 的代码我没有贴出来，它是上边 AbstractJobHandler 的一个实现类，主要就是添加了 name 和 index 这两个方便 debug 的信息，然后实现了 MessageHandleErrorCallback 接口，在发送失败的时候进行善后处理。

主要需要关注的一个属性是 capacity ，它定义了我们工作线程的个数。

比如我们将其初始为10 ，就有10个消费者，一共有10条队列，每次有新消息的时候，我们使用 sessionId 对 capacity 取模，丢到对应的队列中，这样就保证了多条消息并行处理，而对同一个 session 的消息因为是同一条 jobHandler 线程处理，所以是同步执行的。整个代码架构大概如下图：

输入图片说明

后记

代码写的不怎么样，当时写的比较仓促，个人能力也还有不足。

在实测中性能较原来有所提升，不过还未达令人满意的程度，这牵涉到另外一些古老的坑，业务代码承载了太多事情，还有优化空间，能想到的有以下几点：

1. 使用了 websocket，然而身份校验，权限验证还使用了无状态的思想，每次收到客户端消息都要重新执行一遍。既然是长连接，这个事情放在 afterHandshake 回调中做一次就好了。
2. 消息的持久化等一些较耗时的IO操作可以做异步执行。
3. 敏感词过滤的实现仍有优化空间。
4. 其他一些不方便透漏的业务逻辑，这个让我觉得团队内技术方案评审的必要性。

总而言之，进一步优化现在要做的就是拆代码，一下想到了“只送大脑”这个梗😂

这个坑本来打算春节的时候填的，然后还是算立了 flag，一直拖到春节后才完成。不过，春节时我也不是什么都没干：我看了部分 Spring 的代码，学习了下他们是怎么解决这个问题的。再立下一个 flag，下篇文章聊聊 Spring-websocket 的 tricky buffering 实现。