springboot中WebSocket实现批量和定向推送

1.添加依赖

注：这是springboot中整合

   <!--字节套实现消息推送-->
   <dependency>
       <groupId>org.springframework.boot</groupId>
       <artifactId>spring-boot-starter-websocket</artifactId>
   </dependency>

2.配置文件

import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.web.socket.config.annotation.EnableWebSocket;
import org.springframework.web.socket.server.standard.ServerEndpointExporter;
@Configuration
@EnableWebSocket
public class WebSocketConfig {
    @Bean
    public ServerEndpointExporter serverEndpointExporter() {
        return new ServerEndpointExporter();
    }
}

工具类

import com.alibaba.fastjson.JSONObject;
import org.springframework.stereotype.Component;
import org.springframework.util.StringUtils;


import javax.websocket.*;
import javax.websocket.server.PathParam;
import javax.websocket.server.ServerEndpoint;
import java.io.IOException;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
@Component
@ServerEndpoint("/websocket/{username}")
public class WebSocketUtil {

    private static int onlineCount = 0;
    private static Map<String, WebSocketUtil> clients = new ConcurrentHashMap<String, WebSocketUtil>();
    private Session session;
    private String username;

    @OnOpen
    public void onOpen(@PathParam("username") String username, Session session) throws IOException {
        this.username = username;
        this.session = session;
        addOnlineCount();
        clients.put(username, this);
        System.out.println(username+"已连接，当前连接数："+getOnlineCount());
    }

    @OnClose
    public void onClose() throws IOException {
        clients.remove(username);
        subOnlineCount();
    }

    @OnMessage
    public void onMessage(String jsonStr) throws IOException {
        if(!StringUtils.isEmpty(jsonStr)){
            JSONObject jsonTo = JSONObject.parseObject(jsonStr);
            //发给一个人
            if (!jsonTo.get("person").equals("all")) {
                sendMessageTo(jsonTo.get("message").toString(), jsonTo.get("person").toString());
            }
            //发给所有的在线用户
            else {
                sendMessageAll(jsonTo.get("message").toString());
            }
        }
    }

    @OnError
    public void onError(Session session, Throwable error) {
        error.printStackTrace();
    }

    public void sendMessageTo(String message, String To) throws IOException {
        for (WebSocketUtil item : clients.values()) {
            String client = item.username.split("_")[0];
            if (client.equals(To)){
                /*
                 * 锁session和使用getBasicRemote，防止高并发时多个线程同时使用同一个session
                 * 同一个session同时全员推送和定向推送时报出过错误，因此必须加锁，若只是一种推送，可用不用锁且使用异步getAsyncRemote
                 */
                synchronized (item.session){
                    item.session.getBasicRemote().sendText(message);
                }
            }
        }
    }

    public void sendMessageAll(String message) throws IOException {
        for (WebSocketUtil item : clients.values()) {
            synchronized (item.session){
                item.session.getBasicRemote().sendText(message);
            }
        }
    }

    public static synchronized int getOnlineCount() {
        return onlineCount;
    }

    public static synchronized void addOnlineCount() {
        WebSocketUtil.onlineCount++;
    }

    public static synchronized void subOnlineCount() {
        WebSocketUtil.onlineCount--;
    }

    public static synchronized Map<String, WebSocketUtil> getClients() {
        return clients;
    }

}

测试

测试时，我使用了定时任务去推送，需要注意同时使用Scheduled定时和WebSocket时会报错，需要给定时注入线程池。

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONObject;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.scheduling.TaskScheduler;
import org.springframework.scheduling.annotation.EnableScheduling;
import org.springframework.scheduling.annotation.Scheduled;
import org.springframework.scheduling.concurrent.ThreadPoolTaskScheduler;

import java.io.IOException;
import java.util.HashMap;
import java.util.Map;

@Configuration
@EnableScheduling
public class Test {
    //同时使用定时和scoket需要实例定时自己的线程池
    @Bean
    public TaskScheduler taskScheduler() {
        ThreadPoolTaskScheduler scheduling = new ThreadPoolTaskScheduler();
        scheduling.setPoolSize(10);
        scheduling.initialize();
        return scheduling;
    }

    //3.添加定时任务
    @Scheduled(cron = "0/5 * * * * ?")
    private void configureTasks() throws IOException {
        Map all=new HashMap();
        Map one=new HashMap();
        System.out.println("发消息了######");
        all.put("person","all");
        all.put("message","所有人看时间："+ System.currentTimeMillis());
        one.put("person","zs");
        one.put("message","张三看时间："+ System.currentTimeMillis());
        WebSocketUtil webSocketUtil = new WebSocketUtil();
        JSONObject allP = new JSONObject(all);
        JSONObject oneP = new JSONObject(one);
        webSocketUtil.onMessage(allP.toJSONString());
        webSocketUtil.onMessage(oneP.toJSONString());
    }

}