fix: wait for abortPendingRequest before inserting new request entry

Author: lwin-kyaw

packages/shield-controller/src/backend.ts

@@ -142,7 +142,7 @@ export class ShieldRemoteBackend implements ShieldBackend {
     };
 
     // clean up the pending coverage result polling
-    this.#pollingWithTimeout.abortPendingRequest(req.signatureRequest.id);
+    await this.#pollingWithTimeout.abortPendingRequest(req.signatureRequest.id);
 
     const res = await this.#fetch(
       `${this.#baseUrl}/v1/signature/coverage/log`,
@@ -166,7 +166,7 @@ export class ShieldRemoteBackend implements ShieldBackend {
     };
 
     // clean up the pending coverage result polling
-    this.#pollingWithTimeout.abortPendingRequest(req.txMeta.id);
+    await this.#pollingWithTimeout.abortPendingRequest(req.txMeta.id);
 
     const res = await this.#fetch(
       `${this.#baseUrl}/v1/transaction/coverage/log`,

packages/shield-controller/src/polling-with-timeout-abort.test.ts

@@ -99,7 +99,30 @@ describe('PollingWithTimeoutAndAbort', () => {
       fnName: 'test',
     });
     await delay(15); // small delay to let the request start
-    pollingWithTimeout.abortPendingRequest('test');
+    await pollingWithTimeout.abortPendingRequest('test');
     await expect(request).rejects.toThrow('test: Request cancelled');
   });
+
+  it('should get the result when the request succeeds', async () => {
+    const pollingWithTimeout = new PollingWithTimeoutAndAbort({
+      timeout: 1000,
+      pollInterval: 20,
+    });
+
+    const requestFn = jest
+      .fn()
+      .mockImplementation(async (_signal: AbortSignal) => {
+        return new Promise((resolve) => {
+          setTimeout(() => {
+            resolve('test result');
+          }, 100);
+        });
+      });
+
+    const request = pollingWithTimeout.pollRequest('test', requestFn, {
+      fnName: 'test',
+    });
+    const result = await request;
+    expect(result).toBe('test result');
+  });
 });

packages/shield-controller/src/polling-with-timeout-abort.ts

@@ -50,16 +50,17 @@ export class PollingWithTimeoutAndAbort {
     const pollInterval = pollingOptions.pollInterval ?? this.#pollInterval;
 
     // clean up the request entry if it exists
-    this.abortPendingRequest(requestId);
+    await this.abortPendingRequest(requestId);
 
     // insert the request entry for the next polling cycle
     const { abortController } = this.#insertRequestEntry(requestId, timeout);
 
     while (!abortController.signal.aborted) {
       try {
         const result = await requestFn(abortController.signal);
-        // polling success, we just need to clean up the request entry and return the result
-        this.#cleanUpOnFinished(requestId);
+        // polling success, we just need to abort the request and return the result
+        // note: this will trigger the abort handler, which will clean up the request entry
+        abortController.abort();
         return result;
       } catch {
         if (abortController.signal.aborted) {
@@ -82,19 +83,30 @@ export class PollingWithTimeoutAndAbort {
   }
 
   /**
-   * Abort the pending requests.
+   * Abort the pending request.
    * This will clean up the request entry if it exists, and abort the pending request if it exists.
+   * Note: this needs to be called before making a new request (with await) to avoid race conditions with the new request.
    *
    * @param requestId - The ID of the request to abort.
+   * @returns A promise that resolves when the request is aborted.
    */
-  abortPendingRequest(requestId: string) {
-    // firstly clean up the request entry if it exists
-    // note: this does not abort the request, it only cleans up the request entry for the next polling cycle
-    const existingEntry = this.#cleanUpRequestEntryIfExists(requestId);
-    // then abort the request if it exists
-    // note: this does abort the request, but it will not trigger the abort handler (hence, {@link cleanUpRequestEntryIfExists} will not be called)
-    // coz the AbortHandler event listener is already removed from the AbortSignal
-    existingEntry?.abortController.abort(this.ABORT_REASON_CANCELLED);
+  async abortPendingRequest(requestId: string) {
+    return new Promise((resolve) => {
+      // firstly clean up the request entry if it exists
+      // note: this does not abort the request, it only cleans up the request entry from the map.
+      const removedEntry = this.#cleanUp(requestId);
+      // if there's no existing entry, then we can just resolve the promise
+      if (!removedEntry) {
+        resolve(undefined);
+        return;
+      }
+      // otherwise, we will abort the existing/pending request
+      removedEntry.abortController.abort(this.ABORT_REASON_CANCELLED);
+      // we want to make sure that the request is actually aborted, before inserting the new request entry
+      if (removedEntry.abortController.signal.aborted) {
+        resolve(undefined);
+      }
+    });
   }
 
   /**
@@ -115,7 +127,7 @@ export class PollingWithTimeoutAndAbort {
 
     // Set the abort handler and listen to the `abort` event
     const abortHandler = () => {
-      this.#cleanUpOnFinished(requestId);
+      this.#cleanUp(requestId);
     };
     abortController.signal.addEventListener('abort', abortHandler);
 
@@ -132,35 +144,21 @@ export class PollingWithTimeoutAndAbort {
   }
 
   /**
-   * Clean up the request entry upon finished (success or failure).
-   * This will remove the abort handler from the AbortSignal, clear the timeout, and remove the request entry.
+   * Clean up the request entry.
+   * This will clear the timeout, remove the abort handler event listener, and remove the request entry from the map.
    *
-   * @param requestId - The ID of the request to clean up for.
-   * @returns The request entry that was cleaned up, if it exists.
+   * @param requestId - The ID of the request to clean up.
+   * @returns The request entry that was cleaned up.
    */
-  #cleanUpOnFinished(requestId: string): RequestEntry | undefined {
-    const requestEntry = this.#cleanUpRequestEntryIfExists(requestId);
+  #cleanUp(requestId: string) {
+    const requestEntry = this.#requestEntries.get(requestId);
     if (requestEntry) {
+      clearTimeout(requestEntry.timerId); // clear the timeout
       requestEntry.abortController.signal.removeEventListener(
         'abort',
         requestEntry.abortHandler,
-      );
-    }
-    return requestEntry;
-  }
-
-  /**
-   * Clean up the request entry if it exists.
-   * This will clear the pending timeout, remove the event listener from the AbortSignal, and remove the request entry.
-   *
-   * @param requestId - The ID of the request to handle the abort for.
-   * @returns The request entry that was aborted, if it exists.
-   */
-  #cleanUpRequestEntryIfExists(requestId: string): RequestEntry | undefined {
-    const requestEntry = this.#requestEntries.get(requestId);
-    if (requestEntry) {
-      clearTimeout(requestEntry.timerId); // Clear the timeout
-      this.#requestEntries.delete(requestId); // Remove the request entry
+      ); // clean up the abort handler event listener
+      this.#requestEntries.delete(requestId); // remove the request entry
     }
     return requestEntry;
   }