Merge pull request #6 from ChasmNetwork/cha-127

handshake retries
ChasmNetwork · Aug 28, 2024 · 6587b95 · 6587b95
2 parents 12a5520 + 82fada5
commit 6587b95
Show file tree

Hide file tree

Showing 2 changed files with 102 additions and 52 deletions.
diff --git a/src/server/express.ts b/src/server/express.ts
@@ -53,12 +53,57 @@ app.get("/", (req, res) => {
   return res.status(200).send("OK");
 });
 
+async function retryWithExponentialBackoff(
+  retries: number,
+  fn: () => Promise<unknown>,
+): Promise<void> {
+  const initialDelay = 1000; // 1s start
+  let delayFactor = 30; // Initial growth factor
+
+  let previousDelay = initialDelay;
+  const delays = [initialDelay];
+
+  for (let attempt = 1; attempt < retries; attempt++) {
+    let currentDelay = previousDelay * delayFactor;
+
+    while (currentDelay <= previousDelay) {
+      delayFactor += 0.5;
+      currentDelay = previousDelay * delayFactor;
+    }
+
+    // 30 mins cap
+    if (currentDelay > 30 * 60 * 1000) {
+      currentDelay = 30 * 60 * 1000;
+    }
+
+    delays.push(currentDelay);
+    previousDelay = currentDelay;
+  }
+
+  logger.info(`Delays: ${delays.map((d) => d / 1000).join(", ")} seconds`);
+
+  for (let attempt = 0; attempt < retries; attempt++) {
+    try {
+      await fn();
+      return;
+    } catch (err: any) {
+      logger.error(`Attempt ${attempt + 1} failed: ${err.message}`);
+      if (attempt < retries - 1) {
+        const delay = delays[attempt];
+        logger.info(`Retrying in ${delay / 1000} seconds...`);
+        await new Promise((resolve) => setTimeout(resolve, delay));
+      } else {
+        logger.error("All retry attempts failed. Exiting gracefully.");
+        process.exit(1);
+      }
+    }
+  }
+}
+
 app.listen(PORT, async () => {
   logger.info(`Server running on port ${PORT}`);
-  await handshake().catch((err) => {
-    logger.error(err.message);
-    process.exit(1);
-  });
+  logger.info("Running exponential backoff handshake...");
+  await retryWithExponentialBackoff(6, handshake);
 });
 
 export default app;
diff --git a/src/server/handshake.ts b/src/server/handshake.ts
@@ -27,58 +27,63 @@ export async function handshake() {
   logger.debug("Connecting to orchestrator at " + orchestrator);
   logger.debug("Your webhook URL: " + WEBHOOK_URL);
 
-  const ws = new WebSocket(orchestrator, {
-    headers: {
-      Authorization: `Bearer ${WEBHOOK_API_KEY}`,
-    },
-  });
+  return new Promise((resolve, reject) => {
+    const ws = new WebSocket(orchestrator, {
+      headers: {
+        Authorization: `Bearer ${WEBHOOK_API_KEY}`,
+      },
+    });
 
-  ws.on("open", function open() {
-    logger.debug("Initiated WS handshake");
+    ws.on("open", function open() {
+      logger.debug("Initiated WS handshake");
 
-    logger.info(`Setting up webhook at: ${WEBHOOK_URL}`);
-    ws.send(
-      JSON.stringify({
-        type: "webhook",
-        webhookUrl: WEBHOOK_URL,
-        ip: IP,
-        version: getVersion(),
-        name: SCOUT_NAME,
-        providers: PROVIDERS,
-        model: MODEL,
-      }),
-    );
-  });
+      logger.info(`Setting up webhook at: ${WEBHOOK_URL}`);
+      ws.send(
+        JSON.stringify({
+          type: "webhook",
+          webhookUrl: WEBHOOK_URL,
+          ip: IP,
+          version: getVersion(),
+          name: SCOUT_NAME,
+          providers: PROVIDERS,
+          model: MODEL,
+        }),
+      );
+    });
 
-  ws.on("error", (error) => {
-    throw new Error(
-      `Handshake failed: ${error.message}\n${JSON.stringify(error, null, 2)}`,
-    );
-  });
+    ws.on("error", (error) => {
+      logger.error(`Handshake failed: ${error.message}`);
+      reject(error);
+    });
 
-  ws.on("close", (code, reason) => {
-    if (code !== 1000) {
-      throw new Error(
-        `❌ Handshake with orchestrator at ${ORCHESTRATOR_URL} failed with error:\n${reason}`,
-      );
-    }
-    logger.debug(`WS connection closed by the server with ${code}`);
-  });
+    ws.on("close", (code, reason) => {
+      if (code !== 1000) {
+        logger.error(
+          `❌ Handshake with orchestrator at ${ORCHESTRATOR_URL} failed with error:\n${reason}`,
+        );
+        reject(new Error(reason.toString()));
+      } else {
+        logger.debug(`WS connection closed by the server with ${code}`);
+        resolve(null);
+      }
+    });
 
-  ws.on("message", function incoming(data) {
-    const { success } = JSON.parse(data.toString());
-    if (success) {
-      logger.info(
-        chalk.green.bold(
-          `✅ Handshake with orchestrator at ${ORCHESTRATOR_URL} complete`,
-        ),
-      );
-      ws.close();
-      return;
-    } else {
-      throw new Error(
-        "❌ Handshake with orchestrator failed. No response from server.",
-      );
-    }
+    ws.on("message", function incoming(data) {
+      const { success } = JSON.parse(data.toString());
+      if (success) {
+        logger.info(
+          chalk.green.bold(
+            `✅ Handshake with orchestrator at ${ORCHESTRATOR_URL} complete`,
+          ),
+        );
+        ws.close();
+        resolve(null);
+      } else {
+        logger.error(
+          "❌ Handshake with orchestrator failed. No response from server.",
+        );
+        reject(new Error("No response from server"));
+      }
+    });
   });
 }