Keep track of reconnection atempts in delay calc

justuswilhelm · justuswilhelm · commit 56b609617207 · 2024-04-20T18:22:11.000+09:00
Track how many times Sarus has already tried to (re)connect.
When exponential backoff is enabled, use the stored number of connection
attempts to calculate the exponential delay.
diff --git a/__tests__/index/retryConnectionDelay.test.ts b/__tests__/index/retryConnectionDelay.test.ts
@@ -65,28 +65,85 @@ describe("retry connection delay", () => {
 });
 
 describe("Exponential backoff delay", () => {
-  it("will never be more than 8000 ms with rate set to 2", () => {
+  describe("with rate 2, backoffLimit 8000 ms", () => {
     // The initial delay shall be 1 s
     const initialDelay = 1000;
     const exponentialBackoffParams: ExponentialBackoffParams = {
       backoffRate: 2,
       // We put the ceiling at exactly 8000 ms
       backoffLimit: 8000,
     };
-    expect(
-      calculateRetryDelayFactor(exponentialBackoffParams, initialDelay, 0),
-    ).toBe(1000);
-    expect(
-      calculateRetryDelayFactor(exponentialBackoffParams, initialDelay, 1),
-    ).toBe(2000);
-    expect(
-      calculateRetryDelayFactor(exponentialBackoffParams, initialDelay, 2),
-    ).toBe(4000);
-    expect(
-      calculateRetryDelayFactor(exponentialBackoffParams, initialDelay, 3),
-    ).toBe(8000);
-    expect(
-      calculateRetryDelayFactor(exponentialBackoffParams, initialDelay, 4),
-    ).toBe(8000);
+    it("will never be more than 8000 ms with rate set to 2", () => {
+      expect(
+        calculateRetryDelayFactor(exponentialBackoffParams, initialDelay, 0),
+      ).toBe(1000);
+      expect(
+        calculateRetryDelayFactor(exponentialBackoffParams, initialDelay, 1),
+      ).toBe(2000);
+      expect(
+        calculateRetryDelayFactor(exponentialBackoffParams, initialDelay, 2),
+      ).toBe(4000);
+      expect(
+        calculateRetryDelayFactor(exponentialBackoffParams, initialDelay, 3),
+      ).toBe(8000);
+      expect(
+        calculateRetryDelayFactor(exponentialBackoffParams, initialDelay, 4),
+      ).toBe(8000);
+    });
+
+    it("should delay reconnection attempts exponentially", async () => {
+      const webSocketSpy = jest.spyOn(global, "WebSocket" as any);
+      webSocketSpy.mockImplementation(() => {});
+      const setTimeoutSpy = jest.spyOn(global, "setTimeout");
+      const sarus = new Sarus({
+        url,
+        exponentialBackoff: exponentialBackoffParams,
+      });
+      expect(sarus.state).toStrictEqual({
+        kind: "connecting",
+        failedConnectionAttempts: 0,
+      });
+      let instance: WebSocket;
+      [instance] = webSocketSpy.mock.instances;
+      if (!instance.onopen) {
+        throw new Error();
+      }
+      instance.onopen(new Event("open"));
+      if (!instance.onclose) {
+        throw new Error();
+      }
+      instance.onclose(new CloseEvent("close"));
+
+      let cb: Sarus["connect"];
+      // We iteratively call sarus.connect() and let it fail, seeing
+      // if it reaches 8000 as a delay and stays there
+      const attempts: [number, number][] = [
+        [1000, 1],
+        [2000, 2],
+        [4000, 3],
+        [8000, 4],
+        [8000, 5],
+      ];
+      attempts.forEach(([delay, failedAttempts]: [number, number]) => {
+        const call =
+          setTimeoutSpy.mock.calls[setTimeoutSpy.mock.calls.length - 1];
+        if (!call) {
+          throw new Error();
+        }
+        expect(call[1]).toBe(delay);
+        cb = call[0];
+        cb();
+        instance =
+          webSocketSpy.mock.instances[webSocketSpy.mock.instances.length - 1];
+        if (!instance.onclose) {
+          throw new Error();
+        }
+        instance.onclose(new CloseEvent("close"));
+        expect(sarus.state).toStrictEqual({
+          kind: "connecting",
+          failedConnectionAttempts: failedAttempts,
+        });
+      });
+    });
   });
 });
diff --git a/__tests__/index/state.test.ts b/__tests__/index/state.test.ts
@@ -15,7 +15,12 @@ describe("state machine", () => {
 
     // In the beginning, the state is "connecting"
     const sarus: Sarus = new Sarus(sarusConfig);
-    expect(sarus.state.kind).toBe("connecting");
+    // Since Sarus jumps into connecting directly, 1 connection attempt is made
+    // right in the beginning, but none have failed
+    expect(sarus.state).toStrictEqual({
+      kind: "connecting",
+      failedConnectionAttempts: 0,
+    });
 
     // We wait until we are connected, and see a "connected" state
     await server.connected;
@@ -24,14 +29,22 @@ describe("state machine", () => {
     // When the connection drops, the state will be "closed"
     server.close();
     await server.closed;
-    expect(sarus.state.kind).toBe("closed");
-
-    // Restart server
-    server = new WS(url);
+    expect(sarus.state).toStrictEqual({
+      kind: "closed",
+      failedConnectionAttempts: 0,
+    });
 
     // We wait a while, and the status is "connecting" again
     await delay(1);
-    expect(sarus.state.kind).toBe("connecting");
+    // In the beginning, no connection attempts have been made, since in the
+    // case of a closed connection, we wait a bit until we try to connect again.
+    expect(sarus.state).toStrictEqual({
+      kind: "connecting",
+      failedConnectionAttempts: 0,
+    });
+
+    // We restart the server and let the Sarus instance reconnect:
+    server = new WS(url);
 
     // When we connect in our mock server, we are "connected" again
     await server.connected;
diff --git a/src/index.ts b/src/index.ts
@@ -179,10 +179,16 @@ export default class Sarus {
    * after the constructor wraps up.
    */
   state:
-    | { kind: "connecting" }
+    | { kind: "connecting"; failedConnectionAttempts: number }
     | { kind: "connected" }
     | { kind: "disconnected" }
-    | { kind: "closed" } = { kind: "connecting" };
+    /**
+     * The closed state carries of the number of failed connection attempts
+     */
+    | { kind: "closed"; failedConnectionAttempts: number } = {
+    kind: "connecting",
+    failedConnectionAttempts: 0,
+  };
 
   constructor(props: SarusClassParams) {
     // Extract the properties that are passed to the class
@@ -377,7 +383,19 @@ export default class Sarus {
    * Connects the WebSocket client, and attaches event listeners
    */
   connect() {
-    this.state = { kind: "connecting" };
+    if (this.state.kind === "closed") {
+      this.state = {
+        kind: "connecting",
+        failedConnectionAttempts: this.state.failedConnectionAttempts,
+      };
+    } else if (
+      this.state.kind === "connected" ||
+      this.state.kind === "disconnected"
+    ) {
+      this.state = { kind: "connecting", failedConnectionAttempts: 0 };
+    } else {
+      // This is a NOOP, we are already connecting
+    }
     this.ws = new WebSocket(this.url, this.protocols);
     this.setBinaryType();
     this.attachEventListeners();
@@ -391,11 +409,22 @@ export default class Sarus {
   reconnect() {
     const self = this;
     const { retryConnectionDelay, exponentialBackoff } = self;
+    // If we are already in a "connecting" state, we need to refer to the
+    // current amount of connection attemps to correctly calculate the
+    // exponential delay -- if exponential backoff is enabled.
+    const failedConnectionAttempts =
+      self.state.kind === "connecting"
+        ? self.state.failedConnectionAttempts
+        : 0;
 
     // If no exponential backoff is enabled, retryConnectionDelay will
     // be scaled by a factor of 1 and it will stay the original value.
     const delay = exponentialBackoff
-      ? calculateRetryDelayFactor(exponentialBackoff, retryConnectionDelay, 0)
+      ? calculateRetryDelayFactor(
+          exponentialBackoff,
+          retryConnectionDelay,
+          failedConnectionAttempts,
+        )
       : retryConnectionDelay;
 
     setTimeout(self.connect, delay);
@@ -543,7 +572,22 @@ export default class Sarus {
         if (eventName === "open") {
           self.state = { kind: "connected" };
         } else if (eventName === "close" && self.reconnectAutomatically) {
-          self.state = { kind: "closed" };
+          const { state } = self;
+          // If we have previously been "connecting", we carry over the amount
+          // of failed connection attempts and add 1, since the current
+          // connection attempt failed. We stay "connecting" instead of
+          // "closed", since we've never been fully "connected" in the first
+          // place.
+          if (state.kind === "connecting") {
+            self.state = {
+              kind: "connecting",
+              failedConnectionAttempts: state.failedConnectionAttempts + 1,
+            };
+          } else {
+            // If we were in a different state, we assume that our connection
+            // freshly closed and have not made any failed connection attempts.
+            self.state = { kind: "closed", failedConnectionAttempts: 0 };
+          }
           self.removeEventListeners();
           self.reconnect();
         }
