How to use ConcurrentWebSocketSessionDecorator with a custom WebSocketHandler

user_294657 0
Viewed 514

I have a custom WebSocket handler (for a custom sub-protocol), registered the usual way:

public class WSConfig implements WebSocketConfigurer {

    @Override
    public void registerWebSocketHandlers(WebSocketHandlerRegistry webSocketHandlerRegistry) {
        webSocketHandlerRegistry
                .addHandler(new PerConnectionWebSocketHandler(CustomProtocolHandler.class), endpointUrl);
    }
}

Inside CustomProtocolHandler all methods inherited from WebSocketHandler (e.g. afterConnectionEstablished, handleMessage etc) receive WebSocketSession which is not thread-safe. The official tutorial says ConcurrentWebSocketSessionDecorator can be used to prevent concurrent writes, but by the time my custom handler is hit it is already too late to wrap the session.

public class CustomProtocolHandler implements WebSocketHandler, SubProtocolCapable {

    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) {
        //This is not safe
        session.sendMessage(...);

        //And this this doesn't make sense
        new ConcurrentWebSocketSessionDecorator(session).sendMessage(...);
    }
}

Now, I see that Spring's own SubProtocolWebSocketHandler (which, like mine, implements WebSocketHandler and SubProtocolCapable) does decorate the session internally, and I was hoping I'd be able to use it as the basis, but to create SubProtocolWebSocketHandler I need MessageChannel which comes from spring-messaging. I neither use Spring Messaging, nor have any clue how I would get a MessageChannel. Should I be implementing a custom one?

So, how is ConcurrentWebSocketSessionDecorator intended to be used? I would expect a way to register session-wrapping logic with the webSocketHandlerRegistry, but nope. The only idea I have is to maintain yet another ConcurrentHashMap that maps the raw session (or its ID) to a wrapper, but that's horrible and the one thing I don't need is more state to manage and clean up.

spring-websocket spring websocket java
1 Answers

I ended up making my own version of PerConnectionWebSocketHandler that, like the original, maintains a session-to-handler map, but mine wraps the handler into a holder that also contains a decorated session for the handler to use.

The original PerConnectionWebSocketHandler looks like this:

public class PerConnectionWebSocketHandler implements WebSocketHandler, BeanFactoryAware {

    private final Map<WebSocketSession, WebSocketHandler> handlers = new ConcurrentHashMap<>();

    ...

    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        WebSocketHandler handler = this.provider.getHandler();
        this.handlers.put(session, handler);
        handler.afterConnectionEstablished(session);
    }
}

And mine now looks like this:

public class CustomPerConnectionWebSocketHandler implements WebSocketHandler {

    private final Map<WebSocketSession, HandlerWrapper> handlers = new ConcurrentHashMap<>();

    ...

    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        WebSocketHandler handler = ...;
        //Here's a chance to decorate the session as needed
        WebSocketSession decoratedSession = new ConcurrentWebSocketSessionDecorator(session, sendTimeLimit, sendBufferSizeLimit);
        HandlerWrapper wrapper = new HandlerWrapper(handler, decoratedSession);
        this.handlers.put(session, wrapper);
        wrapper.afterConnectionEstablished();
    }


    @Override
    public void handleMessage(WebSocketSession session, WebSocketMessage<?> message) throws Exception {
         //Delegate all the calls in this fashion
         handlers.get(session).handleMessage(message);
    }

    ...

    private static class HandlerWrapper {

        private final WebSocketHandler handler;
        private final WebSocketSession session;

        HandlerWrapper(WebSocketHandler handler, WebSocketSession session) {
            this.handler = handler;
            this.session = session;
        }

        void afterConnectionEstablished() throws Exception {
            handler.afterConnectionEstablished(session);
        }

        void handleMessage(WebSocketMessage<?> message) throws Exception {
            handler.handleMessage(session, message);
        }

        ... //Other delegating methods
    }
}
user_2946570
Related