Spring WebFlux를 이용한 chat 프로그램 - WebSocket 설정

로그인 처리까지 완료가 되었으니 이제는 채팅을 위해 WebSocket 설정을 진행한다.

WebSocketConfig.java

package com.company.config.websocket;

import java.util.HashMap;
import java.util.Map;

import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.web.reactive.HandlerMapping;
import org.springframework.web.reactive.handler.SimpleUrlHandlerMapping;
import org.springframework.web.reactive.socket.WebSocketHandler;
import org.springframework.web.reactive.socket.server.support.WebSocketHandlerAdapter;

import com.company.config.websocket.handler.ChatWebSocketHandler;

/**
 * WebSocket 핸들러를 등록하고, Spring WebFlux 환경에서 WebSocket 요청을 처리할 수 있도록
 * 필요한 설정을 구성하는 Configuration
 *
 * 주요 역할:
 * - WebSocket 요청 경로("/ws/chat")에 대한 핸들러(ChatWebSocketHandler) 등록
 * - WebSocket 요청 처리를 위한 핸들러 어댑터(WebSocketHandlerAdapter) 등록
 */
@Configuration
public class WebSocketConfig {

	/**
	 * WebSocket 요청 경로에 따라 핸들러 매핑 설정
	 * 클라이언트가 "/ws/chat" 경로로 websocket 연결 요청을 보내면
	 * 해당 요청은 ChatWebSocketHandler에서 처리하도록 한다.
	 *
	 * @param handler
	 * @return
	 */
	@Bean HandlerMapping handlerMapping(ChatWebSocketHandler handler) {
		Map<String, WebSocketHandler> map = new HashMap<>();
		map.put("/ws/chat/{roomId}", handler);	// 동적 경로 지원은 불가 → 접두어만 등록

		SimpleUrlHandlerMapping mapping = new SimpleUrlHandlerMapping();
		mapping.setUrlMap(map);
		mapping.setOrder(-1);
		return mapping;
	}

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

}

채팅방을 만들고 채팅방마다 들어가는 방식으로 처리할 예정이기 때문 roomId 까지를 접두어로 등록한다.

---

ChatWebSocketHandler.java

package com.company.config.websocket.handler;

import java.time.LocalDateTime;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

import org.springframework.stereotype.Component;
import org.springframework.web.reactive.socket.WebSocketHandler;
import org.springframework.web.reactive.socket.WebSocketMessage;
import org.springframework.web.reactive.socket.WebSocketSession;

import com.company.chat.service.ChatMessageService;
import com.company.common.database.document.ChatMessage;
import com.fasterxml.jackson.databind.ObjectMapper;

import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import reactor.core.publisher.Sinks;
import reactor.core.publisher.Sinks.EmitResult;

/**
 * WebSocket을 통해 들어오는 메시지를 브로드캐스트(모든 클라이언트에게 전송)하는 핸들러 
 * - 클라이언트가 메시지를 보내면 Sinks.Many를 통해 모든 연결된 세션에 해당 메시지 전송
 */
@Component
@Slf4j
@RequiredArgsConstructor
public class ChatWebSocketHandler implements WebSocketHandler {

    // 채팅방별 메시지 브로드캐스트 sink를 저장하는 맵
    private final Map<String, Sinks.Many<String>> roomSinkMap = new ConcurrentHashMap<>();

    private final ChatMessageService chatMessageService;

    private final ObjectMapper objectMapper = new ObjectMapper(); // JSON 파싱

    /**
     * 클라이언트와 WebSocket 연결이 수립되면 호출되는 메서드
     *
     * @param session 현재 연결된 WebSocket 세션 객체
     * @return Mono<Void> 비동기적으로 연결을 처리
     */
    @Override
    public Mono<Void> handle(WebSocketSession session) {

        String roomId = extractRoomId(session);

        // 해당 채팅방 sink가 없으면 생성
        Sinks.Many<String> roomSink = roomSinkMap.computeIfAbsent(
                                                    roomId,
                                                    key -> Sinks.many().replay().limit(1) // 마지막 메시지 1개까지 리플레이
                                                );

        // 서버가 클라이언트로 전송할 메시지 스트림
        Flux<WebSocketMessage> outgoing = roomSink.asFlux() // Sink로부터 메시지 Flux 스트림 획득
                                                    .map(session::textMessage); // 문자열을 WebSocket 텍스트 메시지로 변환

        // 클라이언트로부터 수신되는 메시지 스트림
        Flux<String> incoming = session.receive()
                                        .map(WebSocketMessage::getPayloadAsText).doOnNext(json -> {
                                            try {

                                                /*
                                                 * json 형태로 전송된 메시지 내용을 Object형태로 변환한다.
                                                 */
                                                ChatMessage msg = objectMapper.readValue(json, ChatMessage.class);
                                                msg.setTimestamp(LocalDateTime.now());
                                                msg.setRoomId(Long.parseLong(roomId));

                                                // MongoDB 저장
                                                chatMessageService.save(msg).subscribe();

                                                // 브로드캐스트
                                                EmitResult result = roomSink.tryEmitNext(json);
                                                log.info("Emit result: {}", result.name()); // 디버깅을 위해

                                            } catch (Exception e) {
                                                log.error("메시지 파싱 실패: {}", e.getMessage());
                                            }
                                        });

        // 서버 → 클라이언트 메시지 전송
        return session.send(outgoing).and(incoming.then());
    }

    /**
     * 세션 경로에서 채팅방 ID 추출 (예: /ws/chat/12345에서 12345 추출)
     *
     * @param session
     * @return
     */
    private String extractRoomId(WebSocketSession session) {
        String path = session.getHandshakeInfo().getUri().getPath(); // ex) /ws/chat/1
        return path.substring(path.lastIndexOf("/") + 1);
    }

}

---

일잔적으로 WebSocket을 사용하는 방식으로 나는 SimpMessagingTemplate를 사용했었다.

그런데 ChatGPT에게 물어 서 작성을 하다 보니 WebSocketHandler을 사용하도록 알려 줘서 여기서는 WebSocketHandler을 사용한다.

둘의 차이점은 

	SimpMessagingTemplate	WebSocketHandler
레벨	고수준 (STOMP, pub/sub)	저수준 (순수 WebSocket)
쓰임새	메시지 발송, 구독, 브로드캐스트 쉽게	수신/송신 직접 제어, 커스텀 처리
필요조건	STOMP/WebSocketMessageBroker 설정 필요	WebSocketConfigurer 등록
코드 스타일	깔끔, 단순	복잡, 자유도 높음
비유	택배회사 호출해서 배송시키기	직접 들고 뛰는 거

 

이렇단다..ㅋ