Websocket Multiple Session

client.py

import asyncio
import websockets
import sys

# 로깅 설정 (클라이언트 측에서도 필요시 사용)
# import logging
# logging.basicConfig(level=logging.INFO)

async def send_messages(websocket):
    """
    stdin으로부터 메시지를 읽어 서버로 전송하는 태스크
    """
    print("서버에 메시지를 보내려면 내용을 입력하고 Enter 키를 누르세요.")
    print("종료하려면 'exit' 또는 'quit'을 입력하세요.")
    while True:
        try:
            # 비동기 환경에서 블로킹 I/O (input)를 안전하게 사용하기 위해 asyncio.to_thread 사용
            message_to_send = await asyncio.to_thread(input, "> ")
            if message_to_send.lower() in ["exit", "quit"]:
                print("연결을 종료합니다...")
                break
            
            await websocket.send(message_to_send)
            # print(f"서버로 전송: {message_to_send}") # 전송 확인 로그 (선택)
        except EOFError: # Ctrl+D 등으로 입력 스트림이 닫혔을 때
            print("입력 스트림이 닫혔습니다. 연결을 종료합니다...")
            break
        except websockets.exceptions.ConnectionClosed:
            print("메시지 전송 중 서버와의 연결이 끊어졌습니다.")
            break
        except Exception as e:
            print(f"메시지 전송 중 오류 발생: {e}")
            break
    # send_messages 루프가 종료되면 receive_messages 태스크도 종료될 수 있도록
    # 웹소켓을 닫거나, 다른 태스크에 종료 신호를 보낼 수 있습니다.
    # 여기서는 websocket.close()가 uri 컨텍스트 매니저에 의해 처리되므로 별도 호출 불필요.

async def receive_messages(websocket):
    """
    서버로부터 메시지를 수신하여 출력하는 태스크
    """
    try:
        async for message in websocket:
            print(f"\n< 서버로부터 수신: {message}")
            print("> ", end="", flush=True) # 다음 입력을 위해 프롬프트 다시 표시
    except websockets.exceptions.ConnectionClosedOK:
        print("서버와의 연결이 정상적으로  종료되었습니다.")
    except websockets.exceptions.ConnectionClosedError:
        print("서버와의 연결이 비정상적으로 끊어졌습니다.")
    except Exception as e:
        print(f"메시지 수신 중 오류 발생: {e}")


async def client_logic():
    uri = "ws://localhost:8765"
    try:
        async with websockets.connect(uri) as websocket:
            print(f"{uri} 에 성공적으로 연결되었습니다.")

            # 메시지 송신 태스크와 수신 태스크를 동시에 실행
            send_task = asyncio.create_task(send_messages(websocket))
            receive_task = asyncio.create_task(receive_messages(websocket))

            # 두 태스크 중 하나라도 완료될 때까지 대기
            # (보통 send_task가 'exit' 입력으로 먼저 종료됨)
            done, pending = await asyncio.wait(
                [send_task, receive_task],
                return_when=asyncio.FIRST_COMPLETED,
            )

            # 완료된 태스크가 있다면, 나머지 보류 중인 태스크를 취소
            for task in pending:
                task.cancel()
            
            # 취소된 태스크가 정리될 시간을 줌
            if pending:
                await asyncio.wait(pending)

    except ConnectionRefusedError:
        print(f"{uri} 에 연결할 수 없습니다. 서버가 실행 중인지 확인하세요.")
    except Exception as e:
        print(f"클라이언트 실행 중 오류 발생: {e}")
    finally:
        print("클라이언트 프로그램이 종료됩니다.")


if __name__ == "__main__":
    try:
        asyncio.run(client_logic())
    except KeyboardInterrupt:
        print("\n클라이언트가 강제 종료됩니다.")

server.py

import asyncio
import websockets
import logging

# 로깅 설정
logging.basicConfig(level=logging.INFO)

# 연결된 클라이언트를 관리하기 위한 집합(set)
CONNECTED_CLIENTS = set()

async def handler(websocket):
    """
    개별 클라이언트 연결을 처리하는 핸들러 함수
    """
    client_address = websocket.remote_address
    logging.info(f"클라이언트 연결됨: {client_address}")
    CONNECTED_CLIENTS.add(websocket)
    logging.info(f"현재 연결된 클라이언트 수: {len(CONNECTED_CLIENTS)}")

    try:
        # 클라이언트로부터 메시지를 계속 수신 대기
        async for message in websocket:
            logging.info(f"클라이언트 {client_address} 로부터 메시지 수신: {message}")

            # 여기에 메시지 처리 로직을 추가할 수 있습니다.
            # 예를 들어, 모든 클라이언트에게 메시지를 브로드캐스트하거나,
            # 특정 조건에 따라 다른 응답을 보낼 수 있습니다.

            response = f"서버가 메시지를 받았습니다: '{message}'"
            await websocket.send(response)
            logging.info(f"클라이언트 {client_address} 에게 응답 전송: {response}")

    except websockets.exceptions.ConnectionClosedOK:
        logging.info(f"클라이언트 {client_address} 연결 정상 종료됨.")
    except websockets.exceptions.ConnectionClosedError as e:
        logging.error(f"클라이언트 {client_address} 연결 비정상 종료됨: {e}")
    except Exception as e:
        logging.error(f"클라이언트 {client_address} 처리 중 오류 발생: {e}")
    finally:
        # 클라이언트 연결이 끊어지면 집합에서 제거
        CONNECTED_CLIENTS.remove(websocket)
        logging.info(f"클라이언트 {client_address} 연결 해제됨. 현재 연결된 클라이언트 수: {len(CONNECTED_CLIENTS)}")

async def main():
    host = "localhost"
    port = 8765
    async with websockets.serve(handler, host, port):
        logging.info(f"WebSocket 서버가 {host}:{port} 에서 실행 중입니다...")
        await asyncio.Future()  # 서버를 계속 실행 상태로 유지

if __name__ == "__main__":
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        logging.info("서버가 종료됩니다.")