""" AUTHOR: Khushal P Soonderji DATE: Thursday, 2nd Jan., 2025 OBJECTIVE: To broadcast live tick updates to connected clients. It doesn't matter which stockbroker we are getting the ticks from as long as we are reading standardized ticks from the Kafka queue. REFERENCES: N/A DOWNLOADS: N/A """ # ***************************************************************************************************************** # ***** **** # *** IMPORT *** # ***** **** # ***************************************************************************************************************** # To make sibling directories accessible for imports: import sys sys.path.append(".") sys.path.append("..") # System-level activities: import io import os # My utils: from utils_v2.string import json from utils_v2.string import regex from utils_v2.system import files from utils_v2.database.async_mongo_v2 import AsyncMongo from utils_v2.queue.kafka.controllers.async_kafka import ConsumerKafka, get_ssl_context from utils_v2.cache.async_redis_cache_v2 import AsyncRedisCache from utils_v2.serialization.json_serializer import JSONSerializer # To make HTTP calls: import httpx # To work with date and time: import datetime import time # Models: from models.core.user import CoreUserInfoModel # To work with SocketIO: import socket import socketio # For asynchronous activities: import asyncio # To work with various datatypes: from typing import List # Debugging: from icecream import IceCreamDebugger # ***************************************************************************************************************** # ***** **** # *** MACROS / ONE-TIME INIT *** # ***** **** # ***************************************************************************************************************** # Debugging: printer = IceCreamDebugger(prefix = "Tick-Out | ", includeContext = True) printer.disable() # To make API calls: http_client = httpx.AsyncClient( limits = httpx.Limits( max_connections = 100, # ............ Maximum number of connections allowed in the pool. max_keepalive_connections = 50, # ... Maximum number of connections that can be kept alive. ), timeout = httpx.Timeout( pool = 120.0, # .... Time to wait for a free connection from the pool. connect = 2.5, # ... Time to wait for establishing a connection to the server. write = 10.0, # .... Time to wait for sending data. read = 9.9 # ....... Time to wait for receiving data. ) ) # General: SERVER_HOSTNAME = str(socket.gethostname()) # For SocketIO: # Namespaces: NAMESPACE_MODULE = None NAMESPACE_PASSTHROUGH = "/passthrough" # Events: EVENT_CONNECT = "connect" EVENT_DISCONNECT = "disconnect" EVENT_ECHO = "echo" EVENT_TICKS = "ticks" # ***************************************************************************************************************** # ***** **** # *** VARIABLES *** # ***** **** # ***************************************************************************************************************** # For SocketIO: ALLOWED_ORIGINS = [] sio = socketio.AsyncServer( cors_ALLOWED_ORIGINS = lambda x: any(regex.match(x, origin) for origin in ALLOWED_ORIGINS), async_mode = "asgi" ) app = socketio.ASGIApp(sio) # Redis: redis_cache: AsyncRedisCache | None = None # For kafka: kafka_consumer: ConsumerKafka | None = None # Session-awareness and maintenance of this script's state: SCRIPT_DATA = {} exclusive_lock = asyncio.Semaphore(1) CONNECTED_CLIENTS = {} FLAGS = { "initDone": False } # ***************************************************************************************************************** # ***** **** # *** FUNCTIONS *** # ***** **** # ***************************************************************************************************************** async def send_ticks(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 sio.emit( event = EVENT_TICKS, data = ticks, namespace = NAMESPACE_MODULE ) # --------------------------------------------------------------------------------------------------------------------- async def ticks_from_kafka( consumer: ConsumerKafka, 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 consumer: The preconfigured Kafka consumer that can listen for ticks in asynchronous mode. :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 """ printer("Starting Kafka consumer (ticks).") # Do the next part infinitely: while True: # Get messages form Kafka: ticks = await consumer.consume( count = fetch_count, timeout = fetch_timeout ) printer(len(ticks)) # If there are no updates to give: if not 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 = [send_ticks(t.value if isinstance(t.value, list) else [t.value]) for t in ticks] results = await asyncio.gather(*tasks) # --------------------------------------------------------------------------------------------------------------------- async def init( script_id: str, debug: bool ): """ 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. :return: True if initialized successfully, else False. """ # Declare the required global variables: global SCRIPT_DATA global ALLOWED_ORIGINS global redis_cache global kafka_consumer # Basic stuff: if debug: printer.enable() printer("Initializing.") # ┏┓ • • # ┃┃┏┓┓┏┓┓┏┓┏ # ┗┛┛ ┗┗┫┗┛┗┛ # ┛ response = await http_client.post( url = r"https://api.thecaoffice.com/ca/get/title", headers = {"Origin": "https://thecaoffice.com/"}, data = { "domainName": "127.0.0.1:1234", "screenWidth": 1920, "screenHeight": 1080 } ) 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", [])] printer(ALLOWED_ORIGINS) if len(ALLOWED_ORIGINS) < 1: print("FATAL: ALLOWED ORIGINS IS EMPTY!") return False # ┏┓ ┓ ┓ ┳┓ # ┃ ┏┓┏┓┏┫ ┏┓┏┓┏┫ ┃┃┏┓╋┏┓ # ┗┛┛ ┗ ┗┻ ┗┻┛┗┗┻ ┻┛┗┻┗┗┻ # Get the script credentials: response = await http_client.get( url = r"https://nexcom.ditscentre.in/internal/cred/get", headers = {"X-Script-Id": script_id} ) if response.status_code not in [200]: print("FATAL: SCRIPT CREDENTIALS LOADING FAILED!") return False script_cred = response.json().get("data") # Get the script data: response = await http_client.get( url = r"https://nexcom.ditscentre.in/internal/data/get", headers = {"X-Script-Id": script_id} ) if response.status_code not in [200]: print("FATAL: SCRIPT DATA LOADING FAILED!") return False SCRIPT_DATA = response.json().get("data") # Done with this step: printer("Cred and Data loaded.") # ┓┏┓ ┏┓ ┏┓┓• # ┃┫ ┏┓╋┃┏┏┓ ┃ ┃┓┏┓┏┓╋┏ # ┛┗┛┗┻┛┛┗┗┻ ┗┛┗┗┗ ┛┗┗┛ # Create the consumer that will listen to changes in watchlist: consumer_creds = script_cred["kafka"]["consumer"] kafka_consumer = ConsumerKafka( topic = consumer_creds["topic"], bootstrap_servers = consumer_creds["config"]["bootstrapServers"], security_protocol = consumer_creds["config"].get("securityProtocol", "PLAINTEXT"), ssl_context = get_ssl_context( ca_file = consumer_creds["config"].get("caFile"), cert_file = consumer_creds["config"].get("certFile"), key_file = consumer_creds["config"].get("keyFile"), ), serializer = JSONSerializer(), debug = debug ) if not await kafka_consumer.connect(): print("FATAL: KAFKA CONSUMER NOT CREATED!") return False printer("Kafka consumer ready.") # ┳┓ ┓• ┏┓ ┓ # ┣┫┏┓┏┫┓┏ ━━ ┃ ┏┓┏┣┓┏┓ # ┛┗┗ ┗┻┗┛ ┗┛┗┻┗┛┗┗ redis_cache = AsyncRedisCache( connection_string = script_cred["redisCache"]["general"]["connectionString"], serializer = JSONSerializer(), debug = debug, debug_prefix = "General Cache | " ) if not await redis_cache.connect(): print("FATAL: REDIS CACHE NOT CREATED!") return False printer("Redis cache ready.") # ┳┓ ┓ ┓ ┏┳┓ ┓ # ┣┫┏┓┏┃┏┏┓┏┓┏┓┓┏┏┓┏┫ ┃ ┏┓┏┃┏┏ # ┻┛┗┻┗┛┗┗┫┛ ┗┛┗┻┛┗┗┻ ┻ ┗┻┛┛┗┛ # ┛ # Start the background task that will receive ticks from the Kafka queue and broadcast them to the respective # connected clients: sio.start_background_task( ticks_from_kafka, consumer = kafka_consumer, fetch_count = 100, fetch_timeout = 1.0 ) # ┳┓ # ┃┃┏┓┏┓┏┓ # ┻┛┗┛┛┗┗ # If everything went well, we return with success: printer("Initialization done.") return True # --------------------------------------------------------------------------------------------------------------------- @sio.on(event = EVENT_CONNECT, namespace = NAMESPACE_MODULE) async def on_connect(sid, environ, *args) -> bool: """ The event handler for when a new connection request comes in. :param sid: The session id of the incoming request (generated by SocketIO). :param environ: The set of headers and other connection-specific values. :param args: Any extra input coming from the connection request. :return: True to accept a connection request, False to reject it. """ # declare the required global variables: global CONNECTED_CLIENTS # Initialize the script if needed: async with exclusive_lock: if not FLAGS.get("initDone"): FLAGS["initDone"] = await init( script_id = os.environ["SCRIPT_ID"], debug = True if os.environ["DEBUG"].lower() == "true" else False ) # If the initialization failed, we cannot accept the incoming request: if not FLAGS.get("initDone"): printer("SOCKET REJECTED: Init. pending.", sid) return False # Get the session token from the incoming request: session_token = environ.get("HTTP_X_SESSION_TOKEN") if not session_token and len(args) > 0: session_token = args[0].get("X-Session-Token") if not session_token: printer("SOCKET REJECTED: No session token.", sid) return False # Get the user's details from the session token: user_info = await redis_cache.get(key = session_token) if not user_info: 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"io_{session_token}" async with exclusive_lock: CONNECTED_CLIENTS[sid] = { "user": user_info, "redisKey": redis_key, "rooms": [] } sid_cached = await redis_cache.set( key = redis_key, value = {"server": SERVER_HOSTNAME, "socket_id": sid} ) # Done here: printer("SOCKET ACCEPTED.", sid, sid_cached) return True # --------------------------------------------------------------------------------------------------------------------- @sio.on(event = EVENT_DISCONNECT, namespace = NAMESPACE_MODULE) async def handle_disconnect(sid, reason) -> None: """ To handle a disconnect event. Automatically triggered when a client disconnects from the server. :param sid: The session id of the client (generated by SocketIO on connecting). :param reason: The hint about why the disconnection happened. :return: None. """ # declare the required global variables: global CONNECTED_CLIENTS # register the disconnect in the global variable, and on the cache server: client_info = {} async with exclusive_lock: client_info = CONNECTED_CLIENTS.pop(sid) sid_uncached = await redis_cache.delete(key = client_info["redisKey"]) printer("SOCKET DISCONNECTED", sid, reason, sid_uncached) # ***************************************************************************************************************** # ***** **** # *** MAIN PROGRAM *** # ***** **** # ***************************************************************************************************************** if __name__ == "__main__": # To get args from the terminal: import argparse # To run the ASGI: import uvicorn from multiprocessing import freeze_support # Get the config from the command-line: parser = argparse.ArgumentParser(description = f"SocketIO to serve live market data (and a general passthrough).") parser.add_argument( "-w", "--workers", type = int, help = "The no. of threads to spin up for this instance!", default = 2 ) parser.add_argument( "-a", "--host", type = str, help = "The host for the app. e.g.: '0.0.0.0' or '127.0.0.1'.", default = "127.0.0.1" ) parser.add_argument( "-p", "--port", type = int, help = "The port no. to bind the app to.", default = 8080 ) parser.add_argument( "-s", "--script-id", type = str, help = "The id of this script (will affect the loaded config)." ) parser.add_argument( "-d", "--debug", action = "store_true", help = "Whether, or not, you want to see debugging messages in the terminal.", default = False ) args = parser.parse_args() # Note down the config; os.environ["SCRIPT_ID"] = args.script_id os.environ["DEBUG"] = str(args.debug) # Startup message: message = "Debugging will be disabled by default." printer.enable() printer(message) printer.disable() print(message) # Run the gateway: freeze_support() uvicorn.run( app = "tick_out:app", workers = args.workers, host = args.host, port = args.port )