Package: fastapi-distributed-websocket
FastAPI Distributed WebsocketA library to implement websocket for distibuted systems based on FastAPI.N.B.: This library is still at an early stage, use it in production at your own risk.What it doesThe main features of this libarary are:Easly implementing broadcasting, pub/sub, chat rooms, etc...Proxy websocket connections to other servers (e.g. from an api gateway)AuthenticationClean exception handlingAn in memory broker for fast developmentProblems of scaling websocket among multiple servers in productionWebsocket is a relatively new protocol for real time communication over HTTP.
It establishes a durable, stateful, full-duplex connection between clients and the server.
It can be used to implement chats, real time notifications, broadcasting and
pub/sub models.Connections from clientsHTTP request/response mechanism fits very well to scale among multiple server
instances in production. Any time a client makes a request, it can connect
to any server instance and it's going to receive the same response. After
the response has been returned to the client, it went disconnected and it can
make another request without the need to hit the same instace as before.
This thanks to the stateless nature of HTTP.However, Websocket establishes a stateful connection between the client and the
server and, if some error occurs and the connection went lost, we have to
ensure that clients are going to hit the same server instance they were connected
before, since that instance was managing the connection state.Stateful means that there is a state that can be manipulated. In particular,
a stateful connection is a connection that heavily relies on its state in
order to workBroadcasting and group messagesAnother problem of scaling Websocket occurs when we need to send messages to
multiple connected clients (i.e. broadcasting a message or sending a message to
all clients subscribed to a specific topic).Imagine that we have a chat server, and that when an user send a message in a
specific chat room, we broadcast it to all users subscribed to that room.
If we have a single server instance, all connection are managed by this instance
so we can safely trust that the message will be delivered to all recipents.
On the other hand, with multiple server instances, users subscribing to a chat
room will probably connect to different instances. This way, if an user send a
message to the chat room'xyz'at the serverA, users subscribed to the same
chat room at the serverBare not receiving it.Documenting Websocket interfacesAnother common problem with Websocket, that's not even related to scaling, is
about documentation. Due to the event driven nature of the Websocket protocol
it does not fit well to be documented withopenapi.
However a new specification for asynchronous, event driven interfaces has been
defined recently. The spec name isasyncapiand I'm
currently studying it. I don't know if this has to be implemented here or it's
better having a separate library for that, however this is surely something
we have to look at.Other problemsWhen I came first to think about this library, I started making a lot of research
of common problems related to Websocket on stackoverflow, reddit, github issues and
so on. I found some interesting resource that are however related to the implementation
itself. I picked up best solutions and elaborated my owns converging all of that in
this library.ExamplesInstallation$ pip install fastapi-distributed-websocketBasic usageThis is a basic example using an in memory broker with a single server instance.fromfastapiimportFastAPI,WebSocket,WebSocketDisconnect,statusfromdistributed_websocketimportConnection,WebSocketManagerapp=FastAPI()manager=WebSocketManager('channel:1',broker_url='memory://')...app.on_event('startup')asyncdefstartup()->None:...awaitmanager.startup()app.on_event('shutdown')asyncdefshutdown()->None:...awaitmanager.shutdown()@app.websocket('/ws/{conn_id}')asyncdefwebsocket_endpoint(ws:WebSocket,conn_id:str,*,topic:Optional[Any]=None,)->None:connection:Connection=awaitmanager.new_connection(ws,conn_id)try:whileTrue:msg=awaitconnection.receive_json()awaitmanager.broadcast(msg)exceptWebSocketDisconnect:awaitmanager.remove_connection(connection)Themanger.new_connectionmethod create a new Connection object and add it to
themanager.active_connectionslist. Note that after aWebSocketDisconnectis raised, we callremove_connection: this method only remove the connection
object from themanager.active_connectionslist, without callingconnection.close, since
the connection is already closed.
If you need to close a connection at any other time, you can usemanager.close_connection.
If you useconnection.iter_json, it already handles theWebSocketDisconnectexception, so
you can simply callmanager.remove_connectionjust after the loop (see next code block).Note that here we are usingmanager.broadcastto send the message to all connections managed
by the WebSocketManager instance. However, this method only work if we have a single server
instance. If we have multiple server instances, we have to usemanager.receive, to properly
send the message to the broker.@app.websocket('/ws/{conn_id}')asyncdefwebsocket_endpoint(ws:WebSocket,conn_id:str,*,topic:Optional[Any]=None,)->None:connection:Connection=awaitmanager.new_connection(ws,conn_id)# This is the preferred way of handling WebSocketDisconnectasyncformsginconnection.iter_json():awaitmanager.receive(connection,msg)awaitmanager.remove_connection(connection)Proxy from an API gatewayLet's say we are developing a chat service and that all our services are behind
an API gateway. If we want to keep our websocket service behind it too, then
fastapi-distributed-websocket provides us withWebSocketProxy.fromdistributed_websocketimportWebSocketProxy# skipped other imports for brevityapp=FastAPI()WS_TARGET_ENDPOINT='ws://websocket_service:8000/wshandler'@app.websocket('/ws')asyncdefwebsocket_proxy(websocket:WebSocket):awaitwebsocket.accept()ws_proxy=WebSocketProxy(websocket,WS_TARGET_ENDPOINT)awaitws_proxy()This will forward all messages from the client to the target endpoint and
all messages from the target endpoint to the client.Now let's assume that our websocket service code is the code of our previous
example. Our API Gateway code will be:fromdistributed_websocketimportWebSocketProxy# skipped other imports for brevityapp=FastAPI()WS_TARGET_ENDPOINT='ws://websocket_service:8000/ws/{}'@app.websocket('/ws/{conn_id}')asyncdefwebsocket_endpoint(ws:WebSocket,conn_id:str,)->None:awaitwebsocket.accept()ws_proxy=WebSocketProxy(websocket,WS_TARGET_ENDPOINT.format(conn_id))awaitws_proxy()API ReferenceConnectionConnection objects wrap the websocket connection and provide a simple interface
to send and receive messages. They have atopicsattribute to store subscriptions
patterns and implement pub/sub models.asyncaccept(self) -> NoneAccept the connection.asyncclose(self, code: int = 1000) -> NoneClose the connection with the specified status.asyncreceive_json(self) -> AnyReceive a JSON message.asyncsend_json(self, data: Any) -> NoneSend a JSON message over the connection.asynciter_json(self) -> AsyncIterator[Any]Iterate over the messages received over the connection.MessagesMessage objects store the message type, the topic and the data. They provides
an easy serialization/deserialization mechanism.
Remeber that messages returned byconnection.iter_jsonare already deserialized
intodictobjects, so here we calldeserializationthe process of converting
adictobject into aMessageobject.type: strThe message type.topic: strThe message topic.conn_id: str | list[str]The connection id or list of connection ids that the message should be sent to.data: AnyThe message data.classmethodfrom_client_message(cls, *, data: Any) -> MessageCreate a message from a client message.__serialize__(self) -> dictSerialize the message into adictobject.SubscriptionsYou can bind topics to connection objects to implement pub/sub models, notification and so on.
Thetopicsattribute is a set of strings that follows the pattern matching syntax of MQTT.
This library share connection objects state between server instances, so you may find
references to terms likechannel,publish,subscribeandunsubscribereferring to
the same concepts but applied to the underlying server/broker communication.This may be confusing, but remember to keep separated the communication between the server
and the clients, that you are developing and the communication between the server and the broker,
that you usually don't deal with.subscribe(connection: Connection, message: Message) -> NoneSubscribe a connection tomessage.topic.unsubscribe(connection: Connection, message: Message) -> NoneUnsubscribe a connection frommessage.topic.hanlde_subscription_message(connection: Connection, message: Message) -> NoneCallssubscribeorunsubscribedepending on the message type.matches(topic: str, patterns: set[str]) -> boolCheck iftopicmatches any of the patterns inpatterns.AuthenticationAuthentication is provided with theWebSocketOAuth2PasswordBearerclass.
It inherits fromFastAPIOAuth2PasswordBearerand overrides__call__method to accept
aWebSocketobject.async__call__(self, websocket: WebSocket) -> str | NoneAuthenticate the websocket connection and return theAuthorizationheader value.If the authentication fails, returnNoneif the objects has been initialized withauto_error=Falseor close the connection with theWS_1008_POLICY_VIOLATIONcode.Exceptions and Exception Handlingfastapi-distributed-websocketprovides exception handling via decorators. You can use the
apposite decorators passing an exception class and a handler callable. Exception handlers
should accept only the exception object as argument.Why this is useful?Because sometimes the same type of exception can be raised by different parts of the application,
this way you can decorate the higer level function in the call stack to handle the exception at
any level.A baseWebSocketExceptionclass is provided to bind connection objects to the exception, so
your handler function can easily access it.
If you need to access connection objects from the exception handler, your custom exceptions
should inherit fromWebSocketException, no matter if they are really network related or not.WebSocketException(self, message: str, *, connection: Connection) -> NoneInvalidSubscription(self, message: str, *, connection: Connection) -> NoneRaised when a subscription pattern use an invalid syntax. Inherits fromWebSocketException.InvalidSubscriptionMessage(self, message: str, *, connection: Connection) -> NoneLikeInvalidSubscriptionit could be raised for bad syntax, but it could also be raisedwhen the message type is notsubscribeorunsubscribe. Inherits fromWebSocketException.handle(exc: BaseException, handler: Callable[..., Any]) -> Callable[..., Any]Decorator to handle exceptions. If you decorate a function with this decorator, at any timean exception of typeexcis raised or propagated to the function, it will be handled byhandler.Use this decorator only if both your handler and the function are not async.asyncahandle( exc: BaseException, handler: Callable[..., Coroutine[Any, Any, Any]] ) -> Callable[..., Any]Decorator to handle exceptions, same adhandle, but the handler is a coroutine function.Use this if your handler is a coroutine function, while the decorated function could beeither a sync or an async function.Broker InterfacesConnections' state is shared between server instances using a pub/sub broker. By default,
the broker is areids.asyncio.Redisinstance (exaioredis.Redis), but you can use any
other implementation.fastapi-distributed-websocketprovides anInMemoryBrokerclass
for development purposes.
You can inherit fromBrokerInterfaceand override the methods to implement your own broker.asyncconnect(self) -> Coroutine[Any, Any, None]Connect to the broker.asyncdisconnect(self) -> Coroutine[Any, Any, None]Disconnect from the broker.asyncsubscribe(self, channel: str) -> Coroutine[Any, Any, None]Subscribe to a channel.asyncunsubscribe(self, channel: str) -> Coroutine[Any, Any, None]Unsubscribe from a channel.asyncpublish(self, channel: str, message: Any) -> Coroutine[Any, Any, None]Publish a message to a channel.asyncget_message(self, **kwargs) -> Coroutine[Any, Any, Message | None]Get a message from the broker.WebSocketManagerTheWebSocketManagerclass is where the main logic of the library is implemented.It keeps track of the connection objects and starts the broker connection.
It spawn a main task, a listener that wait (non-blocking) for messages from the broker,
and send them to the connection objects (broadcasting or checking for subscriptions)
spawning a new task for each send.The broker initialisation is done in the constructor while calls tobroker.connectandbroker.disconnectare handled in thestartupandshutdownmethods.asyncnew_connection( self, websocket: WebSocket, conn_id: str, topic: str | None = None ) -> Coroutine[Any, Any, Connection]Create a new connection object, add it toself.active_connectionsand return it.asyncclose_connection( self, connection: Connection, code: int = status.WS_1000_NORMAL_CLOSURE ) -> Coroutine[Any, Any, None]Close a connection object and remove it fromself.active_connections.remove_connection(self, connection: Connection) -> NoneRemove a connection object fromself.active_connections.set_conn_id(self, connection: Connection, conn_id: str) -> NoneSet the connection id and notify the client.send(self, topic: str, message: Any) -> NoneSend a message to all the connection objects subscribed totopic.It spawns a new task wrapping the coroutine resulting fromself._send.broadcast(self, message: Any) -> NoneSend a message to all the connection objects.It spawns a new task wrapping the coroutine resulting fromself._broadcast.send_by_conn_id(self, conn_id: str | list[str], message: Any) -> NoneSend a message to all the connection objects with the given connection id.It spawns a new task wrapping the coroutine resulting fromself._send_by_conn_idifconn_idis a string or from_send_multi_by_conn_idif it is a list.send_msg(self, message: Message) -> NoneBased on the message type, it callssend,send_by_conn_idorbroadcast.asyncreceive( self, connection: Connection, message: Any ) -> Coroutine[Any, Any, None]Receive a message from a connection object. It passes the message down toa private method that handle eventual subscriptions and then publish the messageto the broker.asyncstartup(self) -> Coroutine[Any, Any, None]Start the broker connection and the listener task.asyncshutdown(self) -> Coroutine[Any, Any, None]Close the broker connection and the listener task.It also takes care to cancel all the tasks spawned bysendandbroadcastandclose all the connection objects before.WebSocketProxyTheWebSocketProxyclass initialise callable objects that can be
used to start proxyng websocket messages from client to a server and viceversa.
It's initialised with a two parameters:client: aWebSocketobjectserver_endpoint: astrcontaining the endpoint of the serverNotice that the target server could be a remote server or the same server that starts the proxy.async__call__(self) -> Coroutine[Any, Any, None]Start a websocket connection toserver_endpointand spawn two tasks:one that forwards the messages from the client to the target and the other thatforwards the messages from the target to the client.