Merge pull request #73 from JmPotato/add_stop_websocket

Add graceful shutdown logic for WebSocket

Author: lmlmt

hyperliquid/info.py

@@ -25,6 +25,8 @@ def __init__(
         if not skip_ws:
             self.ws_manager = WebsocketManager(self.base_url)
             self.ws_manager.start()
+        else:
+            self.ws_manager = None
         if meta is None:
             meta = self.meta()
 
@@ -43,6 +45,12 @@ def __init__(
             if name not in self.name_to_coin:
                 self.name_to_coin[name] = spot_info["name"]
 
+    def disconnect_websocket(self):
+        if self.ws_manager is None:
+            raise RuntimeError("Cannot call disconnect_websocket since skip_ws was used")
+        else:
+            self.ws_manager.stop()
+
     def user_state(self, address: str) -> Any:
         """Retrieve trading details about a user.
 

hyperliquid/websocket_manager.py

@@ -73,16 +73,25 @@ def __init__(self, base_url):
         ws_url = "ws" + base_url[len("http") :] + "/ws"
         self.ws = websocket.WebSocketApp(ws_url, on_message=self.on_message, on_open=self.on_open)
         self.ping_sender = threading.Thread(target=self.send_ping)
+        self.stop_event = threading.Event()
 
     def run(self):
-        self.ping_sender.start()
         self.ws.run_forever()
+        self.ping_sender.start()
 
     def send_ping(self):
-        while True:
-            time.sleep(50)
+        while not self.stop_event.wait(50):
+            if not self.ws.keep_running:
+                break
             logging.debug("Websocket sending ping")
             self.ws.send(json.dumps({"method": "ping"}))
+        logging.debug("Websocket ping sender stopped")
+
+    def stop(self):
+        self.stop_event.set()
+        self.ws.close()
+        if self.ping_sender.is_alive():
+            self.ping_sender.join()
 
     def on_message(self, _ws, message):
         if message == "Websocket connection established.":