java
复制代码
@Getter
@ToString
public class WebSocketEvent extends ApplicationEvent {
private final WebSocketType state;
private final String bizId;
private final Object data;
private final LocalDateTime dateTime;
private final String sessionId;
public WebSocketEvent(Object source, WebSocketType state, String sessionId, String bizId, Object data, LocalDateTime dateTime) {
super(source);
this.state = state;
this.sessionId = sessionId;
this.bizId = bizId;
this.data = data;
this.dateTime = dateTime;
}
public WebSocketEvent(Object source, WebSocketType state, String sessionId, String bizId, Object data) {
this(source, state, sessionId, bizId, data, LocalDateTime.now());
}
public WebSocketEvent(Object source, WebSocketType state, String sessionId, String bizId) {
this(source, state, sessionId, bizId, null, LocalDateTime.now());
}
public WebSocketEvent(Object source, WebSocketType state) {
this(source, state, null, null);
}
}
java
复制代码
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.collections4.CollectionUtils;
import org.apache.commons.lang3.StringUtils;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.ApplicationContext;
import org.springframework.stereotype.Component;
import org.springframework.web.util.UriComponentsBuilder;
import javax.websocket.*;
import javax.websocket.server.ServerEndpoint;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.Consumer;
import static java.util.Objects.isNull;
import static java.util.Objects.nonNull;
@Slf4j
@Component
@ServerEndpoint(value = "/websocket")
public class WebSocketServer {
private static final long MAX_IDLE_TIMEOUT = TimeUnit.SECONDS.toMillis(45);
private static final AtomicInteger ONLINE_COUNT = new AtomicInteger(0);
private static final ConcurrentHashMap<String, SessionCache> SESSION_MAP = new ConcurrentHashMap<>();
private static ApplicationContext context;
@Autowired
public void setApplicationContext(ApplicationContext context) {
WebSocketServer.context = context;
}
/**
* 连接建立成功调用的方法
*/
@OnOpen
public void onOpen(Session session) {
final String id = session.getId();
final List<String> allow = getQueryParams(session).get("allow");
if (CollectionUtils.isNotEmpty(allow)) {
final String bizId = allow.get(0);
if (StringUtils.isNotBlank(bizId)) {
final SessionCache put = SESSION_MAP.put(id, new SessionCache(session, bizId));
int cnt = isNull(put) ? ONLINE_COUNT.incrementAndGet() : ONLINE_COUNT.get();
log.info("连接 [{} ({})] 已加入,当前连接数 : {}", id, bizId, cnt);
session.setMaxIdleTimeout(MAX_IDLE_TIMEOUT);
context.publishEvent(new WebSocketEvent(this, WebSocketType.ON_OPEN, id, bizId));
return;
}
}
closeQuietly(session, CloseReason.CloseCodes.CANNOT_ACCEPT);
log.warn("连接 [{}] 被拒绝,没有设置业务标识", id);
}
/**
* 连接关闭调用的方法
*/
@OnClose
public void onClose(Session session) {
final String id = session.getId();
SessionCache remove = SESSION_MAP.remove(id);
final boolean nonNull = nonNull(remove);
final String bizId = nonNull ? remove.bizId : null;
final int cnt = nonNull ? ONLINE_COUNT.decrementAndGet() : ONLINE_COUNT.get();
log.info("连接 [{} ({})] 已断开,当前连接数 : {}", id, bizId, cnt);
context.publishEvent(new WebSocketEvent(this, WebSocketType.ON_CLOSE, id, bizId));
}
/**
* 出现错误
*
* @param session
* @param error
*/
@OnError
public void onError(Session session, Throwable error) {
final String id = session.getId();
final SessionCache sessionCache = SESSION_MAP.get(id);
final String bizId = nonNull(sessionCache) ? sessionCache.bizId : null;
log.warn("连接 [{} ({})] 有错误 : {}\n{}", id, bizId, error.getMessage(), error.getStackTrace());
}
/**
* 收到客户端消息后调用的方法
*
* @param message
* @param session
*/
@OnMessage
public void onMessage(String message, Session session) {
final String id = session.getId();
final SessionCache sessionCache = SESSION_MAP.get(id);
final String bizId = nonNull(sessionCache) ? sessionCache.bizId : null;
log.info("连接 [{} ({})] 有消息 : {}", id, bizId, message);
if (StringUtils.isNotBlank(message)) {
context.publishEvent(new WebSocketEvent(this, WebSocketType.ON_MESSAGE, id, bizId, message));
}
}
public static boolean isOnline(SessionCache cache, long timestamp) {
return nonNull(cache) && cache.session.isOpen() && !cache.isTimeout(timestamp);
}
public static boolean isOnline(String sessionId, long timestamp) {
if (StringUtils.isNotBlank(sessionId)) {
return isOnline(SESSION_MAP.get(sessionId), timestamp);
}
return false;
}
public static void removeTimeout() {
final long now = System.currentTimeMillis();
SESSION_MAP.forEach((key, value) -> {
if (value.isTimeout(now) || !value.session.isOpen()) {
closeQuietly(value.session, CloseReason.CloseCodes.GOING_AWAY);
SESSION_MAP.remove(key);
log.warn("主动断开 Timeout 连接 [{} ({})]", key, value.bizId);
}
});
}
public static void refreshTimestamp(String sessionId) {
if (StringUtils.isNotBlank(sessionId)) {
final SessionCache cache = SESSION_MAP.get(sessionId);
if (nonNull(cache)) {
cache.refreshTimestamp();
}
}
}
public static void removeUnAuthorization(String sessionId) {
final SessionCache cache = SESSION_MAP.get(sessionId);
if (nonNull(cache)) {
closeQuietly(cache.session, CloseReason.CloseCodes.CANNOT_ACCEPT);
SESSION_MAP.remove(sessionId);
log.warn("主动断开 UnAuthorization 连接 [{} ({})]", sessionId, cache.bizId);
}
}
/**
* 发送消息,实践表明,每次浏览器刷新,session会发生变化。
*
* @param sessionId
* @param message
*/
public static void send(String sessionId, String message) {
send(sessionId, message, null);
}
public static void send(String sessionId, String message, Consumer<SendResult> callback) {
final long now = System.currentTimeMillis();
if (StringUtils.isNotBlank(sessionId)) {
final SessionCache cache = SESSION_MAP.get(sessionId);
if (isOnline(cache, now)) {
cache.session.getAsyncRemote().sendText(message, sendResult -> {
if (!sendResult.isOK()) {
Throwable ex = sendResult.getException();
final String bizId = cache.bizId;
log.error("向 连接 [{} ({})] 发送数据 出错 : {}\n{}", sessionId, bizId, ex.getMessage(), ex.getStackTrace());
}
Optional.ofNullable(callback).ifPresent(x -> x.accept(sendResult));
});
}
}
}
public static void broadcast(String message) {
SESSION_MAP.forEach((key, value) -> send(key, message));
}
public static void broadcast(String message, Consumer<SendResult> callback) {
SESSION_MAP.forEach((key, value) -> send(key, message, callback));
}
public static void broadcast(String bizId, String message) {
broadcast(bizId, message, null);
}
public static void broadcast(String bizId, String message, Consumer<SendResult> callback) {
if (StringUtils.isNotBlank(bizId)) {
final long now = System.currentTimeMillis();
SESSION_MAP.forEach((key, value) -> {
if (bizId.equals(value.bizId) && isOnline(value, now)) {
value.session.getAsyncRemote().sendText(message, sendResult -> {
if (!sendResult.isOK()) {
Throwable ex = sendResult.getException();
final String sessionId = value.session.getId();
log.error("向 连接 [{} ({})] 发送数据 出错 : {}\n{}", sessionId, bizId, ex.getMessage(), ex.getStackTrace());
}
Optional.ofNullable(callback).ifPresent(x -> x.accept(sendResult));
});
}
});
}
}
private static Map<String, List<String>> getQueryParams(Session session) {
// return session.getRequestParameterMap();
return UriComponentsBuilder.fromUri(session.getRequestURI()).build().getQueryParams();
}
private static void closeQuietly(Session session) {
closeQuietly(session, CloseReason.CloseCodes.NO_STATUS_CODE);
}
private static void closeQuietly(Session session, CloseReason.CloseCodes closeCode) {
try {
if (session.isOpen()) {
session.close(new CloseReason(closeCode, ""));
}
} catch (Exception ignored) {
}
}
private static final class SessionCache {
private Session session;
private String bizId;
private long timestamp;
public SessionCache(Session session) {
this(session, "");
}
public SessionCache(Session session, String bizId) {
this.session = session;
this.bizId = bizId;
this.timestamp = System.currentTimeMillis();
}
public Session getSession() {
return session;
}
public void setSession(Session session) {
this.session = session;
}
public String getBizId() {
return bizId;
}
public void setBizId(String bizId) {
this.bizId = bizId;
}
public long getTimestamp() {
return timestamp;
}
public void setTimestamp(long timestamp) {
this.timestamp = timestamp;
}
public void refreshTimestamp() {
setTimestamp(System.currentTimeMillis());
}
public boolean isTimeout(long now) {
return now - getTimestamp() > MAX_IDLE_TIMEOUT;
}
}
}