Spring/spring-framework
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

OrderedMessageChannelDecorator doesn't preclude send limits

Browse Source 
Closes gh-25581
This commit is contained in:
Rossen Stoyanchev
2020-08-28 20:40:15 +01:00
parent c4f4fbc003
commit bb941b6180
8 changed files with 417 additions and 84 deletions
Download Patch File Download Diff File Expand all files Collapse all files

19
spring-websocket/src/main/java/org/springframework/web/socket/handler/ConcurrentWebSocketSessionDecorator.java
View File

@@ -22,10 +22,12 @@ import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReentrantLock;
import java.util.function.Consumer;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.springframework.lang.Nullable;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.WebSocketMessage;
import org.springframework.web.socket.WebSocketSession;
@@ -54,6 +56,10 @@ public class ConcurrentWebSocketSessionDecorator extends WebSocketSessionDecorat
private final OverflowStrategy overflowStrategy;
@Nullable
private Consumer<WebSocketMessage<?>> preSendCallback;
private final Queue<WebSocketMessage<?>> buffer = new LinkedBlockingQueue<>();
private final AtomicInteger bufferSize = new AtomicInteger();
@@ -130,6 +136,15 @@ public class ConcurrentWebSocketSessionDecorator extends WebSocketSessionDecorat
return (start > 0 ? (System.currentTimeMillis() - start) : 0);
}
/**
* Set a callback invoked after a message is added to the send buffer.
* @param callback the callback to invoke
* @since 5.3
*/
public void setMessageCallback(Consumer<WebSocketMessage<?>> callback) {
this.preSendCallback = callback;
}
@Override
public void sendMessage(WebSocketMessage<?> message) throws IOException {
@@ -140,6 +155,10 @@ public class ConcurrentWebSocketSessionDecorator extends WebSocketSessionDecorat
this.buffer.add(message);
this.bufferSize.addAndGet(message.getPayloadLength());
if (this.preSendCallback != null) {
this.preSendCallback.accept(message);
}
do {
if (!tryFlushMessageBuffer()) {
if (logger.isTraceEnabled()) {

9
spring-websocket/src/main/java/org/springframework/web/socket/messaging/StompSubProtocolHandler.java
View File

@@ -39,6 +39,7 @@ import org.springframework.messaging.simp.SimpAttributes;
import org.springframework.messaging.simp.SimpAttributesContextHolder;
import org.springframework.messaging.simp.SimpMessageHeaderAccessor;
import org.springframework.messaging.simp.SimpMessageType;
import org.springframework.messaging.simp.broker.OrderedMessageChannelDecorator;
import org.springframework.messaging.simp.stomp.BufferingStompDecoder;
import org.springframework.messaging.simp.stomp.StompCommand;
import org.springframework.messaging.simp.stomp.StompDecoder;
@@ -57,6 +58,7 @@ import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.ConcurrentWebSocketSessionDecorator;
import org.springframework.web.socket.handler.SessionLimitExceededException;
import org.springframework.web.socket.handler.WebSocketSessionDecorator;
import org.springframework.web.socket.sockjs.transport.SockJsSession;
@@ -461,6 +463,13 @@ public class StompSubProtocolHandler implements SubProtocolHandler, ApplicationE
payload = errorMessage.getPayload();
}
}
Runnable task = OrderedMessageChannelDecorator.getNextMessageTask(message);
if (task != null) {
Assert.isInstanceOf(ConcurrentWebSocketSessionDecorator.class, session);
((ConcurrentWebSocketSessionDecorator) session).setMessageCallback(m -> task.run());
}
sendToClient(session, accessor, payload);
}