KR_WebSocket - somaz94/python-study GitHub Wiki

import websockets
import asyncio

# 기본 WebSocket 서버
async def echo(websocket, path):
    async for message in websocket:
        await websocket.send(f"Echo: {message}")

# 서버 시작
async def start_server():
    server = await websockets.serve(echo, "localhost", 8765)
    await server.wait_closed()

# 실행
asyncio.run(start_server())

# 기본 WebSocket 클라이언트
async def connect_to_server():
    async with websockets.connect("ws://localhost:8765") as websocket:
        await websocket.send("Hello!")
        response = await websocket.recv()
        print(response)

class WebSocketServer:
    def __init__(self):
        self.clients = set()
    
    async def register(self, websocket):
        self.clients.add(websocket)
        try:
            await self.handle_client(websocket)
        finally:
            self.clients.remove(websocket)
    
    async def handle_client(self, websocket):
        try:
            async for message in websocket:
                await self.broadcast(message)
        except websockets.exceptions.ConnectionClosed:
            pass
    
    async def broadcast(self, message):
        if self.clients:
            await asyncio.gather(
                *[client.send(message) for client in self.clients]
            )

class WebSocketClient:
    def __init__(self, uri):
        self.uri = uri
        self.websocket = None
        self.retry_count = 0
        self.max_retries = 5
    
    async def connect(self):
        while self.retry_count < self.max_retries:
            try:
                self.websocket = await websockets.connect(self.uri)
                self.retry_count = 0
                return True
            except websockets.exceptions.ConnectionClosed:
                self.retry_count += 1
                await asyncio.sleep(2 ** self.retry_count)
        return False

import json

class WebSocketProtocol:
    def __init__(self):
        self.handlers = {}
    
    def register_handler(self, event_type, handler):
        self.handlers[event_type] = handler
    
    async def handle_message(self, websocket, message):
        try:
            data = json.loads(message)
            event_type = data.get('type')
            payload = data.get('payload')
            
            if event_type in self.handlers:
                response = await self.handlers[event_type](payload)
                await websocket.send(json.dumps({
                    'type': f"{event_type}_response",
                    'payload': response
                }))
        except json.JSONDecodeError:
            await websocket.send(json.dumps({
                'type': 'error',
                'payload': '잘못된 JSON 형식'
            }))

import jwt
from datetime import datetime, timedelta

class SecureWebSocket:
    def __init__(self, secret_key):
        self.secret_key = secret_key
    
    def create_token(self, user_id):
        payload = {
            'user_id': user_id,
            'exp': datetime.utcnow() + timedelta(hours=1)
        }
        return jwt.encode(payload, self.secret_key, algorithm='HS256')
    
    async def authenticate_connection(self, websocket, path):
        try:
            token = await websocket.recv()
            user_id = self.verify_token(token)
            return user_id
        except ValueError as e:
            await websocket.close()
            return None
            
    def verify_token(self, token):
        try:
            payload = jwt.decode(token, self.secret_key, algorithms=['HS256'])
            return payload['user_id']
        except jwt.PyJWTError:
            raise ValueError("유효하지 않은 토큰")

import asyncio
import orjson  # orjson은 기본 json보다 빠른 처리 제공
import zlib
from concurrent.futures import ProcessPoolExecutor

class OptimizedWebSocket:
    def __init__(self):
        self.clients = set()
        self.process_pool = ProcessPoolExecutor(max_workers=4)
        self.message_queue = asyncio.Queue(maxsize=1000)
        
    async def register(self, websocket):
        self.clients.add(websocket)
        try:
            await self.handle_client(websocket)
        finally:
            self.clients.remove(websocket)
            
    async def handle_client(self, websocket):
        # 메시지 처리 루프
        try:
            async for message in websocket:
                # 압축된 메시지 처리
                if websocket.extensions and 'permessage-deflate' in websocket.extensions:
                    message = zlib.decompress(message)
                    
                # 메시지 큐에 추가
                await self.message_queue.put((websocket, message))
        except websockets.exceptions.ConnectionClosed:
            pass
            
    async def process_message_queue(self):
        """메시지 큐 처리 루프"""
        while True:
            websocket, message = await self.message_queue.get()
            try:
                # 무거운 처리는 프로세스 풀에서 실행
                data = await asyncio.get_event_loop().run_in_executor(
                    self.process_pool, 
                    self.process_message,
                    message
                )
                # 결과 전송
                await websocket.send(orjson.dumps(data))
            except Exception as e:
                print(f"메시지 처리 오류: {str(e)}")
            finally:
                self.message_queue.task_done()
                
    def process_message(self, message):
        """CPU 집약적인 메시지 처리 (별도 프로세스에서 실행)"""
        # 실제 구현에서는 복잡한 연산 수행
        return {"status": "processed", "original": message}
    
    async def start_server(self, host='localhost', port=8765):
        """서버 시작"""
        # 메시지 큐 처리 태스크 시작
        asyncio.create_task(self.process_message_queue())
        
        # 서버 시작
        server = await websockets.serve(
            self.register, 
            host, 
            port,
            compression=None,  # 클라이언트 지원 시 자동으로 압축 활성화
            max_size=10 * 1024 * 1024,  # 최대 메시지 크기 (10MB)
            max_queue=64  # 내부 메시지 큐 크기
        )
        
        return server

from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.middleware.cors import CORSMiddleware
from starlette.websockets import WebSocketState
import asyncio
import json
import logging

# FastAPI 애플리케이션 설정
app = FastAPI(title="WebSocket API")

# CORS 설정
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # 실제 환경에서는 특정 출처만 허용
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

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

# 연결 관리자
class ConnectionManager:
    def __init__(self):
        self.active_connections: dict = {}
        self.lock = asyncio.Lock()
        
    async def connect(self, websocket: WebSocket, client_id: str):
        await websocket.accept()
        async with self.lock:
            self.active_connections[client_id] = websocket
        logger.info(f"클라이언트 연결: {client_id}")
        
    async def disconnect(self, client_id: str):
        async with self.lock:
            if client_id in self.active_connections:
                del self.active_connections[client_id]
                logger.info(f"클라이언트 연결 해제: {client_id}")
    
    async def send_personal_message(self, message: str, client_id: str):
        async with self.lock:
            if client_id in self.active_connections:
                websocket = self.active_connections[client_id]
                if websocket.client_state == WebSocketState.CONNECTED:
                    await websocket.send_text(message)
                    return True
        return False
                
    async def broadcast(self, message: str, exclude: str = None):
        tasks = []
        async with self.lock:
            for client_id, websocket in self.active_connections.items():
                if exclude != client_id and websocket.client_state == WebSocketState.CONNECTED:
                    tasks.append(websocket.send_text(message))
        
        if tasks:
            await asyncio.gather(*tasks)

# 연결 관리자 인스턴스 생성
manager = ConnectionManager()

# WebSocket 엔드포인트
@app.websocket("/ws/{client_id}")
async def websocket_endpoint(websocket: WebSocket, client_id: str):
    await manager.connect(websocket, client_id)
    try:
        while True:
            # 메시지 수신
            data = await websocket.receive_text()
            
            # 메시지 파싱
            try:
                message_data = json.loads(data)
                message_type = message_data.get("type")
                
                # 메시지 유형에 따른 처리
                if message_type == "chat":
                    # 채팅 메시지 브로드캐스트
                    formatted_message = json.dumps({
                        "type": "chat",
                        "sender": client_id,
                        "content": message_data.get("content", ""),
                        "timestamp": asyncio.get_event_loop().time()
                    })
                    await manager.broadcast(formatted_message)
                    
                elif message_type == "private":
                    # 개인 메시지
                    target_id = message_data.get("target")
                    if target_id:
                        formatted_message = json.dumps({
                            "type": "private",
                            "sender": client_id,
                            "content": message_data.get("content", ""),
                            "timestamp": asyncio.get_event_loop().time()
                        })
                        success = await manager.send_personal_message(formatted_message, target_id)
                        # 전송 결과 알림
                        await websocket.send_text(json.dumps({
                            "type": "status",
                            "status": "delivered" if success else "failed",
                            "target": target_id
                        }))
                
            except json.JSONDecodeError:
                await websocket.send_text(json.dumps({
                    "type": "error", 
                    "message": "잘못된 JSON 형식"
                }))
                
    except WebSocketDisconnect:
        await manager.disconnect(client_id)
        # 연결 해제 알림 브로드캐스트
        await manager.broadcast(json.dumps({
            "type": "system",
            "content": f"클라이언트 {client_id}가 나갔습니다"
        }))

# 일반 HTTP 엔드포인트
@app.get("/")
async def get_root():
    return {"message": "WebSocket 서버가 실행 중입니다. '/ws/{client_id}'로 연결하세요."}

# 실행 스크립트
if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)

# 실시간 주식 시세 업데이트 시스템
import asyncio
import json
import random
import logging
import websockets
from datetime import datetime
from typing import Dict, List, Set, Any

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

class StockTickerServer:
    def __init__(self):
        # 클라이언트 관리 
        self.clients: Set[websockets.WebSocketServerProtocol] = set()
        # 클라이언트 구독 정보
        self.subscriptions: Dict[str, Set[websockets.WebSocketServerProtocol]] = {}
        # 주식 시세 데이터
        self.stock_data: Dict[str, Dict[str, Any]] = {
            "AAPL": {"price": 150.0, "change": 0.0, "volume": 0},
            "MSFT": {"price": 250.0, "change": 0.0, "volume": 0},
            "GOOGL": {"price": 2800.0, "change": 0.0, "volume": 0},
            "AMZN": {"price": 3300.0, "change": 0.0, "volume": 0},
            "META": {"price": 330.0, "change": 0.0, "volume": 0},
        }
        
    async def register(self, websocket: websockets.WebSocketServerProtocol):
        """새 클라이언트 등록"""
        self.clients.add(websocket)
        logger.info(f"클라이언트 연결: {len(self.clients)}개 연결 활성화")
        
    async def unregister(self, websocket: websockets.WebSocketServerProtocol):
        """클라이언트 등록 해제"""
        self.clients.discard(websocket)
        
        # 구독에서도 제거
        for symbol in list(self.subscriptions.keys()):
            if websocket in self.subscriptions[symbol]:
                self.subscriptions[symbol].discard(websocket)
                # 비어있는 구독 정리
                if not self.subscriptions[symbol]:
                    del self.subscriptions[symbol]
                    
        logger.info(f"클라이언트 연결 해제: {len(self.clients)}개 연결 활성화")
        
    async def process_message(self, websocket: websockets.WebSocketServerProtocol, message: str):
        """클라이언트 메시지 처리"""
        try:
            data = json.loads(message)
            action = data.get("action")
            
            if action == "subscribe":
                # 주식 구독
                symbols = data.get("symbols", [])
                for symbol in symbols:
                    if symbol in self.stock_data:
                        if symbol not in self.subscriptions:
                            self.subscriptions[symbol] = set()
                        self.subscriptions[symbol].add(websocket)
                        
                        # 현재 시세 즉시 전송
                        await websocket.send(json.dumps({
                            "type": "stock_update",
                            "symbol": symbol,
                            "data": self.stock_data[symbol],
                            "timestamp": datetime.now().isoformat()
                        }))
                
                await websocket.send(json.dumps({
                    "type": "subscription_success",
                    "symbols": symbols
                }))
                        
            elif action == "unsubscribe":
                # 구독 해제
                symbols = data.get("symbols", [])
                for symbol in symbols:
                    if symbol in self.subscriptions and websocket in self.subscriptions[symbol]:
                        self.subscriptions[symbol].discard(websocket)
                        
                await websocket.send(json.dumps({
                    "type": "unsubscription_success",
                    "symbols": symbols
                }))
                
            elif action == "get_available_stocks":
                # 사용 가능한 주식 목록 요청
                await websocket.send(json.dumps({
                    "type": "available_stocks",
                    "symbols": list(self.stock_data.keys())
                }))
                
        except json.JSONDecodeError:
            await websocket.send(json.dumps({
                "type": "error", 
                "message": "잘못된 JSON 형식"
            }))
        except Exception as e:
            logger.error(f"메시지 처리 오류: {str(e)}")
            await websocket.send(json.dumps({
                "type": "error", 
                "message": f"서버 오류: {str(e)}"
            }))
    
    async def handle_client(self, websocket: websockets.WebSocketServerProtocol):
        """클라이언트 연결 처리"""
        await self.register(websocket)
        try:
            async for message in websocket:
                await self.process_message(websocket, message)
        except websockets.exceptions.ConnectionClosedError:
            pass
        finally:
            await self.unregister(websocket)
    
    async def update_stocks(self):
        """주식 시세 업데이트 생성"""
        while True:
            # 모든 주식 시세 업데이트
            for symbol in self.stock_data:
                # 랜덤 가격 변동
                change_percent = (random.random() - 0.5) * 2.0  # -1.0% ~ +1.0%
                price = self.stock_data[symbol]["price"]
                change = price * change_percent / 100.0
                
                # 데이터 업데이트
                self.stock_data[symbol] = {
                    "price": round(price + change, 2),
                    "change": round(change, 2),
                    "change_percent": round(change_percent, 2),
                    "volume": random.randint(100, 10000),
                    "last_update": datetime.now().isoformat()
                }
                
                # 해당 주식을 구독한 클라이언트에게 업데이트 전송
                if symbol in self.subscriptions:
                    update_message = json.dumps({
                        "type": "stock_update",
                        "symbol": symbol,
                        "data": self.stock_data[symbol]
                    })
                    
                    for client in self.subscriptions[symbol].copy():
                        try:
                            await client.send(update_message)
                        except websockets.exceptions.ConnectionClosed:
                            # 닫힌 연결 처리는 handle_client에서 수행
                            pass
            
            # 업데이트 주기 (1초)
            await asyncio.sleep(1)
    
    async def start_server(self, host: str = 'localhost', port: int = 8765):
        """서버 시작"""
        # 주식 업데이트 태스크 시작
        asyncio.create_task(self.update_stocks())
        
        # WebSocket 서버 시작
        server = await websockets.serve(
            self.handle_client, 
            host, 
            port,
            ping_interval=20,
            ping_timeout=60
        )
        
        logger.info(f"주식 시세 서버 시작: {host}:{port}")
        return server

# 서버 실행 코드
async def main():
    server = StockTickerServer()
    await server.start_server()
    await asyncio.Future()  # 무한 실행

if __name__ == "__main__":
    asyncio.run(main())