(202501001) New Socket.IO system started with provision for namespaces and more.
This commit is contained in:
@@ -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)
|
||||
|
||||
|
||||
# *****************************************************************************************************************
|
||||
# ***** ****
|
||||
|
||||
Reference in New Issue
Block a user