구현한 코드에 대한 PR#

Spring WebSocket 도입 배경#

프로젝트 진행중 외국인 재학생과 한국인 재학생이 소통할 수 있는 기능을 구현하였다. 간단하게 이미지로 보면 다음과 같다.

사용자는 채팅 목록을 확인할 수 있어야 한다
특정 채팅방에 입장하여 사용자와 채팅을 할 수 있어야 한다.

처음에 HTTP 기반 채팅을 구현하고자 하였다.

하지만 HTTP는 요청과 응답이라는 구조로 통신이 이루어지므로, 실시간으로 바뀌는 정보에 대해서는 지속적으로 요청을 해야되는 불편함이 있다. 이런 식으로 매번 연결을 맺고 끊는 작업은 꽤나 많은 비용이 드는 작업으로 비효율적이다.

따라서 실시간으로 사용성을 높이기 위해 채팅을 구현하기 위해 socket 방식을 사용하였다.

구현 과정#

WebSocket의 핵심 흐름은 다음과 같다

클라이언트가 WebSocket 연결을 요청 → 서버에서 인증 및 연결 설정
연결이 확립되면 지속적인 양방향 통신 가능
메시지 전송 및 수신 처리

WebSocketConfig#

@RequiredArgsConstructor
@EnableWebSocket
@Slf4j
public class WebSocketConfig implements WebSocketConfigurer {

    private final ChatHandler chatHandler;
    private final WebSocketAuthInterceptor authInterceptor;

    @Override
    public void registerWebSocketHandlers(WebSocketHandlerRegistry registry) {
        registry.addHandler(chatHandler, "/ws/chat/**")
                .addInterceptors(authInterceptor)
                .setAllowedOrigins("*");
    }
}

위 설정을 통해 클라이언트가 /ws/chat/{roomId} 엔드포인트로 WebSocket 요청을 보내면 ChatHandler가 이를 처리하게 된다.

또한, 프로젝트에서 JWT 인증을 사용하고 있기 때문에, WebSocket 연결 시에도 토큰 검증이 필요하다. 그러나 메시지를 보낼 때마다 토큰을 검증하면 성능 저하가 발생할 수 있으므로, 핸드셰이크 과정에서 한 번만 검증하고 이후 통신에서는 검증을 생략하도록 설계했다.

채팅방 접근 권한 확인#

@Override
public void afterConnectionEstablished(WebSocketSession session) throws Exception {
    Long userId = (Long) session.getAttributes().get("userId");
    Long roomId = (Long) session.getAttributes().get("roomId");

    if (userId == null || roomId == null) {
        session.close(CloseStatus.BAD_DATA);
        return;
    }

    service.addSessionToRoom(roomId, session);
}

afterConnectionEstablished() 메서드는 WebSocket 연결이 성공적으로 이루어졌을 때 호출된다.이때, userId와 roomId를 확인하여 해당 채팅방에 참여할 권한이 있는지 검증한다.(예: A 사용자가 채팅방 1번에만 참여할 수 있는데, 임의로 2번에 접근하려 하면 오류 발생)

권한 검증이 완료되면 service.addSessionToRoom(roomId, session)을 통해 WebSocket 세션을 저장한다.

채팅 메시지 전송#

위 과정이 완료되면 클라이언트는 WebSocket에 정상적으로 연결된 상태이며, 이제 메시지를 주고받을 수 있다.

@Component
@Slf4j
@RequiredArgsConstructor
public class ChatHandler extends TextWebSocketHandler {

    private final ObjectMapper mapper;
    private final ChatService service;

    @Override
    public void handleMessage(WebSocketSession session, WebSocketMessage<?> message) {
        try {
            String payload = (String) message.getPayload();
            ChatMessage chatMessage = mapper.readValue(payload, ChatMessage.class);
            chatMessage.setTime(LocalDateTime.now());
            service.handleAction(chatMessage);
        } catch (Exception e) {
            handleServerError(session, ERROR_UNEXPECTED);
        }
    }
}

클라이언트가 메시지를 보내면 handleMessage() 메서드가 호출된다.이때, JSON 형태의 메시지를 ChatMessage 객체로 변환하고, service.handleAction(chatMessage)을 호출하여 메시지를 처리한다.

public class ChatMessage {

    private MessageType type;
    private Long roomId;
    private Long userId;
    private String message;
    private LocalDateTime time;
}

ChatMessage의 type은 text인지, image인지 구분하고 roomId, userId(보낸 사람Id), 메시지 내용을 포함하여 전달된다.

이제 service.handleAction(chatMessage)을 통해 해당 채팅방에 있는 사용자들에게 메시지를 전송한다.

public void handleAction(ChatMessage chatMessage) throws IllegalArgumentException {
    Chatroom chatRoom = chatRoomRepository.findById(chatMessage.getRoomId())
            .orElseThrow(() -> new IllegalArgumentException(ERROR_CHATROOM_NOT_FOUND));
    if (chatMessage.getType().equals(MessageType.TALK)) {
        saveMessageAndHandleUnreadCount(chatMessage);
        sendMessageToRoom(chatRoom.getId(), chatMessage);
    }
}

public void sendMessageToRoom(Long roomId, ChatMessage message) {
    Set<WebSocketSession> sessions = sessionsPerRoom.get(roomId);
    if (sessions != null) {
        sessions.removeIf(session -> !sendMessage(session, message));
    }
}

public boolean sendMessage(WebSocketSession session, ChatMessage message) {
    try {
         session.sendMessage(new TextMessage(mapper.writeValueAsString(message)));
        return true;
    } catch (IOException e) {
         return false;
    }
}

handleAction에서는 채팅방이 존재하는지 확인 후 메시지를 DB에 저장 후 해당 채팅방의 모든 세션에 전송한다.

현재 종료 방식에서 문제#

클라이언트에서 명시적으로 연결을 종료하면(채팅방 페이지에서 뒤로 가기 버튼을 눌러 채팅 목록 페이지 가기기), 웹소켓 서버에서는 종료되는 순간 afterConnectionClosed 메서드에서 클라이언트 종료를 감지할 수 있다.

@Override
public void afterConnectionClosed(WebSocketSession session, CloseStatus status) {
        service.removeSession(session);
}

하지만 사용자가 명시적을 종료하는 경우가 아닌, 예상하지 못하게 네트워크가 끊긴 경우에도 위와 같이 서버에서 알 수 있을까?

소켓 서버에 접속하고 인터넷 접속을 끄면 다음과 같은 상황이 생긴다

클라이언트 : 인터넷을 껐기 때문에 서버와 연결이 끊어진 상황이다.
서버 : 클라이언트에서 연결이 끊겼는지 실시간으로 감지할 수 없다.

즉 명시적으로 종료 처리를 하지 않았기 때문에 서버에서는 끊겼는지 여부를 실시간으로 모르고 있는 상황이다.

문제
웹소켓 세션을 메모리에 보관하기 때문에 클라이언트에서 네트워크 오류가 발생하여 웹 소켓 세션을 제거하지 않으면 메모리에 계속 쌓이는 문제가 발생한다. (서버를 내리기 전까지 계속 쌓임)

구글링을 해보니 소켓 타임아웃 설정을 하는 방법이 있었지만 연결이 끊겼을 때 빠르게 감지하기 어려웠고, Spring WebSocket 구현 시 isOpen 메서드를 제공하지만, isOpen메서드 만으로 클라이언트의 상태를 완벽히 확인하기 어렵다는 테스트가 있었다.

따라서 ping/pong을 주고 받는 방법을 도입하였다.

세션 관리를 위한 PING-PONG#

간단하게 흐름을 살펴보면 다음과 같다.

서버가 일정 주기마다 “PING” 메시지를 클라이언트에게 전송
클라이언트가 “PONG” 메시지를 응답하면 서버가 웹소켓 세션 만료 시간 업데이트
일정 시간 동안 “PONG”을 받지 못한 세션은 비활성화된 세션으로 간주하고 제거

@Scheduled(fixedRate = 30000)
public void sendPingMessages() {
    long currentTime = System.currentTimeMillis();
    sessionsPerRoom.forEach((roomId, sessions) -> {
        Set<WebSocketSession> invalidSessions;
        try {
            invalidSessions = handlePingAndValidateSessions(roomId, sessions, currentTime);
        } catch (IOException e) {
            return;
        }
        cleanupInvalidSessions(roomId, sessions, invalidSessions);
    });
}

private Set<WebSocketSession> handlePingAndValidateSessions(Long roomId, Set<WebSocketSession> sessions, long currentTime) throws IOException {
    Set<WebSocketSession> invalidSessions = new HashSet<>();
    for (WebSocketSession session : sessions) {
        if (!session.isOpen() || isSessionTimedOut(session, currentTime)) {
            invalidSessions.add(session);
            continue;
        }
        session.sendMessage(new TextMessage("PING"));
    }
    return invalidSessions;
}

@Scheduled를 활용해 30초마다 일정 주기마다 모든 웹소켓 세션의 상태를 확인한다.

만약 이미 timeout된 웹소켓 세션이라면 비정상적으로 종료된 세션이라 판단하고 종료하기 위해 Set<WebSocketSession> invalidSessions에 저장한다.

정상적인 세션들에게는 session.sendMessage(new TextMessage("PING"))를 통해 PING 메시지를 보내준다.

그 후 cleanupInvalidSessions 메서드에서 session.close()을 통해 비정상적 종료라고 판단하고 저장한 세션들을 닫아준다. 이제 비정상적으로 종료된 세션들이 메모리에 남아있지 않아 편안하다.

그리고 이제 PING 메시지를 받은 클라이언트는 연결되어 있다면 PONG 응답을 보낼 것이다.

그럼 서버는 PONG메시지를 보낸 세션들의 timeout 시간을 갱신해주고 만료되지 않은 살아있는 연결이구나를 판단할 수 있게 된다.

@Override
    public void handleMessage(WebSocketSession session, WebSocketMessage<?> message) {
    try {
    String payload = (String) message.getPayload();
        if ("PONG".equalsIgnoreCase(payload.trim())) {
            service.updateLastPongTimestamp(session);
            return;
        }
        ChatMessage chatMessage = mapper.readValue(payload, ChatMessage.class);
        chatMessage.setTime(LocalDateTime.now());
        service.handleAction(chatMessage);
    } catch (IllegalArgumentException e) {
        handleClientError(session, ERROR_INVALID_MESSAGE_FORMAT);
    } 
}

public void updateLastPongTimestamp(WebSocketSession session){
    session.getAttributes().put("timeout", System.currentTimeMillis() + PONG_TIMEOUT);
}

위 코드는 소켓을 통해 메시지를 처리하는 부분이며 "PONG".equalsIgnoreCase(payload.trim()) 조건문을 활용해 PONG 응답인 경우 updateLastPongTimestamp을 통해 만료 시간을 갱신해주었다.

실제 소켓을 통해 단순 실시간 통신만 구현하는 것은 어렵지 않았다.

실제 사용자의 관점에서 통신이 갑자기 끊어지면 웹소켓 세션이 어떻게 관리가 될까?, 리소스가 계속 낭비되는 것은 아닐까? 와 같은 보이지 않은 부분을 생각하는 것이 쉽지 않았던 것 같다.

또한 이론적으로만 알았던 소켓 통신 방법을 구현으로 직접 해보니 동작 원리를 자세히 익힐 수 있었다.

구현한 코드에 대한 PR#

Spring WebSocket 도입 배경#

구현 과정#

WebSocketConfig#

채팅방 접근 권한 확인#

채팅 메시지 전송#

현재 종료 방식에서 문제#

세션 관리를 위한 PING-PONG#

참고 자료#