Web消息推送之SSE
@RestController
@CrossOrigin
@RequestMapping("/sse")
public class SseEmitterServer {
private static final Logger logger = LoggerFactory.getLogger(SseEmitterServer.class);
/**
* 当前连接数,不准确，如果用户多次刷新会导致偏少
*/
private static AtomicInteger count = new AtomicInteger(0);
/**
* 使用map对象，便于根据userId来获取对应的SseEmitter，或者放redis里面
*/
private static Map<String, SseEmitter> sseEmitterMap = new ConcurrentHashMap<>();
/**
* 创建用户连接并返回 SseEmitter
*
* @param userId 用户ID
* @return SseEmitter
*/
@GetMapping(path = "/subscribe/{userId}", produces = {MediaType.TEXT_EVENT_STREAM_VALUE})
public SseEmitter subscribe(@PathVariable("userId") String userId) throws IOException {
// 设置超时时间，0表示不过期。默认30秒，超过时间未完成会抛出异常：AsyncRequestTimeoutException
SseEmitter sseEmitter = new SseEmitter(30_000L);
// 设置前端的重试时间为15s，如果不加这个发送一下，前端就不会显示连接成功
sseEmitter.send("连接成功");
// 注册回调
sseEmitter.onCompletion(completionCallBack(userId));
sseEmitter.onError(errorCallBack(userId));
sseEmitter.onTimeout(timeoutCallBack(userId));
sseEmitterMap.put(userId, sseEmitter);
// 数量+1
count.getAndIncrement();
logger.info("创建新的sse连接，当前用户：{}", userId);
return sseEmitter;
}
/**
* 给指定用户发送信息 -- 单播
*/
@GetMapping(path = "sendMessage")
public void sendMessage(String userId, String message) {
if (sseEmitterMap.containsKey(userId)) {
try {
// (userId).send(message, MediaType.APPLICATION_JSON);
System.out.println(userId + "==" + message);
sseEmitterMap.get(userId).send(message);
} catch (IOException e) {
logger.error("用户[{}]推送异常:{}", userId, e.getMessage());
removeUser(userId);
}
}
}
/**
* 向多人发布消息 -- 组播
* @param groupId 开头标识
* @param message 消息内容
*/
public void groupSendMessage(String groupId, String message) {
if (MapUtils.isNotEmpty(sseEmitterMap)) {
/*Set<String> ids = ().stream().filter(m -> (groupId)).collect(());
batchSendMessage(message, ids);*/
sseEmitterMap.forEach((k, v) -> {
try {
if (k.startsWith(groupId)) {
v.send(message, MediaType.APPLICATION_JSON);
}
} catch (IOException e) {
logger.error("用户[{}]推送异常:{}", k, e.getMessage());
removeUser(k);
}
});
}
}
/**
* 群发所有人 -- 广播
*/
public void batchSendMessage(String message) {
sseEmitterMap.forEach((k, v) -> {
try {
v.send(message, MediaType.APPLICATION_JSON);
} catch (IOException e) {
logger.error("用户[{}]推送异常:{}", k, e.getMessage());
removeUser(k);
}
});
}
/**
* 群发消息
*/
public void batchSendMessage(String message, Set<String> ids) {
ids.forEach(userId -> sendMessage(userId, message));
}
/**
* 移除用户连接
*/
@GetMapping("/close/{userId}")
public void removeUser(@PathVariable("userId") String userId) {
sseEmitterMap.remove(userId);
// 数量-1
count.getAndDecrement();
logger.info("移除用户：{}", userId);
}
/**
* 获取当前连接信息
*/
public List<String> getIds() {
return new ArrayList<>(sseEmitterMap.keySet());
}
/**
* 获取当前连接数量
*/
@GetMapping("/getUserCount")
public int getUserCount() {
return count.intValue();
}
private Runnable completionCallBack(String userId) {
return () -> {
logger.info("结束连接：{}", userId);
// removeUser(userId);
};
}
private Runnable timeoutCallBack(String userId) {
return () -> {
logger.info("连接超时：{}", userId);
removeUser(userId);
};
}
private Consumer<Throwable> errorCallBack(String userId) {
return throwable -> {
logger.info("连接异常：{}", userId);
removeUser(userId);
};
}
}