I am try to get message from when task is done from Celery and push message to redis but issue is when i try to broadcast message on websocket they are close or partial active which make it revert and many solutions said used ping function which send empty message all the time to client until he task is done but it not make sense.
below is code example
import asyncio
import redis.asyncio as aioredis
import json
from fastapi import WebSocket
from fastapi.middleware.cors import CORSMiddleware
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
import json
import argparse
from typing import Dict, Any
import enum
from fastapi import APIRouter
from app.api.utils.verifyJWTToken import verify_jwt_token
from fastapi.exceptions import HTTPException
import threading
app = FastAPI()
web_socket_router = APIRouter()
class WebSocketState(enum.Enum):
CONNECTING = 0
CONNECTED = 1
DISCONNECTED = 2
RESPONSE = 3
class RedisPubSubManager:
"""
Initializes the RedisPubSubManager.
Args:
host (str): Redis server host.
port (int): Redis server port.
"""
def __init__(self, host="localhost", port=6379):
self.redis_host = host
self.redis_port = port
self.pubsub = None
async def _get_redis_connection(self) -> aioredis.Redis:
"""
Establishes a connection to Redis.
Returns:
aioredis.Redis: Redis connection object.
"""
return aioredis.Redis(
host=self.redis_host, port=self.redis_port, auto_close_connection_pool=False
)
async def connect(self) -> None:
"""
Connects to the Redis server and initializes the pubsub client.
"""
self.redis_connection = await self._get_redis_connection()
self.pubsub = self.redis_connection.pubsub()
async def _publish(self, chat_id: str, message: str) -> None:
"""
Publishes a message to a specific Redis channel.
Args:
chat_id (str): Channel or chat ID.
message (str): Message to be published.
"""
await self.redis_connection.publish(chat_id, message)
async def subscribe(self, chat_id: str) -> aioredis.Redis:
"""
Subscribes to a Redis channel.
Args:
chat_id (str): Channel or chat ID to subscribe to.
Returns:
aioredis.ChannelSubscribe: PubSub object for the subscribed channel.
"""
await self.pubsub.subscribe(chat_id)
return self.pubsub
async def unsubscribe(self, chat_id: str) -> None:
"""
Unsubscribes from a Redis channel.
Args:
chat_id (str): Channel or chat ID to unsubscribe from.
"""
await self.pubsub.unsubscribe(chat_id)
class WebSocketManager:
_instance = None
_lock = threading.Lock()
def __new__(cls, *args, **kwargs):
with cls._lock:
if not cls._instance:
cls._instance = super(WebSocketManager, cls).__new__(cls)
cls._instance.chats = {}
cls._instance.users = {}
cls._instance.pubsub_client = RedisPubSubManager()
return cls._instance
def __init__(self):
"""
Initializes the WebSocketManager.
Attributes:
chats (dict): A dictionary to store WebSocket connections in different chats.
pubsub_client (RedisPubSubManager): An instance of the RedisPubSubManager class for pub-sub functionality.
"""
# self.chats: dict = {}
# self.users: Dict[str, Any] = {}
# self.pubsub_client = RedisPubSubManager()
async def connect(self, websocket: WebSocket):
await websocket.accept()
async def set_client_web_socket(self, user_id, websocket):
self.users[user_id] = websocket
print("send_data::user_id",self.users)
# await self.users[user_id].send_text('{"message":"connected"}')
await self.send_data(user_id,'{"message":"connected"}')
async def send_data(self, user_id, message: str):
try:
websocket = self.users.get(user_id)
print("websocket.application_state",websocket.application_state)
# if websocket and not websocket.application_state == "closed":
await websocket.send_text(message)
# else:
# print("WebSocket is closed or does not exist.")
except Exception as e:
print(f"Failed to send message: {e}")
async def ping(self, websocket: WebSocket):
try:
# if websocket and not websocket.application_state == "closed":
await websocket.send_text("nil")
# else:
# print("WebSocket is closed or does not exist.")
except Exception as e:
print(f"Failed to send message: {e}")
def get_client_web_socket(self, user_id):
return self.users[user_id]
async def disconnect(self, message,code, websocket: WebSocket):
await websocket.send_text(message)
await websocket.close(code)
async def add_user_id_to_chat(self, chat_id: str, user_id:str) -> None:
"""
Adds a user's WebSocket connection to a chat.
Args:
chat_id (str): chat ID or channel name.
websocket (WebSocket): WebSocket connection object.
"""
websocket= self.get_client_web_socket(user_id)
if chat_id in self.chats:
self.chats[chat_id].append(websocket)
else:
self.chats[chat_id] = [websocket]
await self.pubsub_client.connect()
pubsub_subscriber = await self.pubsub_client.subscribe(chat_id)
print("self.chats : ",self.chats)
message = {
"user_id": user_id,
"chat_id": chat_id,
"message": f"User {user_id} connected to chat - {chat_id}",
}
await socket_manager.send_data(chat_id, json.dumps(message))
asyncio.create_task(self._pubsub_data_reader(pubsub_subscriber))
async def add_user_to_chat(self, chat_id: str, websocket: WebSocket) -> None:
"""
Adds a user's WebSocket connection to a chat.
Args:
chat_id (str): chat ID or channel name.
websocket (WebSocket): WebSocket connection object.
"""
await websocket.accept()
if chat_id in self.chats:
self.chats[chat_id].append(websocket)
else:
self.chats[chat_id] = [websocket]
await self.pubsub_client.connect()
pubsub_subscriber = await self.pubsub_client.subscribe(chat_id)
asyncio.create_task(self._pubsub_data_reader(pubsub_subscriber))
async def broadcast_to_chat(self, chat_id: str, message: str) -> None:
"""
Broadcasts a message to all connected WebSockets in a chat.
Args:
chat_id (str): chat ID or channel name.
message (str): Message to be broadcasted.
"""
await self.pubsub_client._publish(chat_id, message)
async def remove_user_from_chat(self, chat_id: str, websocket: WebSocket) -> None:
"""
Removes a user's WebSocket connection from a chat.
Args:
chat_id (str): chat ID or channel name.
websocket (WebSocket): WebSocket connection object.
"""
self.chats[chat_id].remove(websocket)
if len(self.chats[chat_id]) == 0:
del self.chats[chat_id]
await self.pubsub_client.unsubscribe(chat_id)
async def _pubsub_data_reader(self, pubsub_subscriber):
"""
Reads and broadcasts messages received from Redis PubSub.
Args:
pubsub_subscriber (aioredis.ChannelSubscribe): PubSub object for the subscribed channel.
"""
while True:
message = await pubsub_subscriber.get_message(
ignore_subscribe_messages=True
)
if message is not None:
# print("message : ",message)
chat_id = message["channel"].decode("utf-8")
print("message : ",chat_id)
all_sockets = self.chats[chat_id]
print("message : ",all_sockets)
for socket in all_sockets:
data = message["data"].decode("utf-8")
print("message : ",socket.application_state)
await socket.send_text(json.dumps(data))
# try:
# if socket.application_state == WebSocketState.CONNECTED:
# await socket.send_text(json.dumps(data))
# except RuntimeError as e:
# print("Failed to send message, likely due to closed WebSocket:", str(e))
# await socket_manager.disconnect(json.dumps(f'{"error: ",str(e)}'), socket)
socket_manager = WebSocketManager()
@web_socket_router.websocket("/connect")
async def websocket_endpoint(websocket: WebSocket):
await websocket.accept()
token = websocket.headers.get("Authorization")
error_message = ""
if token:
try:
token_data = await verify_jwt_token(token)
# print(token_data) # Use token_data as needed
except HTTPException as http_exc:
error_message = http_exc.detail
else:
error_message = "Authorization token is missing"
if error_message:
await socket_manager.disconnect(error_message, 4001, websocket)
return
try:
user_id = token_data["user"]["_id"]
print("user_id",user_id)
await socket_manager.set_client_web_socket(user_id, websocket)
while True:
await socket_manager.ping(websocket)
except WebSocketDisconnect:
print("disconnect")
await asyncio.sleep(3)
await socket_manager.disconnect("Error while receive message", websocket)
# await socket_manager.broadcast(f"Client #{user_id} left the chat")
@web_socket_router.websocket("/send")
async def websocket_endpoint(websocket: WebSocket):
await socket_manager.connect(websocket)
token = websocket.headers.get("Authorization")
error_message = ""
if token:
try:
token_data = await verify_jwt_token(token)
# print(token_data) # Use token_data as needed
except HTTPException as http_exc:
error_message = http_exc.detail
else:
error_message = "Authorization token is missing"
if error_message:
await socket_manager.disconnect(error_message, 4001, websocket)
return
try:
user_id = token_data["user"]["_id"]
print("user_id",user_id)
await socket_manager.send_data(user_id,'{"message":"new message"}')
except WebSocketDisconnect:
print("disconnect")
await socket_manager.disconnect("Error while receive message", websocket)
# await socket_manager.broadcast(f"Client #{user_id} left the chat")
@web_socket_router.websocket("/api/v1/ws/{chat_id}/{user_id}")
async def websocket_endpoint(websocket: WebSocket, chat_id: str, user_id: int):
await socket_manager.add_user_to_chat(chat_id, websocket)
message = {
"user_id": user_id,
"chat_id": chat_id,
"message": f"User {user_id} connected to chat - {chat_id}",
}
await socket_manager.broadcast_to_chat(chat_id, json.dumps(message))
try:
while True:
data = await websocket.receive_text()
message = {"user_id": user_id, "chat_id": chat_id, "message": data}
await socket_manager.broadcast_to_chat(chat_id, json.dumps(message))
except WebSocketDisconnect:
await socket_manager.remove_user_from_chat(chat_id, websocket)
message = {
"user_id": user_id,
"chat_id": chat_id,
"message": f"User {user_id} disconnected from chat - {chat_id}",
}
await socket_manager.broadcast_to_chat(chat_id, json.dumps(message))
@web_socket_router.websocket("/api/v2/ws/{chat_id}/{user_id}")
async def websocket_endpoint(websocket: WebSocket, chat_id: str, user_id: int):
# await socket_manager.add_user_to_chat(chat_id, websocket)
message = {
"user_id": user_id,
"chat_id": chat_id,
"message": f"User {user_id} connected to chat - {chat_id}",
}
await socket_manager.broadcast_to_chat(chat_id, json.dumps(message))
try:
while True:
data = await websocket.receive_text()
message = {"user_id": user_id, "chat_id": chat_id, "message": data}
await socket_manager.broadcast_to_chat(chat_id, json.dumps(message))
except WebSocketDisconnect:
await socket_manager.remove_user_from_chat(chat_id, websocket)
message = {
"user_id": user_id,
"chat_id": chat_id,
"message": f"User {user_id} disconnected from chat - {chat_id}",
}
await socket_manager.broadcast_to_chat(chat_id, json.dumps(message))
# Register the WebSocket router
app.include_router(web_socket_router)