websocket消息丢失解决方案

后台在使用websocket给前端传消息时，有时消息量过大会有数据丢失的偶发情况，websocket源码中未查看到获取消息发送成功的状态，可以如下解决。

一、整体思路

1、后台通过websocket传输给前端消息，并且后台生成校验此消息的定时任务，设置每5秒重发
2、前端接收到消息后将消息通过websocket传输给后台
3、后台如接收到前端的消息则删除对应的发送消息定时任务，如未收到消息则继续发送，设置最多发送5次（超过5次默认认为此条消息记录有误）
4、建议：建议websocket发送消息单独为一个模块，防止定时任务过多抢占服务内存情况发生。

二、代码示例

1、线程池的配置

创建一个配置类，注入线程池的相关配置

@Configuration
public class WebConfig {
    @Bean("threadPoolTaskScheduler")
    public ThreadPoolTaskScheduler getThreadPoolTaskScheduler() {
        // 定时任务线程池
        ThreadPoolTaskScheduler executor = new ThreadPoolTaskScheduler();
        // 线程池大小
        executor.setPoolSize(10);
        // 线程执行前缀
        executor.setThreadNamePrefix("ThreadPoolTaskScheduler-");
        // executor.setWaitForTasksToCompleteOnShutdown(true);
        // executor.setAwaitTerminationSeconds(60);
        executor.initialize();
        return executor;
    }

2、消息实体类

@Data
public class TestEntity {
    private String key;//每条消息key要保持唯一
    private Object value;//发送的消息内容
    private Integer sendNum;//同一条消息发送次数
}

3、手动注入所需工具类

package com.media.common.utils;

import org.springframework.beans.BeansException;
import org.springframework.context.ApplicationContext;
import org.springframework.context.ApplicationContextAware;
import org.springframework.stereotype.Component;

@Component
public class SpringContextUtils implements ApplicationContextAware {
    /**
     * 应用上下文
     */
    private static ApplicationContext applicationContext;

    @Override
    public void setApplicationContext(ApplicationContext applicationContext) throws BeansException {
        SpringContextUtils.applicationContext = applicationContext;
    }

    public static ApplicationContext getApplicationContext(){
        return applicationContext;
    }

    @SuppressWarnings("unchecked")
    public static <T> T getBean(String name) throws BeansException {
        return (T)applicationContext.getBean(name);
    }

    public static <T> T getBean(Class<T> clz) throws BeansException {
        return (T)applicationContext.getBean(clz);
    }
}

4、定时任务编辑类

将所有的定时任务放入一个队列中，如想要停止此定时任务，直接将队列中对应的key删除即可（不同的消息要保持key唯一）

package com.media.msg.websockettest;

import com.alibaba.fastjson2.JSONObject;
import com.media.common.dto.common.ApiResult;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.scheduling.concurrent.ThreadPoolTaskScheduler;
import org.springframework.scheduling.support.CronTrigger;
import org.springframework.web.bind.annotation.*;

import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ScheduledFuture;

@RestController
public class TestController {
    private final Logger log = LoggerFactory.getLogger(this.getClass());

    private final String cron = "0/5 * * * * ?";//5秒重发消息

    // 线程池
    @Autowired
    private ThreadPoolTaskScheduler threadPoolTaskScheduler;

    // 任务队列管理
    @SuppressWarnings("rawtypes")
    private ConcurrentHashMap<String, ScheduledFuture> futureMap = new ConcurrentHashMap<String, ScheduledFuture>();

    // 加入新的任务进来
    @SuppressWarnings({"rawtypes"})
    @PostMapping("/addSchedule")
    public ApiResult addSchedule(@RequestBody TestEntity t) {
        DelayTaskExecTest task = new DelayTaskExecTest(t);
        ScheduledFuture<?> schedule = threadPoolTaskScheduler.schedule(task, new CronTrigger(cron));
        System.out.println("新消息已添加到定时任务：" + JSONObject.toJSONString(t));
        // 加入到队列中，
        futureMap.put(t.getKey(), schedule);

        return ApiResult.success();
    }

    // 移除已有的一个任务
    @SuppressWarnings("rawtypes")
    @PostMapping("/removeSchedule")
    public ApiResult removeSchedule(@RequestParam("key") String key) {
        ScheduledFuture scheduledFuture = futureMap.get(key);
        if (scheduledFuture != null) {
            // 取消定时任务
            scheduledFuture.cancel(true);
            // 如果任务取消需要消耗点时间
            boolean cancelled = scheduledFuture.isCancelled();
            while (!cancelled) {
                scheduledFuture.cancel(true);
                System.out.println(key + "取消中");
            }
            System.out.println(key + "任务移除成功");
            // 最后从队列中删除
            futureMap.remove(key);
        }
        return ApiResult.success();
    }
}

5、定时任务处理类

在run方法中写相关逻辑，我这里是调用了websocket的消息发送接口。
此外，要注意实现了Runnable的接口@Autowired注入会为null，所以需要手动注入

package com.media.msg.websockettest;

import com.media.common.utils.SpringContextUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

public class DelayTaskExecTest implements Runnable {
    private static final Logger log = LoggerFactory.getLogger(DelayTaskExecTest.class);

    //在Runnable @Autowired注入会null  所以需要手动注入
    private TestController testController;

    TestEntity testEntity;

    public DelayTaskExecTest(TestEntity testEntity) {
        this.testEntity = testEntity;
    }

    @Override
    public void run() {
        //执行具体的定时任务业务逻辑
        log.info("发送websocket消息，key={}，第{}次发送", testEntity.getKey(), testEntity.getSendNum());
        Integer newSendNum = testEntity.getSendNum() + 1;
        testEntity.setSendNum(newSendNum);

        //手动注入
        testController = SpringContextUtils.getApplicationContext().getBean(TestController.class);
        //这里根据自己websocket调用方法，调用对应的发送消息的接口即可
//        webSocketController.sendObjMessage("1",testEntity);

        if (newSendNum > 5) {
            //如果次数大于5，则直接关闭此消息的定时任务发送
            testController.removeSchedule(testEntity.getKey());
        }

    }
}

6、websocket消息接收处理

如果后台接收到了前端传来的消息，则将此消息在队列中删除

    @OnMessage
    public void onMessage(String message, Session session) throws Exception{

        System.out.println("接收到了前端的消息：" + message);
        //接收到前端的消息，转化为消息实体，删除对应的定时任务
        try {
            TestEntity testEntity= JSONObject.parseObject(message, TestEntity.class);
            webSocketController.removeSchedule(testEntity.getKey());
            return;
        } catch (Exception e) {
            e.printStackTrace();
        }
        return;
    }

7、结果测试

7.1 测试定时任务超过次数自动关闭

定时任务同一条消息发送超过5次则自动关闭，直接调用生成定时任务的接口即可

结果如下

7.2 测试websocket接收消息删除定时任务

websocket接收到前端传来的消息后，删除发送消息的定时任务，用websocket在线连接方式测试

结果如下

综上，整体的流程已经完成，可以根据需要自行修改定时间隔和次数，有意见和建议欢迎留言！