Spring websocket并发发送消息异常的解决

@Configuration
@EnableWebSocket // 启动Websocket
public class WebSocketConfig implements WebSocketConfigurer {
    @Override
    public void registerWebSocketHandlers(WebSocketHandlerRegistry registry) {
        registry.addHandler(myHandler(), "/myHandler/**")
            // 添加拦截器，可以获取连接的param和 header 用作认证鉴权
            .addInterceptors(new LakerSessionHandshakeInterceptor())
            // 设置运行跨域
            .setAllowedOrigins("*");
    }
    @Bean
    public ServletServerContainerFactoryBean createWebSocketContainer() {
        ServletServerContainerFactoryBean container = new ServletServerContainerFactoryBean();
        // 设置默认会话空闲超时 以毫秒为单位 非正值意味着无限超时，默认值 0 ，默认没10s检查一次空闲就关闭
        container.setMaxSessionIdleTimeout(10 * 1000L);
        // 设置异步发送消息的默认超时时间 以毫秒为单位 非正值意味着无限超时 ，默认值-1，还没看到作用
        container.setAsyncSendTimeout(10 * 1000L);
        // 设置文本消息的默认最大缓冲区大小 以字符为单位，默认值 8 * 1024
        container.setMaxTextMessageBufferSize(8 * 1024);
        // 设置二进制消息的默认最大缓冲区大小 以字节为单位，默认值 8 * 1024
        container.setMaxBinaryMessageBufferSize(8 * 1024);
        return container;
    }
    @Bean
    public WebSocketHandler myHandler() {
        return new MyHandler();
    }
}

public class LakerSessionHandshakeInterceptor extends HttpSessionHandshakeInterceptor {
    // 拦截器返回false，则不会进行websocket协议的转换
    // 可以获取请求参数做认证鉴权
    @Override
    public boolean beforeHandshake(ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler wsHandler, Map<String, Object> attributes) throws Exception {
        HttpServletRequest servletRequest = ((ServletServerHttpRequest) request).getServletRequest();
        // 将所有查询参数复制到 WebSocketSession 属性映射
        Enumeration<String> parameterNames = servletRequest.getParameterNames();
        while (parameterNames.hasMoreElements()) {
            String parameterName = parameterNames.nextElement();
            // 放入的值可以从WebSocketHandler里面的WebSocketSession拿出来
            attributes.put(parameterName, servletRequest.getParameter(parameterName));
        }
        if (servletRequest.getHeader(HttpHeaders.AUTHORIZATION) != null) {
            setErrorResponse(response, HttpStatus.UNAUTHORIZED);
            return false;
        }
        return true;
    }
    @Override
    public void afterHandshake(ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler wsHandler, Exception exception) {
    }
    private void setErrorResponse(ServerHttpResponse response, HttpStatus status, String errorMessage) {
        response.setStatusCode(status);
        if (errorMessage != null) {
            try {
                objectMapper.writeValue(response.getBody(), new ErrorDetail(status.value(), errorMessage));
            } catch (IOException ioe) {
            }
        }
    }
}

public class MyHandler extends AbstractWebSocketHandler {
     private final Map<String, WebSocketSession> webSocketSessionMap = new ConcurrentHashMap<>();
    //成功连接时
    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        super.afterConnectionEstablished(session);
        // 设置最大报文大小，如果报文过大则会报错的,可以将限制调大。
        // 会覆盖config中的配置。
        session.setBinaryMessageSizeLimit(8 * 1024);
        session.setTextMessageSizeLimit(8 * 1024);
        webSocketSessionMap.put(session.getId(), session);
    }
    //处理textmessage
    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
        super.handleTextMessage(session, message);
        // 有消息就广播下
        for (Map.Entry<String, WebSocketSession> entry : webSocketSessionMap.entrySet()) {
            String s = entry.getKey();
            WebSocketSession webSocketSession = entry.getValue();
            if (webSocketSession.isOpen()) {
            	webSocketSession.sendMessage(new TextMessage(s + ":" + message.getPayload()));
            	log.info("send to {} msg:{}", s, message.getPayload());
            }
        }
    }
    //报错时
    @Override
    public void handleTransportError(WebSocketSession session, Throwable exception) throws Exception {
        super.handleTransportError(session, exception);
          log.warn("handleTransportError:: sessionId: {} {}", session.getId(), exception.getMessage(), exception);
        if (session.isOpen()) {
            try {
                session.close();
            } catch (Exception ex) {
                // ignore
            }
        }
    }
   //连接关闭时
    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
        super.afterConnectionClosed(session, status);
        WebSocketSession session =  webSocketSessionMap.remove(session.getId());
        if (session.isOpen()) {
            try {
                session.close();
            } catch (Exception ex) {
                // ignore
            }
        }
    }
    //处理binarymessage
    @Override
    protected void handleBinaryMessage(WebSocketSession session, BinaryMessage message) throws Exception {
        super.handleBinaryMessage(session, message);
    }
    //处理pong
    @Override
    protected void handlePongMessage(WebSocketSession session, PongMessage message) throws Exception {
        super.handlePongMessage(session, message);
    }
    //是否支持报文拆包
    // 决定是否接受半包，因为websocket协议比较底层，好像Tcp协议一样，如果发送大消息可能会拆成多个小报文。如果不希望处理不完整的报文，希望底层帮忙聚合成完整消息将此方法返回false,这样底层会等待完整报文到达聚合后才回调。
    @Override
    public boolean supportsPartialMessages() {
        return super.supportsPartialMessages();
    }
}

    @Override
    public void afterConnectionEstablished(WebSocketSession session) {
        log.info("{} Connection established!", session.getId());
        // webSocketSessionMap.put(session.getId(), session);
        // 把线程安全的session代理装饰类放到容器里
        webSocketSessionMap.put(session.getId(), new ConcurrentWebSocketSessionDecorator(session, 10 * 1000, 64000));
    }

// 包装一个WebSocketSession以保证一次只有一个线程可以发送消息。
// 如果发送速度慢，后续尝试从其他线程发送更多消息将无法获取刷新锁，消息将被缓冲。届时，将检查指定的缓冲区大小限制和发送时间限制，如果超过限制，会话将关闭。
public class ConcurrentWebSocketSessionDecorator extends WebSocketSessionDecorator {
	private static final Log logger = LogFactory.getLog(ConcurrentWebSocketSessionDecorator.class);
	private final int sendTimeLimit;
	private final int bufferSizeLimit;
	private final OverflowStrategy overflowStrategy;
	@Nullable
	private Consumer<WebSocketMessage<?>> preSendCallback;
	private final Queue<WebSocketMessage<?>> buffer = new LinkedBlockingQueue<>();
	private final AtomicInteger bufferSize = new AtomicInteger();
	private volatile long sendStartTime;
	private volatile boolean limitExceeded;
	private volatile boolean closeInProgress;
	private final Lock flushLock = new ReentrantLock();
	private final Lock closeLock = new ReentrantLock();
	public ConcurrentWebSocketSessionDecorator(WebSocketSession delegate, int sendTimeLimit, int bufferSizeLimit) {
		this(delegate, sendTimeLimit, bufferSizeLimit, OverflowStrategy.TERMINATE);
	}
    // delegate 需要代理的session
    // sendTimeLimit 表示发送**单个消息**的最大时间
    // bufferSizeLimit 表示发送消息的队列最大字节数，不是消息的数量而是消息的总大小
    // overflowStrategy 表示超过限制时的策略有2个
                 // - 断开连接（默认选项）
                 // - 丢弃最老的数据，直到大小满足
	public ConcurrentWebSocketSessionDecorator(
			WebSocketSession delegate, int sendTimeLimit, int bufferSizeLimit, OverflowStrategy overflowStrategy) {
		super(delegate);
		this.sendTimeLimit = sendTimeLimit;
		this.bufferSizeLimit = bufferSizeLimit;
		this.overflowStrategy = overflowStrategy;
	}
    // 返回自当前发送开始以来的时间（毫秒），如果当前没有发送正在进行则返回 0。
    // 即花费的耗时
	public long getTimeSinceSendStarted() {
		long start = this.sendStartTime;
		return (start > 0 ? (System.currentTimeMillis() - start) : 0);
	}
	// 设置在将消息添加到发送缓冲区后调用的回调
	public void setMessageCallback(Consumer<WebSocketMessage<?>> callback) {
		this.preSendCallback = callback;
	}
	@Override
	public void sendMessage(WebSocketMessage<?> message) throws IOException {
        //检查超限了就不发了
		if (shouldNotSend()) {
			return;
		}
		// 消息放到buffer队列
		this.buffer.add(message);
        // 增加bufferSize用于后面判断是不是超限了
		this.bufferSize.addAndGet(message.getPayloadLength());
		// 发送缓冲区后调用的回调
		if (this.preSendCallback != null) {
			this.preSendCallback.accept(message);
		}
		do {
            // 尝试获取锁，发送消息，只有一个线程负责发送所有消息
			if (!tryFlushMessageBuffer()) {
				// 没获取锁的线程，对当前buffer和时间检查，
                // 检查不过就抛异常 然后框架自己会抓取异常关闭当前的连接
				checkSessionLimits();
				break;
			}
		}
		while (!this.buffer.isEmpty() && !shouldNotSend());
	}
	// 超限了 不能发送了
	private boolean shouldNotSend() {
		return (this.limitExceeded || this.closeInProgress);
	}
	// 尝试获取锁 并发送所有缓存的消息
	private boolean tryFlushMessageBuffer() throws IOException {
		if (this.flushLock.tryLock()) {
			try {
                 // 循环发送消息
				while (true) {
                    // 一次拉一个消息
					WebSocketMessage<?> message = this.buffer.poll();
                    // 没消息了 或者 超限了
					if (message == null || shouldNotSend()) {
                        // 退出 完活
						break;
					}
                    // 释放bufferSize
					this.bufferSize.addAndGet(-message.getPayloadLength());
                    // 用于判断单个消息是否发送超时的
					this.sendStartTime = System.currentTimeMillis();
                    // 发送消息
					getDelegate().sendMessage(message);
                    // 重置开始时间
					this.sendStartTime = 0;
				}
			}
			finally {
				this.sendStartTime = 0;
				this.flushLock.unlock();
			}
			return true;
		}
		return false;
	}
   //检查是否超时，是否超过buffer限制
	private void checkSessionLimits() {
        // 应该发送 且 获取到关闭锁
		if (!shouldNotSend() && this.closeLock.tryLock()) {
			try {
                 //检测是否发送超时
				if (getTimeSinceSendStarted() > getSendTimeLimit()) {
					String format = "Send time %d (ms) for session '%s' exceeded the allowed limit %d";
					String reason = String.format(format, getTimeSinceSendStarted(), getId(), getSendTimeLimit());
					limitExceeded(reason);
				}
                 //检测buffer大小，根据策略要么抛异常关闭连接，要么丢弃数据
				else if (getBufferSize() > getBufferSizeLimit()) {
					switch (this.overflowStrategy) {
                        // 关闭连接，抛出异常框架就会关闭连接    
						case TERMINATE:
							String format = "Buffer size %d bytes for session '%s' exceeds the allowed limit %d";
							String reason = String.format(format, getBufferSize(), getId(), getBufferSizeLimit());
							limitExceeded(reason);
							break;
                        // 丢弃老数据    
						case DROP:
							int i = 0;
							while (getBufferSize() > getBufferSizeLimit()) {
								WebSocketMessage<?> message = this.buffer.poll();
								if (message == null) {
									break;
								}
								this.bufferSize.addAndGet(-message.getPayloadLength());
								i++;
							}
							if (logger.isDebugEnabled()) {
								logger.debug("Dropped " + i + " messages, buffer size: " + getBufferSize());
							}
							break;
						default:
							// Should never happen..
							throw new IllegalStateException("Unexpected OverflowStrategy: " + this.overflowStrategy);
					}
				}
			}
			finally {
				this.closeLock.unlock();
			}
		}
	}
	private void limitExceeded(String reason) {
		this.limitExceeded = true;
		throw new SessionLimitExceededException(reason, CloseStatus.SESSION_NOT_RELIABLE);
	}
	@Override
	public void close(CloseStatus status) throws IOException {
		this.closeLock.lock();
		try {
			if (this.closeInProgress) {
				return;
			}
			if (!CloseStatus.SESSION_NOT_RELIABLE.equals(status)) {
				try {
					checkSessionLimits();
				}
				catch (SessionLimitExceededException ex) {
					// Ignore
				}
				if (this.limitExceeded) {
					if (logger.isDebugEnabled()) {
						logger.debug("Changing close status " + status + " to SESSION_NOT_RELIABLE.");
					}
					status = CloseStatus.SESSION_NOT_RELIABLE;
				}
			}
			this.closeInProgress = true;
			super.close(status);
		}
		finally {
			this.closeLock.unlock();
		}
	}
	public enum OverflowStrategy {
		TERMINATE,
		DROP
	}
}