refactor: separate the in-memory implementation

cmd/server/handler/chat_handler.go

@@ -12,10 +12,11 @@ import (
 )
 
 type ChatHandler struct {
-	wsServer    *chat.WebSocketServer
-	upgrader    websocket.Upgrader
-	authService service.AuthService
-	chatService service.ChatService
+	wsServer     *chat.WebSocketServer
+	upgrader     websocket.Upgrader
+	stateManager *chat.StateManager
+	authService  service.AuthService
+	chatService  service.ChatService
 }
 
 var upgrader = websocket.Upgrader{
@@ -25,13 +26,17 @@ var upgrader = websocket.Upgrader{
 }
 
 func NewChatController(
-	wsServer *chat.WebSocketServer, authService service.AuthService, chatService service.ChatService,
+	wsServer *chat.WebSocketServer,
+	stateManager chat.StateManager,
+	authService service.AuthService,
+	chatService service.ChatService,
 ) *ChatHandler {
 	return &ChatHandler{
-		wsServer:    wsServer,
-		upgrader:    upgrader,
-		authService: authService,
-		chatService: chatService,
+		wsServer:     wsServer,
+		upgrader:     upgrader,
+		stateManager: &stateManager,
+		authService:  authService,
+		chatService:  chatService,
 	}
 }
 
@@ -64,7 +69,7 @@ func (h *ChatHandler) initializeOrUpdateClient(conn *websocket.Conn, userData *u
 	client := h.wsServer.FindClientByUID(userData.FirebaseUID)
 	if client == nil {
 		// 클라이언트를 찾지 못한 경우 새로운 클라이언트를 생성
-		client = chat.NewClient(conn, h.wsServer, userData.Nickname, userData.FirebaseUID)
+		client = chat.NewClient(conn, *h.stateManager, userData.Nickname, userData.FirebaseUID)
 		// 새 클라이언트를 웹소켓 서버에 등록
 		h.wsServer.RegisterClient(client)
 	} else {

cmd/server/router.go

@@ -70,7 +70,7 @@ func NewRouter(app *firebaseinfra.FirebaseApp) (*echo.Echo, error) {
 	wsServer := chat.NewWebSocketServer(stateManager)
 	go wsServer.Run()
 	chat.InitializeWebSocketServer(ctx, wsServer, chatService)
-	chatHandler := handler.NewChatController(wsServer, authService, *chatService)
+	chatHandler := handler.NewChatController(wsServer, stateManager, authService, *chatService)
 
 	// RegisterChan middlewares
 	logger := zerolog.New(os.Stdout)

internal/chat/client.go

@@ -23,26 +23,23 @@ const (
 	maxMessageSize = 10000
 )
 
-// 각 메시지의 끝을 나타내기 위해 사용
 var newline = []byte{'\n'}
 
 type Client struct {
-	Conn            *websocket.Conn  `json:"-"`
-	WebSocketServer *WebSocketServer `json:"-"`
-	MessageSender   chan []byte      `json:"-"`
-	FbUID           string           `json:"id"`
-	Name            string           `json:"name"`
-	Rooms           map[int64]*Room  `json:"-"`
+	Conn          *websocket.Conn `json:"-"`
+	StateManager  StateManager    `json:"-"`
+	MessageSender chan []byte     `json:"-"`
+	FbUID         string          `json:"id"`
+	Name          string          `json:"name"`
 }
 
-func NewClient(conn *websocket.Conn, wsServer *WebSocketServer, name, fbUID string) *Client {
+func NewClient(conn *websocket.Conn, stateManager StateManager, name, fbUID string) *Client {
 	return &Client{
-		FbUID:           fbUID,
-		Name:            name,
-		Conn:            conn,
-		WebSocketServer: wsServer,
-		MessageSender:   make(chan []byte, 256),
-		Rooms:           make(map[int64]*Room),
+		FbUID:         fbUID,
+		Name:          name,
+		Conn:          conn,
+		StateManager:  stateManager,
+		MessageSender: make(chan []byte, 256),
 	}
 }
 
@@ -151,10 +148,9 @@ func (client *Client) sendPing() *pnd.AppError {
 }
 
 func (client *Client) disconnect() *pnd.AppError {
-	client.WebSocketServer.UnregisterChan <- client
-	for roomID := range client.Rooms {
-		room := client.Rooms[roomID]
-		room.UnregisterChan <- client
+	client.StateManager.UnregisterClient(client)
+	for roomID := range client.StateManager.GetClientRooms(client.FbUID) {
+		client.StateManager.LeaveRoom(roomID, client.FbUID)
 	}
 	close(client.MessageSender)
 	if err := client.Conn.Close(); err != nil {
@@ -172,7 +168,7 @@ func (client *Client) handleNewMessage(jsonMessage []byte, chatService *service.
 	switch message.Action {
 	case SendMessageAction:
 		roomID := message.Room.GetID()
-		if room := client.WebSocketServer.StateManager.FindRoomByID(roomID); room != nil {
+		if room := client.StateManager.FindRoomByID(roomID); room != nil {
 			room.BroadcastChan <- &message
 		}
 	case JoinRoomAction:
@@ -184,9 +180,6 @@ func (client *Client) handleNewMessage(jsonMessage []byte, chatService *service.
 }
 
 func (client *Client) handleJoinRoomMessage(message Message, chatService *service.ChatService) *pnd.AppError {
-	if client.WebSocketServer == nil {
-		return pnd.NewAppError(nil, http.StatusInternalServerError, pnd.ErrCodeUnknown, "WebSocket 서버가 nil입니다.")
-	}
 	if message.Room == nil {
 		return pnd.NewAppError(nil, http.StatusBadRequest, pnd.ErrCodeInvalidBody, "채팅방 정보가 nil입니다.")
 	}
@@ -200,12 +193,12 @@ func (client *Client) handleJoinRoomMessage(message Message, chatService *servic
 		return pnd.NewAppError(nil, http.StatusBadRequest, pnd.ErrCodeInvalidBody, "보낸 사람이 nil입니다.")
 	}
 
-	if _, ok := client.Rooms[message.Room.GetID()]; !ok {
+	if !client.StateManager.IsClientInRoom(client.FbUID, message.Room.GetID()) {
 		if room.RegisterChan == nil {
 			return pnd.NewAppError(nil, http.StatusInternalServerError, pnd.ErrCodeUnknown, "방 등록 채널이 nil입니다.")
 		}
 
-		client.Rooms[message.Room.GetID()] = room
+		client.StateManager.JoinRoom(room.ID, client.FbUID)
 		room.RegisterChan <- client
 		err := client.notifyRoomJoined(room, message.Sender)
 		if err != nil {
@@ -217,30 +210,29 @@ func (client *Client) handleJoinRoomMessage(message Message, chatService *servic
 }
 
 func (client *Client) CreateRoomIfNotExists(message Message, chatService *service.ChatService) (*Room, *pnd.AppError) {
-	room := client.WebSocketServer.StateManager.FindRoomByID(message.Room.GetID())
+	if client.StateManager == nil {
+		return nil, pnd.NewAppError(nil, http.StatusInternalServerError, pnd.ErrCodeUnknown, "StateManager가 nil입니다.")
+	}
+	room := client.StateManager.FindRoomByID(message.Room.GetID())
 	if room == nil {
 		log.Info().Msgf("ID %d의 방을 찾을 수 없어 새 방을 생성합니다.", message.Room.GetID())
 		var err *pnd.AppError
-		room, err = client.WebSocketServer.StateManager.CreateRoom(message.Room.Name, message.Room.RoomType, chatService)
+		room, err = client.StateManager.CreateRoom(message.Room.Name, message.Room.RoomType, chatService, client.StateManager)
 		if err != nil {
 			log.Error().Err(err.Err).Msg("방 생성에 실패했습니다.")
 			return nil, err
 		}
+		return room, nil
 	}
 	return room, nil
 }
 
 func (client *Client) handleLeaveRoomMessage(roomID int64) {
-	room := client.WebSocketServer.StateManager.FindRoomByID(roomID)
-	delete(client.Rooms, room.GetID())
+	client.StateManager.LeaveRoom(roomID, client.FbUID)
+	room := client.StateManager.FindRoomByID(roomID)
 	room.UnregisterChan <- client
 }
 
-func (client *Client) isInRoom(room *Room) bool {
-	_, ok := client.Rooms[room.GetID()]
-	return ok
-}
-
 func (client *Client) notifyRoomJoined(room *Room, sender *Client) *pnd.AppError {
 	message := Message{
 		Action: RoomJoinedAction,

internal/chat/in_memory_state_manager.go

@@ -1,6 +1,7 @@
 package chat
 
 import (
+	"context"
 	"net/http"
 	"sync"
 
@@ -10,15 +11,19 @@ import (
 )
 
 type InMemoryStateManager struct {
-	clients map[string]*Client
-	rooms   map[int64]*Room
-	mutex   sync.RWMutex
+	clients     map[string]*Client
+	rooms       map[int64]*Room
+	clientRooms map[string]map[int64]struct{}
+	roomClients map[int64]map[string]*Client
+	mutex       sync.RWMutex
 }
 
 func NewInMemoryStateManager() *InMemoryStateManager {
 	return &InMemoryStateManager{
-		clients: make(map[string]*Client),
-		rooms:   make(map[int64]*Room),
+		clients:     make(map[string]*Client),
+		rooms:       make(map[int64]*Room),
+		clientRooms: make(map[string]map[int64]struct{}),
+		roomClients: make(map[int64]map[string]*Client),
 	}
 }
 
@@ -57,14 +62,16 @@ func (m *InMemoryStateManager) FindRoomByID(roomID int64) *Room {
 }
 
 func (m *InMemoryStateManager) CreateRoom(
-	name string, roomType chat.RoomType, roomService *service.ChatService,
+	name string, roomType chat.RoomType, roomService *service.ChatService, stateManager StateManager,
 ) (*Room, *pnd.AppError) {
 	m.mutex.Lock()
 	defer m.mutex.Unlock()
-	room, err := NewRoom(name, roomType, roomService)
+	ctx := context.Background()
+	row, err := roomService.CreateRoom(ctx, name, roomType)
 	if err != nil {
-		return nil, err
+		return nil, pnd.NewAppError(err, http.StatusInternalServerError, pnd.ErrCodeRoomCreationFailed, "채팅방 생성에 실패했습니다.")
 	}
+	room := NewRoom(row.ID, row.Name, row.RoomType, stateManager)
 	go room.RunRoom(roomService)
 	m.rooms[room.GetID()] = room
 	return room, nil
@@ -79,6 +86,47 @@ func (m *InMemoryStateManager) BroadcastToClients(message []byte) *pnd.AppError
 	return nil
 }
 
+func (m *InMemoryStateManager) JoinRoom(roomID int64, clientID string) *pnd.AppError {
+	m.mutex.Lock()
+	defer m.mutex.Unlock()
+	if _, ok := m.clientRooms[clientID]; !ok {
+		m.clientRooms[clientID] = make(map[int64]struct{})
+	}
+	m.clientRooms[clientID][roomID] = struct{}{}
+	if _, ok := m.roomClients[roomID]; !ok {
+		m.roomClients[roomID] = make(map[string]*Client)
+	}
+	m.roomClients[roomID][clientID] = m.clients[clientID]
+	return nil
+}
+
+func (m *InMemoryStateManager) LeaveRoom(roomID int64, clientID string) *pnd.AppError {
+	m.mutex.Lock()
+	defer m.mutex.Unlock()
+	delete(m.clientRooms[clientID], roomID)
+	delete(m.roomClients[roomID], clientID)
+	return nil
+}
+
+func (m *InMemoryStateManager) IsClientInRoom(clientID string, roomID int64) bool {
+	m.mutex.RLock()
+	defer m.mutex.RUnlock()
+	_, ok := m.clientRooms[clientID][roomID]
+	return ok
+}
+
+func (m *InMemoryStateManager) GetClientRooms(clientID string) map[int64]struct{} {
+	m.mutex.RLock()
+	defer m.mutex.RUnlock()
+	return m.clientRooms[clientID]
+}
+
+func (m *InMemoryStateManager) GetRoomClients(roomID int64) map[string]*Client {
+	m.mutex.RLock()
+	defer m.mutex.RUnlock()
+	return m.roomClients[roomID]
+}
+
 func (m *InMemoryStateManager) SetRoom(room *Room) *pnd.AppError {
 	m.mutex.Lock()
 	defer m.mutex.Unlock()

internal/chat/init.go

@@ -23,24 +23,22 @@ func InitializeWebSocketServer(
 		// 클라이언트를 생성하거나 기존 클라이언트를 재사용
 		client, exists := clientMap[row.UserInfo.FirebaseUID]
 		if !exists {
-			client = NewClient(nil, wsServer, row.UserInfo.Nickname, row.UserInfo.FirebaseUID)
+			client = NewClient(nil, wsServer.StateManager, row.UserInfo.Nickname, row.UserInfo.FirebaseUID)
 			wsServer.RegisterClient(client)
 			clientMap[row.UserInfo.FirebaseUID] = client
 		}
 
 		// 방을 생성하거나 기존 방을 불러옴
 		room := wsServer.StateManager.FindRoomByID(row.RoomInfo.ID)
-
 		if room == nil {
-			room = room.InitRoom(row.RoomInfo.ID, row.RoomInfo.Name, row.RoomInfo.RoomType)
+			room = NewRoom(row.RoomInfo.ID, row.RoomInfo.Name, row.RoomInfo.RoomType, wsServer.StateManager)
 			wsServer.StateManager.SetRoom(room)
 			go room.RunRoom(chatService)
 		}
 
 		// 클라이언트를 방에 등록
-		if !client.isInRoom(room) {
-			client.Rooms[room.ID] = room
-			room.RegisterChan <- client
+		if !wsServer.StateManager.IsClientInRoom(client.FbUID, room.ID) {
+			wsServer.StateManager.JoinRoom(room.ID, client.FbUID)
 		}
 	}
 	return nil