WebSocket in Spring Boot / Spring Framework

Lately I got a chance to get my hands dirty with WebSocket in Spring Boot. The idea was to broadcast the progress of an Async task that took a while to complete its operation.

At first, it felt like a nightmare to me (I over-reacted, yes) but by the time I completed my research of implementing web sockets in my spring boot application, it turned out to be a piece of cake.

“Why another Spring boot web sockets tutorial?” You might be having this question in your mind right now. The reason behind me writing this blog is to share the complete solution to the problem of implementing web socket in spring boot application with an embedded container.

The issue I faced was, after adding WebSockets, the client wasn’t be able to access the server endpoint to my web-sockets. 404, 404, 404… and 404 (Not found). After a continuous struggle of reading the documentation and multiple discussion forums, I was able to find a solution to this.

Let’s start with the implementation:

1. Maven Dependency

First of all you need to add the following dependency in the pom.xml file of your spring boot application:

<dependency>
    <groupId>org.springframework.boot</groupId>
    <artifactId>spring-boot-starter-websocket</artifactId>
</dependency>

We need this dependency as it is a Maven project, and this is going to help us define our ServerEndpoint

2. Create a Socket Endpoint

First I created a class named Socket. This is going to be my socket/endpoint class.

@Component
@ServerEndpoint(value = "/webSocket")
public class Socket {
 private Session session;
 public static Set<Socket> listeners = new CopyOnWriteArraySet<>();
}

After creating the class, I added two annotations on my socket.

a. @Component
The is to define this class as a spring component. It enables a class to be auto-detected when using annotation-based configuration.
b. @ServerEndpoint
ServerEndpoint annotation allows us to enable the web-socket on our socket class. I defined the endpoint with value = “/webSocket”. This is the end-point where our clients/listeners are going to establish their socket connections.

3. Socket Endpoint Configuration

We can configure our endpoint class in two ways:

i. Extend your class with javax.websocket.Endpoint and then customize it by overriding its life-cycle methods.
ii. Use method-level annotations for each method in our web socket.

Spring annotations allow us to write cleaner code, so we’ll use this method-level annotations to define configure our web socket. In order to configure our socket, I will create four methods with socket annotations.

i. @OnOpen

@OnOpen
public void onOpen(Session session) {
    this.session = session;
    listeners.add(this);
}

The method annotated with @OnOpen is invoked whenever a new listener or client is connected to your socket. We have used this method to add the incoming connection to our Set of Socket listeners.

ii. @OnClose

@OnClose
public void onClose(Session session) {
    listeners.remove(this);
}

The method having @OnClose annotation is invoked every time a socket client is disconnected or closes its connection with our socket. We have used this method to remove the session from the Set of our active Socket listeners.

iii. @OnMessage

This annotation allows the client to communicate with the server. A method annotated with @OnMessage will always be invoked whenever a message is sent by the listener to your socket endpoint i.e. your server.

iv. @OnError

The @OnError annotated method will be invoked whenever an error is occurred while trying to communicate with the client.

I’ve a custom static method in our socket class. The method allows you to broadcast messages to all your active listeners.

public static void broadcast(String message) {
    for (Socket listener : listeners) {
        listener.sendMessage(message);
    }
} 

private void sendMessage(String message) {
    try {
        this.session.getBasicRemote().sendText(message);
    } catch (IOException e) {
        log.error("Caught exception while sending message to Session +
this.session.getId(), e.getMessage(), e);
    }
}

4. Setting up Encoder and Decoder

In order to make the communication possible for both the peers, we need to set up the encoder and decoder for our socket and the client. For this, I will create MessageEncoder and MessageDecoder.

The WebSocket can use a Decoder to transform a text message into a Java object and then handle it in the @OnMessage method. The Encoder will be used to convert the object to text and send it to the client.

MessageEncoder.java

Our Encoder will implement Encoder.Text<T> interface, with the following implementation

public class MessageEncoder implements Encoder.Text<String> {

    @Override
    public String encode(String message) throws EncodeException {
        return message;
    }

    @Override
    public void init(EndpointConfig endpointConfig) {
    }

    @Override
    public void destroy() {
        // Close resources (if any used)
    }
}

MessageDecoder.java

For our Decoder, I have implemented the Decoder.Text<T> interface, with the following implementation

public class MessageDecoder implements Decoder.Text<String> {

    @Override
    public String decode(String s) throws DecodeException {
        return s;
    }

    @Override
    public boolean willDecode(String s) {
        return (s != null);
    }

    @Override
    public void init(EndpointConfig endpointConfig) {
        // Custom initialization logic
    }

    @Override
    public void destroy() {
        // Close resources (if any used)
    }
}

Adding Encoder and Decoder to ServerEndpoint

Let’s put together our encoder and decoder by adding them to our @ServerEndpoint annotation.

@Component
@ServerEndpoint(value = "/webSocket", 
                encoders = MessageEncoder.class, 
                decoders = MessageDecoder.class)

After adding our Encoder and Decoder, this is what our Socket class looks like:

@ServerEndpoint(value = "/webSocket", 
                encoders = MessageEncoder.class,
                 decoders = MessageDecoder.class)
public class Socket {
    private Session session;
    public static Set<Socket> listeners = new CopyOnWriteArraySet<>();

    @OnOpen
    public void onOpen(Session session) {
        this.session = session;
        listeners.add(this);
    }

    @OnMessage //Allows the client to send message to the socket.
    public void onMessage(String message) {
        broadcast(message);
    }

    @OnClose
    public void onClose(Session session) {
        listeners.remove(this);
    }

    @OnError
    public void onError(Session session, Throwable throwable) {
        //Error
    }

    public static void broadcast(String message) {
        for (Socket listener : listeners) {
            listener.sendMessage(message);
        }
    }

    private void sendMessage(String message) {
        try {
            this.session.getBasicRemote().sendText(message);
        } catch (IOException e) {
            log.error("Caught exception while sending message to Session Id: " + 
      this.session.getId(), e.getMessage(), e);
        }
    }
}

5. Exporting our WebSocket

This is what took most of my (research) time while trying to implement WebSocket in my application. I tried (almost everything) that I found on the internet in order to make my WebSocket work, but nothing really helped me. I wasn’t able to connect to my Server Endpoint by any of the WebSocket clients.

Finally, I was able to debug the root cause all of that was because I did not have this bean in my application.

What was the issue? Tomcat uses ServletContainerInitializer to find any classes annotated with ServerEndpoint in an application. Spring Boot, on the other hand, doesn’t support ServletContainerInitializer when you’re using any embedded web container.

Therefore, we need to export our ServerEndpoint by creating a bean of ServerEndpointExporter. Let’s do that by creating this Configuration WebSocketConfig class in our application.

@Configuration
public class WebSocketConfig {
    @Bean
    public ServerEndpointExporter serverEndpointExporter() {
        return new ServerEndpointExporter();
    }
}

And we are done!

IMPORTANT: If you are implementing WebSockets in Spring Boot application that doesn’t have an embedded container (e.g. Tomcat for Spring Boot) on it, you won’t be needed to create a ServerEndpointExporter Bean. Step 4 will be the last one for you! 😉

I have added the complete source code to my GitHub repository here.

Feel free to leave your questions and feedback in the comments section below! I would definitely love to hear from you!

Also, if you find this read good enough, do share it with your friend, colleagues or fellow geeks!