Skip to content

Commit

Permalink
fix: websocket messages using flask were sometimes corrupted.
Browse files Browse the repository at this point in the history
It seems the underlying library wsproto is not thread safe.
Using a mutex to protect the websocket calls seems to fix the issue.
  • Loading branch information
maartenbreddels committed Sep 29, 2023
1 parent 348ab5a commit 9697470
Showing 1 changed file with 8 additions and 3 deletions.
11 changes: 8 additions & 3 deletions solara/server/flask.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@
import logging
import mimetypes
import os
import threading
from http.server import HTTPServer
from pathlib import Path
from typing import Any
Expand Down Expand Up @@ -55,15 +56,19 @@ class WebsocketWrapper(websocket.WebsocketWrapper):

def __init__(self, ws: simple_websocket.Server) -> None:
self.ws = ws
self.lock = threading.Lock()

def close(self):
self.ws.close()
with self.lock:
self.ws.close()

def send_text(self, data: str) -> None:
self.ws.send(data)
with self.lock:
self.ws.send(data)

def send_bytes(self, data: bytes) -> None:
self.ws.send(data)
with self.lock:
self.ws.send(data)

async def receive(self):
from anyio import to_thread
Expand Down

0 comments on commit 9697470

Please sign in to comment.