Spring에서 WebSocket 사용해보기

wujin·2024년 6월 24일

Spring network

WebSocket

WebSocket은 클라이언트와 서버를 연결하고, 실시간으로 양방향 통신을 가능하게 하는 Stateful(상태 유지)한 프로토콜이다.
기존 HTTP 통신은 요청을 보내야만 응답을 받는 단방향 통신인 Stateless 방식이다.
WebSocket은 최초 연결이 이루어지면 클라이언트가 별도의 요청을 보내지 않아도 응답을 받을 수 있다.
기존 HTTP 같이 양쪽 방향으로 송수신이 가능한 양방향 통신이지만 한 번에 하나의 전송만 이루어지도록 설정된 것을 반이중 통신(Half Duplex)이라 하고, WebSocket와 같이 데이터를 동시에 양방향으로 송수신 할 수 있는 것을 전이중 통신(Full Deplex)이라고 한다.

주요 특징과 작동 방식

1. 연결 설정 (Handshake)

WebSocket 연결은 초기에 HTTP 프로토콜을 통해 설정된다.
클라이언트가 서버에 WebSocket 연결을 요청하면, 서버는 이를 확인하고 응답한다. 이 과정을 핸드셰이크라고 부른다.
핸드셰이크가 성공적으로 이루어지면, HTTP 연결은 WebSocket 연결로 업그레이드된다.

2. 양방향 통신 (Full-Deplex Communication)

WebSocket은 클라이언트와 서버 간의 양방향 통신을 지원한다.
한쪽에서 다른 쪽으로 메세지를 보낼 때마다 별도의 요청이나 응답이 필요하지 않다.
채팅, 실시간 알림, 게임 등 실시간 데이터 전송이 필요한 애플리케이션에 매우 유용하다.

3. 지속적인 연결 (Persistent Connection)

WebSocket 연결은 클라이언트나 서버가 명시적으로 연결을 종료할 때까지 유지된다.
이는 HTTP의 비연결 지향적 성격과 달리, 지속적인 연결을 유지하여 필요한 경우 언제든지 데이터를 전송할 수 있게 한다.

4. 경량 프로토콜 (Lightweight Protocol)

WebScoket은 헤더 정보가 적고, 데이터 프레임의 오버헤드가 낮아 효율적인 데이터 전송이 가능하다.
이는 대용량 데이터 전송 시 성능 이점을 제공한다.

5. 프레임 기반 통신 (Frame-based Commuication)

WebScoket은 데이터를 메세지 프레임으로 나누어 전송한다.
각 프레임에는 데이터와 함께 데이터 유형, 데이터의 끝 여부 등을 나타내는 메타 정보가 포함된다.

Spring 예제

WebSocketConfig.java

import com.example.websocker.handler.WebSocketHandler;
import lombok.RequiredArgsConstructor;
import org.springframework.context.annotation.Configuration;
import org.springframework.web.socket.config.annotation.EnableWebSocket;
import org.springframework.web.socket.config.annotation.WebSocketConfigurer;
import org.springframework.web.socket.config.annotation.WebSocketHandlerRegistry;

/**
 * WebSocketConfigurer는 WebSocket을 활성화하고, WebSocketHandler를 등록할 수 있게 해준다.
 * registerWebSocketHandlers() 메서드를 통해 WebSocketHandler를 등록할 수 있다.
 * 또한, registerWebSocketHandlers에 연결할 WebSocket 엔드 포인트("ws/chat")을 등록할 수 있다.
 * setAllowedOrigin은 지정한 Origin에서 오는 요청만 허용한다.
 */

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

    private final WebSocketHandler webSocketHandler;

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

WebSocketHandler.java

import com.example.websocker.service.ChatService;
import com.example.websocker.vo.ChatMessage;
import com.example.websocker.vo.ChatRoom;
import com.fasterxml.jackson.databind.ObjectMapper;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;

/**
 * WebSocket 연결에 대한 로직처리 담당.
 * 클라이언트나 서버에서 메세지를 보내면, WebSocket에서 처리 후 다시 메세지를 보내는 역할 담당
 * 대표 메서드
 *      - afterConnectionEstablished(WebSocketSession session) : WebSocket 연결이 맺어진 이후 이를 처리하는 메서드
 *      - handleMessage(WebSocketSession session, WebSocketMessage<?> message) : 메세지를 받았을때 이를 처리하기 위해 사용
 * WebSession : WebSocket 연결에 대한 세션 정보 제공
 */

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

    private final ObjectMapper objectMapper; // payload를 ChatMessage 객체로 만들어 주기 위한 objectMapper
    private final ChatService chatService;

    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
        // 메세지 가져옴
        String payload = message.getPayload();

        // payload를 ChatMessage 객체로 생성
        ChatMessage chatMessage = objectMapper.readValue(payload, ChatMessage.class);

        // ChatMessage 객체에서 roomId를 가져와 일치하는 room 주입
        ChatRoom chatRoom = chatService.findRoomById(chatMessage.getRoomId());

        chatRoom.handlerActions(session, chatMessage, chatService);
    }
}

ChatMessage.java

import lombok.Getter;
import lombok.Setter;

@Getter
@Setter
public class ChatMessage {

    public enum MessageType {
        ENTER, TALK
    }

    private MessageType type;
    private String roomId;
    private String sender;
    private String message;
}

ChatRoom.java

import java.util.HashSet;
import java.util.Set;

@Getter
public class ChatRoom {

    private final String roomId;
    private final String name;
    private Set<WebSocketSession> sessions = new HashSet<>();

    @Builder
    public ChatRoom(String roomId, String name) {
        this.roomId = roomId;
        this.name = name;
    }

    public void handlerActions(WebSocketSession session, ChatMessage chatMessage, ChatService chatService) {
        // 방에 처음 들어온 경우
        if (chatMessage.getType().equals(ChatMessage.MessageType.ENTER)) {
            sessions.add(session);
            chatMessage.setMessage(chatMessage.getSender() + "님이 입장했습니다.");
        }

        // 메세지 전송
        sendMessage(chatMessage, chatService);
    }

    // 채팅방에 입장해 있는 모든 클라이언트에게 메세지 전송
    private <T> void sendMessage(T message, ChatService chatService) {
        sessions.parallelStream()
                .forEach(session -> chatService.sendMessage(session, message));
    }
}

ChatService.java

import com.example.websocker.vo.ChatRoom;
import com.fasterxml.jackson.databind.ObjectMapper;
import jakarta.annotation.PostConstruct;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Service;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;

import java.io.IOException;
import java.util.*;

@Slf4j
@RequiredArgsConstructor
@Service
public class ChatService {

    private final ObjectMapper objectMapper;
    private Map<String, ChatRoom> chatRooms; // key : roomId, value : chatRoom

    // Bean 의존성 주입이 완료되고 실행되어야 하는 메서드에 사용
    @PostConstruct
    private void init() {
        chatRooms = new LinkedHashMap<>();
    }

    public List<ChatRoom> findAllRoom() {
        return new ArrayList<>(chatRooms.values());
    }

    public ChatRoom findRoomById(String roomId) {
        return chatRooms.get(roomId);
    }

    public ChatRoom createRoom(String name) {
        String randomId = UUID.randomUUID().toString();

        ChatRoom chatRoom = ChatRoom.builder()
                .roomId(randomId)
                .name(name)
                .build();

        // 방 생성 후 방 목록에 추가
        chatRooms.put(randomId, chatRoom);

        return chatRoom;
    }

    public <T> void sendMessage(WebSocketSession session, T message) {
        try {
            session.sendMessage(new TextMessage(objectMapper.writeValueAsString(message)));
        } catch (IOException e) {
            log.error(e.getMessage(), e);
        }
    }
}

ChatController.java

import com.example.websocker.service.ChatService;
import com.example.websocker.vo.ChatRoom;
import lombok.RequiredArgsConstructor;
import org.springframework.web.bind.annotation.*;

import java.util.List;

@RequiredArgsConstructor
@RestController
@RequestMapping("/chat")
public class ChatController {

    private final ChatService chatService;

    /**
     * 방 생성
     */
    @PostMapping
    public ChatRoom createRoom(@RequestBody String name) {
        return chatService.createRoom(name);
    }

    /**
     * 모든 방 목록 조회
     */
    @GetMapping
    public List<ChatRoom> findAllRoom() {
        return chatService.findAllRoom();
    }
}