341 lines
13 KiB
Python
341 lines
13 KiB
Python
"""
|
|
|
|
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
|