Spring Boot WebSocket: Building a Multichannel Chat in Java

As you may have already guessed from the title, the topic for today will be Spring Boot WebSockets. Some time ago, I provided an example of WebSocket chat based on Akka toolkit libraries. However, this chat will have somewhat more features, and a quite different design.

I will skip some parts so as not to duplicate too much content from the previous article. Here you can find a more in-depth intro to WebSockets. Please note that all the code that’s used in this article is also available in the GitHub repository.

Spring Boot WebSocket: Tools Used

Let’s start the technical part of this text with a description of the tools that will be further used to implement the whole application. As I cannot fully grasp how to build a real WebSocket API with classic Spring STOMP overlay, I decided to go for Spring WebFlux and make everything reactive.

• Spring Boot – No modern Java app based on Spring can exist without Spring Boot; all the autoconfiguration is priceless.
• Spring WebFlux – A reactive version of classic Spring, provides quite a nice and descriptive toolkit for handling both WebSockets and REST. I would dare to say that it is the only way to actually get WebSocket support in Spring.
• Mongo – One of the most popular NoSQL databases, I am using it for storing message history.
• Spring Reactive Mongo – Spring Boot starter for handling Mongo access in a reactive fashion. Using reactive in one place but not the other is not the best idea. Thus, I decided to make DB access reactive as well.

Let’s start the implementation!

Spring Boot WebSocket: Implementation

Dependencies and Config

pom.xml

<dependencies>
    <!--Compile-->
    <dependency>
        <groupId>org.springframework.boot</groupId>
        <artifactId>spring-boot-starter-webflux</artifactId>
    </dependency>
    <dependency>
        <groupId>org.springframework.boot</groupId>
        <artifactId>spring-boot-starter-data-mongodb-reactive</artifactId>
    </dependency>
</dependencies>

application.properties

spring.data.mongodb.uri=mongodb://chats-admin:admin@localhost:27017/chats

I prefer .properties over .yml — In my honest opinion, YAML is not readable and non-maintainable on a larger scale.

WebSocketConfig

@Configuration
class WebSocketConfig {

    @Bean
    ChatStore chatStore(MessagesStore messagesStore) {
        return new DefaultChatStore(Clock.systemUTC(), messagesStore);
    }

    @Bean
    WebSocketHandler chatsHandler(ChatStore chatStore) {
        return new ChatsHandler(chatStore);
    }

    @Bean
    SimpleUrlHandlerMapping handlerMapping(WebSocketHandler wsh) {
    Map<String, WebSocketHandler> paths = Map.of("/chats/{id}", wsh);
        return new SimpleUrlHandlerMapping(paths, 1);
    }

    @Bean
    WebSocketHandlerAdapter webSocketHandlerAdapter() {
        return new WebSocketHandlerAdapter();
    }
}

And surprise, all four beans defined here are very important.

• ChatStore – Custom bean for operating on chats, I will go into more details on this bean in the following steps.
• WebSocketHandler – Bean that will store all the logic related to handling WebSocket sessions.
• SimpleUrlHandlerMapping – Responsible for mapping URLs to correct handler full URL for this one will look more or less like this ws://localhost:8080/chats/{id}.
• WebSocketHandlerAdapter – A kind of capability bean it adds WebSockets handling support to Spring Dispatcher Servlet.

ChatsHandler

class ChatsHandler implements WebSocketHandler {

    private final Logger log = LoggerFactory.getLogger(ChatsHandler.class);

    private final ChatStore store;

    ChatsHandler(ChatStore store) {
      this.store = store;
    }

    @Override
    public Mono handle(WebSocketSession session) {
        String[] split = session.getHandshakeInfo()
            .getUri()
            .getPath()
            .split("/");
        String chatIdStr = split[split.length - 1];
        int chatId = Integer.parseInt(chatIdStr);
        ChatMeta chatMeta = store.get(chatId);
        if (chatMeta == null) {
            return session.close(CloseStatus.GOING_AWAY);
        }
        if (!chatMeta.canAddUser()) {
            return session.close(CloseStatus.NOT_ACCEPTABLE);
        }

        String sessionId = session.getId();
        store.addNewUser(chatId, session);
        log.info("New User {} join the chat {}", sessionId, chatId);
        return session
               .receive()
               .map(WebSocketMessage::getPayloadAsText)
               .flatMap(message -> store.addNewMessage(chatId, sessionId, message))
               .flatMap(message -> broadcastToSessions(sessionId, message, store.get(chatId).sessions())
               .doFinally(sig -> store.removeSession(chatId, session.getId()))
               .then();
    }

    private Mono broadcastToSessions(String sessionId, String message, List sessions) {
        return sessions
        .stream()
        .filter(session -> !session.getId().equals(sessionId))
        .map(session -> session.send(Mono.just(session.textMessage(message))))
        .reduce(Mono.empty(), Mono::then);
    }
}

As I mentioned above, here you can find all the logic related to handling WebSocket sessions. First, we parse the ID of a chat from the URL to get the target chat. Responding with different statuses depends on the context present for a particular chat. 

Additionally, I am also broadcasting the message to all the sessions related to particular chat — for users to actually exchange the messages. I have also added doFinally trigger that will clear closed sessions from the chatStore, to reduce redundant communication. As a whole, this code is reactive; there are some restrictions I need to follow. I have tried to make it as simple and readable as possible, if you have any idea how to improve it I am open.

ChatsRouter

@Configuration(proxyBeanMethods = false)
    class ChatRouter {

    private final ChatStore chatStore;

    ChatRouter(ChatStore chatStore) {
        this.chatStore = chatStore;
    }

    @Bean
    RouterFunction routes() {
        return RouterFunctions
        .route(POST("api/v1/chats/create"), e -> create(false))
        .andRoute(POST("api/v1/chats/create-f2f"), e -> create(true))
        .andRoute(GET("api/v1/chats/{id}"), this::get)
        .andRoute(DELETE("api/v1/chats/{id}"), this::delete);
    }
}

WebFlux's approach to defining REST endpoints is quite different from the classic Spring. Above, you can see the definition of 4 endpoints for managing chats. As similar as in the case of Akka implementation I want to have a REST API for managing Chats and WebSocket API for actual handling chats. I will skip the function implementations as they are pretty trivial; you can see them on GitHub.

ChatStore

First, the interface:

public interface ChatStore {
    int create(boolean isF2F);
    void addNewUser(int id, WebSocketSession session);
    Mono addNewMessage(int id, String userId, String message);
    void removeSession(int id, String session);
    ChatMeta get(int id);
    ChatMeta delete(int id);

Then the implementation:

public class DefaultChatStore implements ChatStore {

    private final Map<Integer, ChatMeta> chats;
    private final AtomicInteger idGen;
    private final MessagesStore messagesStore;
    private final Clock clock;

    public DefaultChatStore(Clock clock, MessagesStore store) {
        this.chats = new ConcurrentHashMap<>();
        this.idGen = new AtomicInteger(0);
        this.clock = clock;
        this.messagesStore = store;
    }

    @Override
    public int create(boolean isF2F) {
        int newId = idGen.incrementAndGet();
        ChatMeta chatMeta = chats.computeIfAbsent(newId, id -> {
        if (isF2F) {
            return ChatMeta.ofId(id);
        }
            return ChatMeta.ofIdF2F(id);
        });
        return chatMeta.id;
    }

    @Override
    public void addNewUser(int id, WebSocketSession session) {
        chats.computeIfPresent(id, (k, v) -> v.addUser(session));
    }

    @Override
    public void removeSession(int id, String sessionId) {
      chats.computeIfPresent(id, (k, v) -> v.removeUser(sessionId));
    }

    @Override
    public Mono addNewMessage(int id, String userId, String message) {
        ChatMeta meta = chats.getOrDefault(id, null);
        if (meta != null) {
            Message messageDoc = new Message(id, userId, meta.offset.getAndIncrement(), clock.instant(), message);
            return messagesStore.save(messageDoc)
                    .map(Message::getContent);
        }
        return Mono.empty();
    }
    // omitted

The base of ChatStore is the ConcurrentHashMap that holds the metadata of all open chats. Most of the methods from the interface are self-explanatory, and there is nothing special behind them.

• create – Creates a new chat with a bool attribute denoting if the chat is f2f or group.
• addNewUser – Adds a new user to existing chats.
• removeUser – Removes a user from the existing chat.
• get – Gets the metadata of a chat with an ID.
• delete – Deletes the chat from CMH.

The only complex method here is addNewMessages. It increments the message counter within the chat and persists message content in MongoDB, for durability.

MongoDB

Message Entity

public class Message {
   @Id
   private String id;
   private int chatId;
   private String owner;
   private long offset;
   private Instant timestamp;
   private String content;

A model for message content stored in a database, there are three important fields here:

1. chatId – Represent chat in which a particular message was sent.
2. ownerId – The userId of the message sender.
3. offset – Ordinal number of message within the chat, for retrieval ordering.

MessageStore

public interface MessagesStore extends ReactiveMongoRepository<Message, String> {}

Nothing special, classic Spring Repository, but in a reactive fashion, provides the same set of features as JpaRepository. It is used directly in ChatStore. Additionally, in the main application class, WebsocketsChatApplication, I am activating reactive repositories by using @EnableReactiveMongoRepositories. Without this annotation messageStore from above would not work. And here we go, we have the whole chat implemented. Let’s test it!

Spring Boot WebSocket: Testing

For tests, I’m using Postman and Simple WebSocket Client.

1. I’m creating a new chat using Postman. In the response body, I got a WebSocket URL to the recently created chat.

   
   

2. Now it is time to use them and check if users can communicate with one another. Simple Web Socket Client comes into play here. Thus, I am connecting to the newly created chat here.

   
   

3. Here we are, everything is working, and users can communicate with each other.

   
   

There is one last thing to do. Let’s spend a moment looking at things that can be done better.

What Can Be Done Better

As what I have just built is the most basic chat app, there are a few (or in fact quite a lot) things that may be done better. Below, I have listed the things I find worthy of improvement:

• Authentication and rejoining support – Right now, everything is based on the sessionId. It is not an optimal approach. It would be better to have some authentication in place and actual rejoining based on user data.
• Sending attachments – For now, the chat only supports simple text messages. While texting is the basic function of a chat, users enjoy exchanging images and audio files, too.
• Tests – There are no tests for now, but why leave it like this? Tests are always a good idea.
• Overflow in offset – Currency, it is a simple int. If we were to track the offset for a very long time, it would overflow sooner or later.

Summary

Et voilà! The Spring Boot WebSocket chat is implemented, and the main task is done. You have some ideas on what to develop in the next steps.

Please keep in mind that this chat case is very simple, and it will require lots of changes and development for any type of commercial project.

Anyway, I hope that you learned something new while reading this article.

Thank you for your time.

These other resources might interest you:

• Lock-Free Programming in Java
• 7 API Integration Patterns