From 9fdb83f2a47d132bb5b41358d9adb4fc1c77e6a1 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Alexis=20M=C3=A9taireau?= <alexis@notmyidea.org>
Date: Fri, 25 Oct 2024 16:18:49 +0200
Subject: [PATCH] feat(sync): Check for websocket connectivity and retry if not
 connected.

Because the WebSockets do not offer a way to know when they failed to
connect, use a delay and check the state of the connection.
---
 umap/static/umap/js/modules/sync/engine.js    | 1 +
 umap/static/umap/js/modules/sync/websocket.js | 8 ++++++++
 umap/static/umap/js/umap.js                   | 6 ++++--
 umap/websocket_server.py                      | 5 ++++-
 4 files changed, 17 insertions(+), 3 deletions(-)

diff --git a/umap/static/umap/js/modules/sync/engine.js b/umap/static/umap/js/modules/sync/engine.js
index fce47910b..061a171a2 100644
--- a/umap/static/umap/js/modules/sync/engine.js
+++ b/umap/static/umap/js/modules/sync/engine.js
@@ -70,6 +70,7 @@ export class SyncEngine {
    * Authenticate with the server and start the transport layer.
    */
   async authenticate() {
+    console.log("authenticating")
     const [response, _, error] = await this._server.get(this._websocketTokenURI)
     if (!error) {
       this.start(response.token)
diff --git a/umap/static/umap/js/modules/sync/websocket.js b/umap/static/umap/js/modules/sync/websocket.js
index 37cd95e73..1748ebb13 100644
--- a/umap/static/umap/js/modules/sync/websocket.js
+++ b/umap/static/umap/js/modules/sync/websocket.js
@@ -1,5 +1,6 @@
 const PONG_TIMEOUT = 5000;
 const PING_INTERVAL = 30000;
+const FIRST_CONNECTION_TIMEOUT = 2000;
 
 export class WebSocketTransport {
   constructor(webSocketURI, authToken, messagesReceiver) {
@@ -21,6 +22,13 @@ export class WebSocketTransport {
       }
     }
 
+    this.ensureOpen = setInterval(() => {
+      if (this.websocket.readyState !== WebSocket.OPEN) {
+        this.websocket.close()
+        clearInterval(this.ensureOpen)
+      }
+    }, FIRST_CONNECTION_TIMEOUT)
+
     // To ensure the connection is still alive, we send ping and expect pong back.
     // Websocket provides a `ping` method to keep the connection alive, but it's
     // unfortunately not possible to access it from the WebSocket object.
diff --git a/umap/static/umap/js/umap.js b/umap/static/umap/js/umap.js
index f6582bfd6..f4fdd49d9 100644
--- a/umap/static/umap/js/umap.js
+++ b/umap/static/umap/js/umap.js
@@ -236,12 +236,14 @@ U.Map = L.Map.extend({
   },
 
   render: function (fields) {
+    console.log("sync.websocketConnected", this.sync.websocketConnected)
+    console.log("options.syncEnabled", this.options.syncEnabled)
     if (this.options.syncEnabled === true) {
-      if (!this.sync.websocketConnected) {
+      if (this.sync.websocketConnected !== true) {
         const template = `
         <h3><i class="icon icon-16"></i><span>${L._('Disconnected')}</span></h3>
         <p>
-        ${L._('This map has enabled real-time synchronization with other users, but you are currently disconnected. It will automatically reconnect when ready.')}
+        ${L._('This map has enabled real-time synchronization with other users, but you are currently disconnected.It will automatically reconnect when ready.')}
         </p>
         `
         this.dialog.open({
diff --git a/umap/websocket_server.py b/umap/websocket_server.py
index fb024c535..6483d6489 100644
--- a/umap/websocket_server.py
+++ b/umap/websocket_server.py
@@ -196,4 +196,7 @@ async def _serve():
             logging.debug(f"Waiting for connections on {host}:{port}")
             await asyncio.Future()  # run forever
 
-    asyncio.run(_serve())
+    try:
+        asyncio.run(_serve())
+    except KeyboardInterrupt:
+        print("Closing WebSocket server")