Injecting spring bean s into WebSocket

Like Listener, websocket does not belong to spring framework and is not a natural bean of spring framework, so spring beans cannot be injected directly into it.

SpringCongifurator

In order to make websocket work better in Spring framework, the configurator of WebSocket Endpoint will inherit Spring Congifurator. We need to introduce relevant jar packages in pom.xml:

<dependency>
     <groupId>org.springframework</groupId>
     <artifactId>spring-websocket</artifactId>
     <version>${spring.framework.version}</version>
     <scope>compile</scope>
</dependency>

We can use Spring Congifurator as a configurator for websocket endpoint

@ServerEndpoint(value = "/chat/{sessionId}",
                configurator = SpringCongifurator.class)
public class ChatEndpoint {
    ... ...
}

If you need a custom configurator, see the following code specifically:

@ServerEndpoint(value = "/chat/{sessionId}",
                encoders = ChatMessageCodec.class,
                decoders = ChatMessageCodec.class,
                configurator = ChatEndpoint.EndpointConfigurator.class)
public class ChatEndpoint {
    @Inject SessionRegistryService sessionRegisterService;
    ... ...

    // Spring Configurator inherits Server Endpoint Config. Configurator (see WebSocket (5): encoder, decoder and configurator#configurator)
    // In this case, we store httpSession in the user Properties of websocket.Session, so we use a custom way. For easy access, we provide two static methods. Of course, we can use webSocketSession.getUserProperties().get(key) to get it.
    // With Spring Configurator, we can use Spring injection in ChatEndpoint, such as the Session Registry Service above, but ChatEndpoint is not a Spring framework and cannot be injected as a bean. If we want to use ChatEndpoint as a Bean, we can configure it in RootContext
    // @Bean
    // public ChatEndpoint chatEndpoint(){
    //      return new ChatEndpoint();
    // }
    // But then the system has only one ChatEndpoint to manage all the connections, rather than one connection per ChatEndpoint, which is much more complicated to handle.
    public static class EndpointConfigurator extends SpringConfigurator{
        private static final String HTTP_SESSION_KEY = "cn.wei.flowingflying.customer_support.http.session";
        private static final String PRINCIPAL_KEY = "cn.wei.flowingflying.customer_support.user.principal";        

        @Override
        public void modifyHandshake(ServerEndpointConfig config, HandshakeRequest request, HandshakeResponse response) {
            super.modifyHandshake(config, request, response);
            HttpSession session = (HttpSession)request.getHttpSession();
            config.getUserProperties().put(HTTP_SESSION_KEY, session);
            config.getUserProperties().put(PRINCIPAL_KEY, UserPrincipal.getPrincipal(session));            
        }

        private static HttpSession getExposedSession(Session session){
            return (HttpSession) session.getUserProperties().get(HTTP_SESSION_KEY);
        }

        private static Principal getExposedPrincipal(Session session){
            return (Principal)session.getUserProperties().get(PRINCIPAL_KEY);
        }       
    }
}

Using Scheduler in websocket

With websocket, in general, a connection corresponds to a websocket instance, that is, we do not configure websocket as a bean in the root context, which results in all connections being corresponded to the same websocket instance. That is, websocket instances are not singleton beans. The @Scheduled you learned before has to load a singleton bean, and we need to find another way.

The following example starts the ping-pong mechanism in websocket through taskSechdule and terminates at the end of the chat. In the example, we followed the example of learning, demonstrating how to end chat when http session is deleted, focusing on how to use consumer.

@ServerEndpoint(value = "/chat/{sessionId}",
                encoders = ChatMessageCodec.class,
                decoders = ChatMessageCodec.class,
                configurator = ChatEndpoint.EndpointConfigurator.class)
public class ChatEndpoint {
    private static final byte[] pongData = "This is PONG country.".getBytes(StandardCharsets.UTF_8);

    private ScheduledFuture<?> pingFuture;

    @Inject SessionRegistryService sessionRegisterService;
    //We can't add @Scheduled directly before sendPing(), pring's @Schedule only supports singleton bean s, and ChatEndpoint is not. We didn't set it as beans in root Context, but each connection has an instance. Therefore, we will use Task Scheduler directly.
    @Inject TaskScheduler taskScheduler;

    // If we don't assign it to an object, the value of this:: httpSession Removed in Method Reference is different each time. To ensure that the same callback function is in registration and cancellation, we need to solidify it.
    private final Consumer<HttpSession> callback = this::httpSessionRemoved;

    /* Used for execution after injection. For classes that support injection, @PostConstruct is supported, not just spring, but Server Endpoint support for WebSocket. When a connection request from websocket client is received, the server endpoint object is created and @postConstruct is executed */
    @PostConstruct
    public void initialize(){
        this.sessionRegisterService.registerOnRemoveCallback(this.callback);
        this.pingFuture = this.taskScheduler.scheduleWithFixedDelay(
                               this::sendPing,
                               new Date(System.currentTimeMillis() + 25_000L), // start time
                               25_000L); //delay 
    }

    // Handling callback function when httpSession is deleted.
    private void httpSessionRemoved(HttpSession httpSession){
        if(httpSession == this.httpSession){
            synchronized(this){
                if(this.closed)
                    return;
                log.info("Chat session ended abruptly by {} logging out.", this.principal.getName());
                this.close(ChatService.ReasonForLeaving.LOGGED_OUT, null);
            }
        }
    }

    //Processing chat Shutdown Specifically
    private void close(ChatService.ReasonForLeaving reason, String unexpected){ 
        this.closed = true;
        if(this.pingFuture.isCancelled())
            this.pingFuture.cancel(true);
        this.sessionRegisterService.deregisterOnRemoveCallback(this.callback);
    } 

    private void sendPing(){//Send ping
        ... ...
        this.wsSession.getBasicRemote().sendPing(ByteBuffer.wrap(ChatEndpoint.pongData));
    } 

    @OnMessage
    public void onPong(PongMessage message){ //Receive pong
        ByteBuffer data = message.getApplicationData();
        if(!Arrays.equals(ChatEndpoint.pongData, data.array()))
             log.warn("Received pong message with incorrect payload.");
        else
             log.debug("Received good pong message.");
   }
    ... ...
}

Related links: My Professional Java for Web Applications related articles