server.py

import asyncio
import tornado.web
import tornado.websocket
import tornado.httpserver
import tornado.ioloop
import json
import logging
import multiprocessing
import os
from realtime_connection_coordinator import RedisRealtimeCoordinator

try:
    import uvloop
    asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
except ImportError:
    pass

logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - PID:%(process)d - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)




######################################### set the redis uri here #########################################
REDIS_URI=""

if not REDIS_URI:
    raise RuntimeError("REDIS_URI is not set. Please set the Redis URI like redis://localhost:6379/0")



redis_coordinator = None

class MessageHandler(tornado.websocket.WebSocketHandler):
    __slots__ = ('client_id', 'current_room', 'redis_coordinator', '_message_count')
    
    def initialize(self):
        global redis_coordinator
        self.client_id = None
        self.current_room = None
        self.redis_coordinator = redis_coordinator
        self._message_count = 0
        
    def check_origin(self, origin):
        return True
    
    def get_compression_options(self):
        return {
            'compression_level': 6,
            'mem_level': 5
        }
    
    async def open(self):
        try:
            self.client_id = await self.redis_coordinator.register_client(
                client_id=None, 
                connection=self
            )
            
            await self.send_json({
                'event': 'connected',
                'data': {
                    'client_id': self.client_id,
                    'message': 'Connected to WebSocket server'
                }
            })
            
            logger.info(f"WebSocket opened: {self.client_id}")
            
        except Exception as e:
            logger.error(f"Error in open: {e}", exc_info=True)
    
    async def on_message(self, message):
        try:
            if not await self.redis_coordinator.check_rate_limit(self.client_id, limit=10000):
                await self.send_json({
                    'event': 'error',
                    'data': {'message': 'Rate limit exceeded'}
                })
                return
            
            self._message_count += 1
            
            data = json.loads(message)
            event = data.get('event')
            payload = data.get('data', {})
            
            if event == 'join':
                await self.handle_join(payload)
            elif event == 'leave':
                await self.handle_leave(payload)
            elif event == 'send_message':
                await self.handle_send_message(payload)
            elif event == 'broadcast':
                await self.handle_broadcast(payload)
            elif event == 'get_rooms':
                await self.handle_get_rooms()
            else:
                logger.warning(f"Unknown event: {event}")
                
        except json.JSONDecodeError:
            logger.error(f"Invalid JSON from {self.client_id}")
            await self.send_json({
                'event': 'error',
                'data': {'message': 'Invalid JSON'}
            })
        except Exception as e:
            logger.error(f"Error handling message: {e}", exc_info=True)
    
    async def handle_join(self, data):
        room = data.get('room')
        if not room:
            return await self.send_json({
                'event': 'error',
                'data': {'message': 'Room name required'}
            })
        
        try:
            await self.redis_coordinator.join(self.client_id, room)
            self.current_room = room
            
            await self.send_json({
                'event': 'joined',
                'data': {'room': room}
            })
            
            await self.redis_coordinator.transmit(
                'user_joined',
                {'client_id': self.client_id, 'room': room},
                room=room,
                exclude=self.client_id
            )
            
            logger.info(f"Client {self.client_id} joined room {room}")
            
        except Exception as e:
            logger.error(f"Error joining room: {e}", exc_info=True)
    
    async def handle_leave(self, data):
        room = data.get('room', self.current_room)
        if not room:
            return
        
        try:
            await self.redis_coordinator.leave(self.client_id, room)
            
            if room == self.current_room:
                self.current_room = None
            
            await self.send_json({
                'event': 'left',
                'data': {'room': room}
            })
            
            await self.redis_coordinator.transmit(
                'user_left',
                {'client_id': self.client_id, 'room': room},
                room=room,
                exclude=self.client_id
            )
            
            logger.info(f"Client {self.client_id} left room {room}")
            
        except Exception as e:
            logger.error(f"Error leaving room: {e}", exc_info=True)
    
    async def handle_send_message(self, data):
        text = data.get('text')
        room = data.get('room', self.current_room)
        
        if not text:
            return await self.send_json({
                'event': 'error',
                'data': {'message': 'Message text required'}
            })
        
        if not room:
            return await self.send_json({
                'event': 'error',
                'data': {'message': 'Not in a room'}
            })
        
        if len(text) > 10000:
            return await self.send_json({
                'event': 'error',
                'data': {'message': 'Message too long'}
            })
        
        try:
            await self.redis_coordinator.transmit(
                'message',
                {
                    'from': self.client_id,
                    'text': text,
                    'room': room,
                    'timestamp': tornado.ioloop.IOLoop.current().time()
                },
                room=room,
                exclude=self.client_id
            )
            
            await self.send_json({
                'event': 'message_sent',
                'data': {'status': 'success'}
            })
            
        except Exception as e:
            logger.error(f"Error sending message: {e}", exc_info=True)
    
    async def handle_broadcast(self, data):
        text = data.get('text')
        
        if not text:
            return
        
        try:
            await self.redis_coordinator.transmit(
                'broadcast',
                {
                    'from': self.client_id,
                    'text': text,
                    'timestamp': tornado.ioloop.IOLoop.current().time()
                },
                exclude=self.client_id
            )
            
        except Exception as e:
            logger.error(f"Error broadcasting: {e}", exc_info=True)
    
    async def handle_get_rooms(self):
        try:
            rooms = await self.redis_coordinator.get_rooms_for_client(self.client_id)
            await self.send_json({
                'event': 'rooms',
                'data': {'rooms': list(rooms)}
            })
        except Exception as e:
            logger.error(f"Error getting rooms: {e}", exc_info=True)
    
    def on_close(self):
        if self.client_id:
            tornado.ioloop.IOLoop.current().add_callback(self._async_cleanup)
            logger.info(f"WebSocket closed: {self.client_id} (messages: {self._message_count})")
    
    async def _async_cleanup(self):
        try:
            if self.current_room:
                await self.redis_coordinator.transmit(
                    'user_disconnected',
                    {'client_id': self.client_id},
                    room=self.current_room,
                    exclude=self.client_id
                )
            
            if self.client_id:
                await self.redis_coordinator.unregister_client(self.client_id)
            
            logger.info(f"Cleanup complete: {self.client_id}")
            
        except Exception as e:
            logger.error(f"Error in cleanup: {e}", exc_info=True)
    
    async def send_json(self, data):
        try:
            self.write_message(json.dumps(data))
        except tornado.websocket.WebSocketClosedError:
            logger.warning(f"Tried to send to closed WebSocket: {self.client_id}")

class MetricsHandler(tornado.web.RequestHandler):
    def get(self):
        global redis_coordinator
        metrics = redis_coordinator.get_metrics()
        self.set_header('Content-Type', 'application/json')
        self.write(json.dumps(metrics, indent=2))


def make_app():
    return tornado.web.Application(
        [
            (r"/ws", MessageHandler),
            (r"/metrics", MetricsHandler),
        ],
        websocket_ping_interval=30,
        websocket_ping_timeout=60,
        websocket_max_message_size=10 * 1024 * 1024,
        debug=False
    )

async def init_redis_coordinator():
    global redis_coordinator
    redis_coordinator = RedisRealtimeCoordinator(REDIS_URI)
    await redis_coordinator.start()
    logger.info(f"Redis corordinator initialized in PID {os.getpid()}: {redis_coordinator.server_id}")

def run_server():
    tornado.ioloop.IOLoop.current().run_sync(init_redis_coordinator)
    
    app = make_app()
    server = tornado.httpserver.HTTPServer(
        app,
        max_buffer_size=104857600,
        idle_connection_timeout=300,
        body_timeout=300,
        xheaders=True
    )
    
    port = 5000
    server.listen(port)
    
    logger.info(f"Server started on port {port} (PID: {os.getpid()})")
    logger.info(f"Server ID: {redis_coordinator.server_id}")
    
    tornado.ioloop.IOLoop.current().start()

def main():
    num_processes = multiprocessing.cpu_count()
    
    if num_processes > 1:
        logger.info(f"Starting {num_processes} worker processes")
        
        app = make_app()
        server = tornado.httpserver.HTTPServer(
            app,
            max_buffer_size=104857600,
            idle_connection_timeout=300,
            body_timeout=300,
            xheaders=True
        )
        
        port = 5000
        server.bind(port)
        server.start(num_processes)
        
        tornado.ioloop.IOLoop.current().run_sync(init_redis_coordinator)
        
        logger.info(f"All processes started on port {port}")
        tornado.ioloop.IOLoop.current().start()
    else:
        run_server()

if __name__ == "__main__":
    main()