diff --git a/playground/sio_client.py b/playground/sio_client.py index a2068a9..316507e 100644 --- a/playground/sio_client.py +++ b/playground/sio_client.py @@ -1,15 +1,36 @@ import socketio import time + sio = socketio.Client() + +@sio.on(event = "echo", namespace = "/echo") +def on_echo_echo(data): + print("ECHO-ECHO:", data) + +@sio.on(event = "echo", namespace = "/ticks") +def on_ticks_echo(data): + print("TICKS-ECHO:", data) + +@sio.on(event = "ticks", namespace = "/ticks") +def on_ticks(data): + print("TICKS-TICKS:", data) + +namespaces = ["/echo", "/ticks"] sio.connect( url = "http://localhost:5214", - namespaces = ["/test"] + namespaces = namespaces, + headers = { + "Origin": "api.thecaoffice.com", + "X-Session-Token": "1d194a63-7a8f-44e1-b656-1d3016921ece" + } ) + while True: - sio.emit( - event = "echo", - data = {"k0": "v0", "k1": "v1"}, - namespace = "/test" - ) - print("Emitted!") - time.sleep(1) \ No newline at end of file + for ns in namespaces: + sio.emit( + event = "echo", + data = {"k0": "v0", "k1": "v1"}, + namespace = ns + ) + print("Emitted!") + time.sleep(1) \ No newline at end of file diff --git a/wsio_v2/app.py b/wsio_v2/app.py index 7b7281e..0b4dd30 100644 --- a/wsio_v2/app.py +++ b/wsio_v2/app.py @@ -49,6 +49,7 @@ import socketio # To maintain the app's state: from wsio_v2.app_state import AppState from wsio_v2.test.echo import EchoNamespace +from wsio_v2.finstitutions.trading.tick_out import TicksNamespace # Debugging: from icecream import IceCreamDebugger @@ -83,8 +84,24 @@ app_state.http_client = httpx.AsyncClient( ) # General: +app_state.SCRIPT_ID = None +app_state.debug = True app_state.SERVER_HOSTNAME = str(socket.gethostname()) app_state.ALLOWED_ORIGINS = [] +app_state.connected_clients = {} + +# Namespaces: +app_state.NAMESPACE_DEFAULT = None +app_state.NAMESPACE_ECHO = "/echo" +app_state.NAMESPACE_TICKS = "/ticks" +app_state.NAMESPACE_ORDER_UPDATES = "/order-updates" + +# Events: +app_state.EVENT_CONNECT = "connect" +app_state.EVENT_DISCONNECT = "disconnect" +app_state.EVENT_ECHO = "echo" +app_state.EVENT_TICKS = "ticks" +app_state.EVENT_ORDER_UPDATE = "order_update" # ***************************************************************************************************************** @@ -102,7 +119,8 @@ sio = socketio.AsyncServer( app = socketio.ASGIApp(sio) # Register the namespaces: -sio.register_namespace(EchoNamespace(namespace = "/test", app_state = app_state)) +sio.register_namespace(EchoNamespace(namespace = app_state.NAMESPACE_ECHO, app_state = app_state)) +sio.register_namespace(TicksNamespace(namespace = app_state.NAMESPACE_TICKS, app_state = app_state)) # ***************************************************************************************************************** @@ -172,8 +190,9 @@ if __name__ == "__main__": # Startup message: app_state.printer.enable() - app_state.printer(str(args.debug)) - if str(args.debug).lower().find("false") >= 0: app_state.printer.disable() + if args.debug != "True": + app_state.printer("Disabling debug print.") + app_state.printer.disable() # Run the gateway: freeze_support() diff --git a/wsio_v2/finstitutions/trading/tick_out.py b/wsio_v2/finstitutions/trading/tick_out.py index bef477b..a4a2355 100644 --- a/wsio_v2/finstitutions/trading/tick_out.py +++ b/wsio_v2/finstitutions/trading/tick_out.py @@ -32,19 +32,37 @@ # 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 + # ***************************************************************************************************************** # ***** **** @@ -73,7 +91,7 @@ from utils_v2.api.response import ResponseModel # ***************************************************************************************************************** -class EchoNamespace(socketio.AsyncNamespace): +class TicksNamespace(socketio.AsyncNamespace): def __init__( self, @@ -105,6 +123,69 @@ class EchoNamespace(socketio.AsyncNamespace): 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): """ @@ -117,6 +198,12 @@ class EchoNamespace(socketio.AsyncNamespace): 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): """ @@ -129,7 +216,7 @@ class EchoNamespace(socketio.AsyncNamespace): 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) -> dict: + async def on_set_watchlist(self, sid, data) -> None: """ To handle the request to set the watchlist for a particular user. @@ -176,6 +263,70 @@ class EchoNamespace(socketio.AsyncNamespace): 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) + # ***************************************************************************************************************** # ***** **** diff --git a/wsio_v2/helpers/helpers.py b/wsio_v2/helpers/helpers.py index 0c2232b..cbbc280 100644 --- a/wsio_v2/helpers/helpers.py +++ b/wsio_v2/helpers/helpers.py @@ -35,6 +35,9 @@ import sys sys.path.append(".") sys.path.append("..") +# For system-level activities: +import os + # My utils: from utils_v2.string import regex from utils_v2.queue.kafka.controllers.async_kafka import ConsumerKafka, get_ssl_context @@ -104,20 +107,18 @@ def origin_is_allowed(origin: str, app_state: AppState) -> bool: async def init( - script_id: str, - debug: bool, app_state: AppState, ): """ To initialize all credentials, instances, and connectivity for this whole script. - :param script_id: The id to use to load cred and data from the internal service. - :param debug: Whether, or not, you would like to print the debug messages. + :param app_state: The app's state with all the shared variables. :return: True if initialized successfully, else False. """ # Basic stuff: - if debug: app_state.printer.enable() + app_state.SCRIPT_ID = os.environ["SCRIPT_ID"] + app_state.debug = True if str(os.environ["DEBUG"]).lower().find("true") >= 0 else False app_state.printer("Initializing.") # ┏┓ • • @@ -137,9 +138,9 @@ async def init( if response.status_code not in [200]: print("FATAL: ALLOWED ORIGINS NOT FETCHED!") return False - ALLOWED_ORIGINS = [origin["domain"] for origin in response.json().get("data", {}).get("rs2", [])] - app_state.printer(ALLOWED_ORIGINS) - if len(ALLOWED_ORIGINS) < 1: + app_state.ALLOWED_ORIGINS = [origin["domain"] for origin in response.json().get("data", {}).get("rs2", [])] + app_state.printer(app_state.ALLOWED_ORIGINS) + if len(app_state.ALLOWED_ORIGINS) < 1: print("FATAL: ALLOWED ORIGINS IS EMPTY!") return False @@ -150,7 +151,7 @@ async def init( # Get the script credentials: response = await app_state.http_client.get( url = r"https://nexcom.ditscentre.in/internal/cred/get", - headers = {"X-Script-Id": script_id} + headers = {"X-Script-Id": app_state.SCRIPT_ID} ) if response.status_code not in [200]: print("FATAL: SCRIPT CREDENTIALS LOADING FAILED!") @@ -160,7 +161,7 @@ async def init( # Get the script data: response = await app_state.http_client.get( url = r"https://nexcom.ditscentre.in/internal/data/get", - headers = {"X-Script-Id": script_id} + headers = {"X-Script-Id": app_state.SCRIPT_ID} ) if response.status_code not in [200]: print("FATAL: SCRIPT DATA LOADING FAILED!") @@ -186,7 +187,7 @@ async def init( key_file = consumer_creds["config"].get("keyFile"), ), serializer = JSONSerializer(), - debug = debug + debug = app_state.debug ) if not await app_state.kafka_consumer.connect(): print("FATAL: KAFKA CONSUMER NOT CREATED!") @@ -198,10 +199,10 @@ async def init( # ┛┗┗ ┗┻┗┛ ┗┛┗┻┗┛┗┗ app_state.redis_cache = AsyncRedisCache( - connection_string = script_cred["redisCache"]["general"]["sentinelJson"], - # connection_string = script_cred["redisCache"]["general"]["connectionString"], + # connection_string = script_cred["redisCache"]["general"]["sentinelJson"], + connection_string = script_cred["redisCache"]["general"]["connectionString"], # serializer = JSONSerializer(), - debug = debug, + debug = app_state.debug, debug_prefix = "General Cache | " ) if not await app_state.redis_cache.connect():