""" AUTHOR: Khushal P Soonderji DATE: Create: Tuesday, 30th Sept., 2025 OBJECTIVE: A namespace to REFERENCES: N/A DOWNLOADS: N/A """ # ***************************************************************************************************************** # ***** **** # *** IMPORT *** # ***** **** # ***************************************************************************************************************** # To make sibling directories accessible for imports: import sys from engineio.base_client import connected_clients sys.path.append(".") sys.path.append("..") # My utils: from utils_v2.date_time import date_time # To work with SocketIO: import socketio # To maintain the app's state: from wsio_v2.app_state import AppState # Helper functions: from wsio_v2.helpers import helpers # For pre-modelled responses: from utils_v2.api.codes import HttpCodes, StatusCodes from utils_v2.api.response import ResponseModel # Models: from models.core.user import CoreUserInfoModel # For asynchronous activities: import asyncio # To work with various datatypes: from typing import List # ***************************************************************************************************************** # ***** **** # *** MACROS / ONE-TIME INIT *** # ***** **** # ***************************************************************************************************************** # --- Nothing Yet # ***************************************************************************************************************** # ***** **** # *** VARIABLES *** # ***** **** # ***************************************************************************************************************** # --- Nothing Yet # ***************************************************************************************************************** # ***** **** # *** FUNCTIONS *** # ***** **** # ***************************************************************************************************************** class TicksNamespace(socketio.AsyncNamespace): def __init__( self, namespace: str, app_state: AppState, ): """ To set up the namespace that will handle tick data. :param namespace: The string that defines the path of the namespace. E.g.: "/test". :param app_state: The AppState that defines the state of the app and carries any custom variables that will be needed across various namespaces. It's a good way to pass instances of database connections and other shared variables. """ super().__init__(namespace) self.app_state = app_state self.app_state.printer("Registered!") async def on_connect(self, sid, environ, *args) -> bool | None: """ The handler that validates and accepts or rejects incoming connection requests. :param sid: The socketio session id. :param environ: The environment information. :param args: The arguments passed to the handler. :returns: True if the connection was accepted, else False. """ self.app_state.printer("Connection Attempt", sid) # Ensure that the app has been initialized: if not self.app_state.init_done: async with self.app_state.exclusive_lock: self.app_state.init_done = await helpers.init(app_state = self.app_state) if self.app_state.init_done: self.server.start_background_task( self.get_ticks_from_kafka, fetch_count = 1_000, fetch_timeout = 1.0 ) # If the attempt to initialize fails: if not self.app_state.init_done: self.app_state.printer("SOCKET REJECTED: Init. pending.", sid) return False # Check the origin of the incoming request: self.app_state.printer("Checking origin.") origin = environ.get("HTTP_ORIGIN", "???") self.app_state.printer(origin) if not helpers.origin_is_allowed(origin, app_state = self.app_state): self.app_state.printer("SOCKET REJECTED: Bad origin.", sid, origin) return False # Get the session token from the incoming request: self.app_state.printer("Checking session token.") session_token = environ.get("HTTP_X_SESSION_TOKEN") session_token_hint = f"...{session_token[-5:]}" self.app_state.printer(session_token_hint) if not session_token and len(args) > 0: session_token = args[0].get("X-Session-Token") if not session_token: self.app_state.printer("SOCKET REJECTED: No session token.", sid) return False # Get the user's details from the session token: self.app_state.printer("Fetching user info.") user_info = await self.app_state.redis_cache.get(key = session_token) if not user_info: self.app_state.printer("SOCKET REJECTED: Invalid session token.", sid) return False user_info = CoreUserInfoModel(**user_info) # Get the user's watchlist and note down the details. # Consider the following structure for a user's info: redis_key = f"wsio_v2_{sid}" async with self.app_state.exclusive_lock: self.app_state.connected_clients[sid] = { "user": user_info, "redisKey": redis_key, "rooms": [] } sid_cached = await self.app_state.redis_cache.set( key = redis_key, value = { "srvr": self.app_state.SERVER_HOSTNAME, "sid": sid, "usr": user_info, } ) # Done here: self.app_state.printer("SOCKET ACCEPTED.", sid, sid_cached) return True async def on_disconnect(self, sid, reason, *args): """ The handler that notes down that the user has been disconnected and frees up any relevant resources. :param sid: The socketio session id. :param reason: The reason for termination. :param args: The arguments passed to the handler. :returns: None. """ self.app_state.printer("On Disconnect", sid) # De-register the disconnect in the app state, and on the cache server: client_info = {} async with self.app_state.exclusive_lock: client_info = self.app_state.connected_clients.pop(sid, None) sid_uncached = await self.app_state.redis_cache.delete(key = client_info["redisKey"]) if client_info else False self.app_state.printer("SOCKET DISCONNECTED", sid, reason, sid_uncached) async def on_echo(self, sid, data): """ A test event that can be used to see if the namespace is up and running, or not. :param sid: The socketio session id. :param data: The data sent by the client. :returns: None. """ self.app_state.printer("Event", sid, data, type(data).__name__) await self.emit("echo", data, to = sid) async def on_set_watchlist(self, sid, data) -> None: """ To handle the request to set the watchlist for a particular user. :param sid: The socketio session id. :param data: The data sent by the client. :returns: None. """ self.app_state.printer("Setting Watchlist", sid) # Start by assuming success: response = ResponseModel( status_code = StatusCodes.OK, message = "Watchlist set successfully.", data = None ) # Prepare a list of tickers from the request: tickers_list = [] # Remove the user from all rooms (tickers): if response.success: try: rooms = self.rooms(sid) for room in rooms: if room != sid: # ... Don't remove them from their private room. await self.leave_room(sid, room) except Exception as e: self.app_state.printer("EXCEPTION!", e) response.status_code = StatusCodes.FAILED response.message = "Failed to reset previous watchlist." # Now add the user back to just those rooms (tickers) that he has requested: if response.success: try: for room in tickers_list: await self.enter_room(sid, room) except Exception as e: self.app_state.printer("EXCEPTION!", e) response.status_code = StatusCodes.FAILED response.message = "Failed to set new watchlist." # Done here: response_dict, response_http_code = response.for_quart() await self.emit("set_watchlist", response_dict, to = sid) async def send_ticks_to_clients(self, ticks: List[dict]) -> None: """ Here's where we decide which client gets which tick and send it out. WARNING: WE ARE ASSUMING THAT NO FURTHER FORMATING/COMPUTATION IS REQUIRED OTHER THAN SELECTING WHICH SUBSETS OF TICKS TO SEND TO WHICH CLIENTS. FOR US THE TICKS ALREADY HAVE ALL THE DATA NEEDED TO BE SEND TO RESPECTIVE CLIENTS. :param ticks: The list of individual tick updates to send out to the clients. :return: None """ # Currently we're just broadcasting # all the data to all the clients: await self.emit( event = self.app_state.EVENT_TICKS, data = ticks ) async def get_ticks_from_kafka( self, fetch_count: int = 100, fetch_timeout: float = 1.0 ) -> None: """ This function must run in the background forever and just keep listening for ticks on Kafka and keep relaying them to all the connected clients as per their watchlists. :param fetch_count: How many messages to consume in one go. :param fetch_timeout: How long to wait (in seconds) while consuming messages from Kafka. :return: None """ self.app_state.printer("Starting Kafka consumer (ticks).") # Do the next part infinitely: while True: # Note the time: now_utc = date_time.get_current_utc_date_time().timestamp() # Get messages form Kafka: ticks = await self.app_state.kafka_consumer.consume( count = fetch_count, timeout = fetch_timeout ) # If there are no updates to give: if not ticks: self.app_state.no_context_printer("No ticks.") continue # Each message must be treated as an array of tick updates (list of dicts). # In case the producer is sending each individual tick as a separate message, # we normalize it to be a list: tasks = [self.send_ticks_to_clients(t.value if isinstance(t.value, list) else [t.value]) for t in ticks] results = await asyncio.gather(*tasks) # Analyze the ticks: latency = [abs(now_utc - t.ts.timestamp()) for t in ticks] avg_latency = sum(latency) / len(latency) total_ticks = len(ticks) ticks_str = f"COUNT: {total_ticks: >5,} | AVG. LATENCY: {avg_latency:,.5f}" # self.app_state.no_context_printer(ticks_str) # ***************************************************************************************************************** # ***** **** # *** MAIN PROGRAM *** # ***** **** # ***************************************************************************************************************** if __name__ == "__main__": pass