Jakarta websocket handler with Spring managed bean

Question

I have the following websocket client and server for peer-to-peer communication:

package network;

import jakarta.websocket.ContainerProvider;
import jakarta.websocket.Session;
import jakarta.websocket.WebSocketContainer;
import java.net.URI;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.Future;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.boot.context.event.ApplicationReadyEvent;
import org.springframework.context.annotation.Bean;
import org.springframework.context.event.EventListener;
import org.springframework.scheduling.annotation.EnableScheduling;
import org.springframework.scheduling.annotation.Scheduled;
import org.springframework.stereotype.Service;
import org.springframework.web.socket.client.WebSocketClient;
import org.springframework.web.socket.client.standard.StandardWebSocketClient;

@Service
public class WebsocketNetwork {


    private final MessageHandler messageHandler;

    @Bean
    public WebSocketClient webSocketClient() {
        return new StandardWebSocketClient();
    }

    @Autowired
    public WebsocketNetwork(MessageHandler messageHandler) {

        this.messageHandler = messageHandler;

    }

    @EventListener(value = ApplicationReadyEvent.class)
    public void establishConnection() {

        WebSocketContainer container = ContainerProvider.getWebSocketContainer();

        URI address = new URI("ws://localhost:15000/websocket");

        Session session = container.connectToServer(this.messageHandler, address);
    }

}

Handler:

package network;

import jakarta.websocket.ClientEndpoint;
import jakarta.websocket.OnClose;
import jakarta.websocket.OnError;
import jakarta.websocket.OnMessage;
import jakarta.websocket.OnOpen;
import jakarta.websocket.Session;
import jakarta.websocket.server.ServerEndpoint;
import java.util.Map;
import lombok.RequiredArgsConstructor;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.ApplicationEventPublisher;
import org.springframework.stereotype.Component;

@Component
@ServerEndpoint(value = "/websocket", decoders = MessageDecoder.class, encoders = MessageEncoder.class)
@ClientEndpoint(decoders = MessageDecoder.class, encoders = MessageEncoder.class)
@RequiredArgsConstructor(onConstructor = @__(@Autowired))
public class MessageHandler {

    // MessageHandler should have no argument constructor else it fails with:
    /*

    Caused by: java.lang.NoSuchMethodException: network.MessageHandler.<init>()
        at java.base/java.lang.Class.getConstructor0(Class.java:3641) ~[na:na]
        at java.base/java.lang.Class.getConstructor(Class.java:2324) ~[na:na]
        at org.apache.catalina.core.DefaultInstanceManager.newInstance(DefaultInstanceManager.java:135) ~[tomcat-embed-core-10.1.7.jar:10.1.7]
        at org.apache.tomcat.websocket.WsSession.<init>(WsSession.java:275) ~[tomcat-embed-websocket-10.1.7.jar:10.1.7]

    */


    @Autowired
    private ApplicationEventPublisher eventPublisher;       // <-- why this Null? And how to fix it?

    @OnClose
    public void onClose(Session session) {
        System.out.println("Disconnected from server: " + session.getId());
    }

    @OnOpen
    public void onOpen(Session session) {
        System.out.println("Connected to server: " + session.getId());
    }

    // incoming messages
    @OnMessage
    public void receive(NetworkMessage message) {

        // handle incoming message
        
    }

    @OnError
    public void onError(Session session, Throwable throwable) {
        // handle error
        System.err.println(throwable.getMessage());
    }

}

In the above handler, the issue is ApplicationEventPublisher is null even though it's Autowired, when I debugged it the websocket container seems to be creating a newInstance without considering its dependencies, here, does anyone know how to fix it?

Answer 1

To use a Spring-managed bean as a WebSocket handler, you can define a custom javax.websocket.server.ServerEndpointConfig.Configurator class and override its getEndpointInstance() method to return the Spring-managed bean instance.

import javax.websocket.server.ServerEndpointConfig;

import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.ApplicationContext;
import org.springframework.stereotype.Component;

@Component
public class WebSocketConfigurator extends ServerEndpointConfig.Configurator {

    private static ApplicationContext applicationContext;

    @Autowired
    public void setApplicationContext(ApplicationContext applicationContext) {
        WebSocketConfigurator.applicationContext = applicationContext;
    }

    @Override
    public <T> T getEndpointInstance(Class<T> endpointClass) throws InstantiationException {
        return applicationContext.getBean(endpointClass);
    }
}

To use this Configurator class in your Jakarta WebSocket application, you can add the following annotation to your WebSocket endpoint class:

import javax.websocket.server.ServerEndpoint;

@ServerEndpoint(value = "/websocket", configurator = WebSocketConfigurator.class)
public class MessageHandler {
    // ...
}