Implementing a Service Registry using Websockets

package com.ferdisonmezay.websocket.server.configuration;

import org.springframework.context.annotation.Configuration;
import org.springframework.messaging.simp.config.MessageBrokerRegistry;
import org.springframework.web.socket.config.annotation.EnableWebSocketMessageBroker;
import org.springframework.web.socket.config.annotation.StompEndpointRegistry;
import org.springframework.web.socket.config.annotation.WebSocketMessageBrokerConfigurer;

@Configuration
@EnableWebSocketMessageBroker
public class WebsocketConfig implements WebSocketMessageBrokerConfigurer {

    @Override
    public void configureMessageBroker(MessageBrokerRegistry config) {
        config.enableSimpleBroker("/client");
        config.setApplicationDestinationPrefixes("/app");
    }

    @Override
    public void registerStompEndpoints(StompEndpointRegistry registry) {
        registry.addEndpoint("/server").withSockJS();
    }
}

package com.ferdisonmezay.websocket.server.listeners;

import com.ferdisonmezay.websocket.server.service.ClientService;
import lombok.extern.slf4j.Slf4j;
import org.springframework.context.ApplicationListener;
import org.springframework.messaging.MessageHeaders;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.messaging.SessionSubscribeEvent;
import java.util.List;
import java.util.Map;

@Slf4j
@Component
public class SessionSubscribeEventListener implements ApplicationListener<SessionSubscribeEvent> {
    
    private final ClientService clientService;

    public SessionSubscribeEventListener(ClientService clientService) {
        this.clientService = clientService;
    }

    @Override
    public void onApplicationEvent(SessionSubscribeEvent event) {

        MessageHeaders headers = event.getMessage().getHeaders();
        String sessionId = (String) headers.get("simpSessionId");
        Map<String, Object> nativeHeaders = (Map<String, Object>) headers.get("nativeHeaders");
        List<String> destinations = (List<String>) nativeHeaders.get("destination");
        String clientId = destinations.get(0).substring(8);

        log.info("Received subscribe event, clientId: {}, sessionId: {}", clientId, sessionId);
        clientService.registerNewClient(sessionId, clientId);
        clientService.publishClientList();
    }
}

package com.ferdisonmezay.websocket.server.listeners;

import com.ferdisonmezay.websocket.server.service.ClientService;
import lombok.extern.slf4j.Slf4j;
import org.springframework.context.ApplicationListener;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.messaging.SessionDisconnectEvent;

@Slf4j
@Component
public class SessionDisconnectEventListener implements ApplicationListener<SessionDisconnectEvent> {

    private final ClientService clientService;

    public SessionDisconnectEventListener(ClientService clientService) {
        this.clientService = clientService;
    }

    @Override
    public void onApplicationEvent(SessionDisconnectEvent event) {
        log.info("Received disconnect event for session {}", event.getSessionId());
        clientService.removeClient(event.getSessionId());
        clientService.publishClientList();
    }
}

package com.ferdisonmezay.websocket.server.service;

import com.fasterxml.jackson.databind.ObjectMapper;
import com.ferdisonmezay.websocket.server.dto.ClientDto;
import lombok.extern.slf4j.Slf4j;
import org.springframework.messaging.simp.SimpMessagingTemplate;
import org.springframework.stereotype.Service;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.Map;
import java.util.NoSuchElementException;

@Slf4j
@Service
public class ClientService {
    private final Map<String, ClientDto> clientRegistry = new HashMap<>();

    private final SimpMessagingTemplate messageTemplate;

    private final ObjectMapper objectMapper;

    public ClientService(SimpMessagingTemplate messageTemplate, ObjectMapper objectMapper) {
        this.messageTemplate = messageTemplate;
        this.objectMapper = objectMapper;
    }

    public void publishToClient(String client, String data) {
        messageTemplate.convertAndSend(createDestination(client), data);
    }

    public void registerNewClient(String connectionSessionId, String clientId) {
        clientRegistry.put(connectionSessionId, new ClientDto(clientId));
        log.info("Client registered, total number of clients: {}", clientRegistry.size());
    }

    public void removeClient(String connectionSessionId) {
        log.info("Removing client with id {}", clientRegistry.get(connectionSessionId).getClientId());
        clientRegistry.remove(connectionSessionId);
    }

    public void publishClientList() {
        String data;
        try {
            data = objectMapper.writeValueAsString(new ArrayList<>(clientRegistry.values()));
            publishToClient("manager-ui", data);
        } catch (Exception e) {
            log.error("Error creating json!");
        }
    }

    public ClientDto getClientBySessionId(String sessionId) {
        if (!clientRegistry.containsKey(sessionId)) {
            throw new NoSuchElementException();
        }
        return clientRegistry.get(sessionId);
    }

    public ClientDto getClientByClientId(String clientId) {
        return clientRegistry.values().stream()
                .filter(i -> i.getClientId().equals(clientId))
                .findAny().orElseThrow(NoSuchElementException::new);
    }

    public String createDestination(String client) {
        return String.join("/", "", "client", client);
    }
}

package com.ferdisonmezay.websocket.server.controller;

import com.ferdisonmezay.websocket.server.dto.ClientDto;
import com.ferdisonmezay.websocket.server.service.ClientService;
import lombok.extern.slf4j.Slf4j;
import org.springframework.messaging.handler.annotation.DestinationVariable;
import org.springframework.messaging.handler.annotation.MessageMapping;
import org.springframework.messaging.handler.annotation.Payload;
import org.springframework.messaging.simp.SimpMessageHeaderAccessor;
import org.springframework.stereotype.Controller;

import java.util.NoSuchElementException;

@Slf4j
@Controller
public class WebsocketController {
    private final ClientService clientService;

    public WebsocketController(ClientService clientService) {
        this.clientService = clientService;
    }

    @MessageMapping("/{clientId}")
    public void processMessageFromClient(@Payload String message, @DestinationVariable String clientId) {
        log.info("Received: {} for client: {}", message, clientId);
        try {
            ClientDto clientDto = clientService.getClientByClientId(clientId);
            clientService.publishToClient(clientDto.getClientId(), message);
        }
        catch (NoSuchElementException e) {
            log.error("Client {} could not be found!", clientId);
        }
    }

    @MessageMapping("/update")
    public void updateClientStatus(@Payload String message, SimpMessageHeaderAccessor headerAccessor) {
        String sessionId = (String) headerAccessor.getHeader("simpSessionId");
        try {
            ClientDto clientDto = clientService.getClientBySessionId(sessionId);
            clientDto.setActive(Boolean.parseBoolean(message));
            clientService.publishClientList();
            log.info("Received: {} from client: {}", message, clientDto.getClientId());
        }
        catch (NoSuchElementException e) {
            log.error("Session {} could not be found!", sessionId);
        }
    }
}

package com.ferdisonmezay.websocket.client;

import org.springframework.boot.SpringApplication;
import org.springframework.boot.autoconfigure.SpringBootApplication;
import org.springframework.scheduling.annotation.EnableScheduling;

@SpringBootApplication
@EnableScheduling
public class ClientApplication {
    public static void main(String[] args) {
        SpringApplication.run(ClientApplication.class, args);
    }
}

package com.ferdisonmezay.websocket.client.service;

import lombok.extern.slf4j.Slf4j;
import org.springframework.messaging.simp.stomp.StompFrameHandler;
import org.springframework.messaging.simp.stomp.StompHeaders;
import org.springframework.messaging.simp.stomp.StompSession;
import org.springframework.messaging.simp.stomp.StompSessionHandlerAdapter;
import org.springframework.scheduling.annotation.Scheduled;
import org.springframework.stereotype.Service;
import org.springframework.util.concurrent.ListenableFuture;
import org.springframework.web.socket.WebSocketHttpHeaders;
import org.springframework.web.socket.client.standard.StandardWebSocketClient;
import org.springframework.web.socket.messaging.WebSocketStompClient;
import org.springframework.web.socket.sockjs.client.SockJsClient;
import org.springframework.web.socket.sockjs.client.Transport;
import org.springframework.web.socket.sockjs.client.WebSocketTransport;
import org.springframework.web.socket.sockjs.frame.Jackson2SockJsMessageCodec;

import java.lang.reflect.Type;
import java.nio.charset.StandardCharsets;
import java.util.Collections;
import java.util.List;
import java.util.UUID;

@Slf4j
@Service
public class WebsocketClient {

    private static final String SERVER_URL = "ws://localhost:8080/server";
    private static boolean connectedToServer = false;

    private static boolean enabled = false; //we will change value of this variable using management UI

    @Scheduled(fixedDelay = (5 * 1000)) // 5 seconds
    public void scheduledConnectionCheck() {
        log.info("Scheduled check for connection! Current value of 'enabled' flag is {}", enabled);
        if(!connectedToServer) {
            connectToServer();
        }
    }

    public void connectToServer() {
        log.info("Trying to connect to the server!");
        String clientId = UUID.randomUUID().toString();
        StompSession stompSession;
        try {
            stompSession = connect().get();
            subscribe(stompSession, clientId);
            connectedToServer = true;
            log.info("Connection established to {}", SERVER_URL);
        } catch (Exception e) {
            connectedToServer = false;
            log.error("Failed to connect to {}", SERVER_URL);
        }
    }

    private void subscribe(StompSession stompSession, String clientId) {
        stompSession.subscribe("/client/" + clientId, new StompFrameHandler() {

            public Type getPayloadType(StompHeaders stompHeaders) {
                return byte[].class;
            }

            @Override
            public void handleFrame(StompHeaders stompHeaders, Object o) {
                String command = new String((byte[]) o);
                log.info("Received: {}", command);
                if (command.equals("INVERT")) {
                    enabled = !enabled;
                    notifyServer(stompSession);
                }
            }
        });
    }

    private void notifyServer(StompSession stompSession) {
        stompSession.send("/app/update", (enabled ? "true" : "false").getBytes(StandardCharsets.UTF_8));
    }

    private ListenableFuture<StompSession> connect() {

        Transport webSocketTransport = new WebSocketTransport(new StandardWebSocketClient());
        List<Transport> transports = Collections.singletonList(webSocketTransport);

        SockJsClient sockJsClient = new SockJsClient(transports);
        sockJsClient.setMessageCodec(new Jackson2SockJsMessageCodec());

        WebSocketStompClient stompClient = new WebSocketStompClient(sockJsClient);

        return stompClient.connect(SERVER_URL, new WebSocketHttpHeaders(), new CustomHandler());
    }

    private static class CustomHandler extends StompSessionHandlerAdapter {
        @Override
        public void handleTransportError(StompSession session, Throwable exception) {
            log.warn("Connection closed!");
            connectedToServer = false;
            super.handleTransportError(session, exception);
        }
    }
}