Cleaned up timeouts

maneesht · maneesht · commit ac3c2b4f911a · 2022-10-07T16:52:02.000-07:00
diff --git a/packages/storage/src/task.ts b/packages/storage/src/task.ts
@@ -94,6 +94,8 @@ export class UploadTask {
   private _metadataErrorHandler: (p1: StorageError) => void;
   private _resolve?: (p1: UploadTaskSnapshot) => void = undefined;
   private _reject?: (p1: StorageError) => void = undefined;
+  /* eslint-disable-next-line @typescript-eslint/no-explicit-any */ // This has to be an "any" type due to incompatibility between Node and browser timers.
+  private pendingTimeout: any = null;
   private _promise: Promise<UploadTaskSnapshot>;
 
   private sleepTime: number;
@@ -191,8 +193,16 @@ export class UploadTask {
             // Happens if we miss the metadata on upload completion.
             this._fetchMetadata();
           } else {
-            setTimeout(() => {
-              this._continueUpload();
+            this.pendingTimeout = setTimeout(() => {
+              this.pendingTimeout = null;
+              switch (this._state) {
+                case InternalTaskState.CANCELING:
+                  this._transition(InternalTaskState.CANCELED);
+                case InternalTaskState.PAUSING:
+                  this._transition(InternalTaskState.PAUSED);
+                default:
+                  this._continueUpload();
+              }
             }, this.sleepTime);
           }
         }
@@ -411,6 +421,10 @@ export class UploadTask {
         this._state = state;
         if (this._request !== undefined) {
           this._request.cancel();
+        } else if (this.pendingTimeout) {
+          clearTimeout(this.pendingTimeout);
+          this.pendingTimeout = null;
+          this.completeTransitions_();
         }
         break;
       case InternalTaskState.PAUSING:
@@ -419,6 +433,10 @@ export class UploadTask {
         this._state = state;
         if (this._request !== undefined) {
           this._request.cancel();
+        } else if (this.pendingTimeout) {
+          clearTimeout(this.pendingTimeout);
+          this.pendingTimeout = null;
+          this.completeTransitions_();
         }
         break;
       case InternalTaskState.RUNNING:
diff --git a/packages/storage/test/unit/task.test.ts b/packages/storage/test/unit/task.test.ts
@@ -14,7 +14,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-import { assert, expect } from 'chai';
+import { assert, expect, use } from 'chai';
 import { FbsBlob } from '../../src/implementation/blob';
 import { Location } from '../../src/implementation/location';
 import { Unsubscribe } from '../../src/implementation/observer';
@@ -31,8 +31,12 @@ import {
 } from './testshared';
 import { injectTestConnection } from '../../src/platform/connection';
 import { Deferred } from '@firebase/util';
-import { retryLimitExceeded } from '../../src/implementation/error';
+import { canceled, retryLimitExceeded } from '../../src/implementation/error';
 import { SinonFakeTimers, useFakeTimers } from 'sinon';
+import * as sinon from 'sinon';
+import sinonChai from 'sinon-chai';
+
+use(sinonChai);
 
 const testLocation = new Location('bucket', 'object');
 const smallBlob = new FbsBlob(new Uint8Array([97]));
@@ -361,7 +365,7 @@ describe('Firebase Storage > Upload Task', () => {
   function handleStateChange(
     requestHandler: RequestHandler,
     blob: FbsBlob
-  ): Promise<TotalState> {
+  ): { promise: Promise<TotalState>; task: UploadTask } {
     const storageService = storageServiceWithHandler(requestHandler);
     const task = new UploadTask(
       new Reference(storageService, testLocation),
@@ -410,7 +414,7 @@ describe('Firebase Storage > Upload Task', () => {
       }
     );
 
-    return deferred.promise;
+    return { promise: deferred.promise, task };
   }
 
   it('Calls callback sequences for small uploads correctly', () => {
@@ -422,7 +426,7 @@ describe('Firebase Storage > Upload Task', () => {
   it('properly times out if large blobs returns a 503 when finalizing', async () => {
     clock = useFakeTimers();
     // Kick off upload
-    const promise = handleStateChange(
+    const { promise } = handleStateChange(
       fake503ForFinalizeServerHandler(),
       bigBlob
     );
@@ -460,7 +464,10 @@ describe('Firebase Storage > Upload Task', () => {
   it('properly times out if large blobs returns a 503 when uploading', async () => {
     clock = useFakeTimers();
     // Kick off upload
-    const promise = handleStateChange(fake503ForUploadServerHandler(), bigBlob);
+    const { promise } = handleStateChange(
+      fake503ForUploadServerHandler(),
+      bigBlob
+    );
     // Run all timers
     await clock.runAllAsync();
     const { events, progress } = await promise;
@@ -478,10 +485,112 @@ describe('Firebase Storage > Upload Task', () => {
     });
     clock.restore();
   });
+  function resumeCancelSetup(): any {
+    clock = useFakeTimers();
+    const fakeSetTimeout = clock.setTimeout;
+    let gotFirstEvent = false;
+
+    const stub = sinon.stub(global, 'setTimeout');
+
+    const timers: { [key: number]: any } = {};
+    const readyToCancel = new Promise(resolve => {
+      stub.callsFake((fn, timeout) => {
+        // @ts-ignore
+        const res = fakeSetTimeout(fn, timeout);
+        //@ts-ignore
+        timers[res.id] = {
+          data: res,
+          timeout
+        };
+        if ((!gotFirstEvent || timeout === 0) && timeout !== 600000) {
+          clock.tick(timeout as number);
+        } else if (timeout !== 600000) {
+          resolve(null);
+        }
+        return res;
+      });
+    });
+    readyToCancel.then(
+      () => stub.restore(),
+      () => {}
+    );
+    return {
+      ...handleStateChange(
+        fake503ForUploadServerHandler(undefined, () => (gotFirstEvent = true)),
+        bigBlob
+      ),
+      readyToCancel
+    };
+  }
+  it('properly errors with a cancel StorageError if a pending timeout remains', async () => {
+    // Kick off upload
+    const { readyToCancel, promise, task } = resumeCancelSetup();
+
+    await readyToCancel;
+    task.cancel();
+
+    // Run all timers
+    const { events, progress } = await promise;
+    expect(events.length).to.equal(2);
+    expect(events[0]).to.deep.equal({ type: 'resume' });
+    expect(events[1].type).to.deep.equal('error');
+    const canceledError = canceled();
+    expect(events[1].data!.name).to.deep.equal(canceledError.name);
+    expect(events[1].data!.message).to.deep.equal(canceledError.message);
+    const blobSize = bigBlob.size();
+    expect(progress.length).to.equal(1);
+    expect(progress[0]).to.deep.equal({
+      bytesTransferred: 0,
+      totalBytes: blobSize
+    });
+    expect(clock.countTimers()).to.eq(0);
+    clock.restore();
+  });
+  it('properly errors with a pause StorageError if a pending timeout remains', async () => {
+    // Kick off upload
+    const { readyToCancel, promise, task } = resumeCancelSetup();
+
+    await readyToCancel;
+
+    task.pause();
+    expect(clock.countTimers()).to.eq(0);
+    task.resume();
+    await clock.runAllAsync();
+
+    // Run all timers
+    const { events, progress } = await promise;
+    expect(events.length).to.equal(4);
+    expect(events[0]).to.deep.equal({ type: 'resume' });
+    expect(events[1]).to.deep.equal({ type: 'pause' });
+    expect(events[2]).to.deep.equal({ type: 'resume' });
+    expect(events[3].type).to.deep.equal('error');
+    const retryError = retryLimitExceeded();
+    expect(events[3].data!.name).to.deep.equal(retryError.name);
+    expect(events[3].data!.message).to.deep.equal(retryError.message);
+    const blobSize = bigBlob.size();
+    expect(progress.length).to.equal(3);
+    expect(progress[0]).to.deep.equal({
+      bytesTransferred: 0,
+      totalBytes: blobSize
+    });
+    expect(progress[1]).to.deep.equal({
+      bytesTransferred: 0,
+      totalBytes: blobSize
+    });
+    expect(progress[2]).to.deep.equal({
+      bytesTransferred: 0,
+      totalBytes: blobSize
+    });
+    expect(clock.countTimers()).to.eq(0);
+    clock.restore();
+  });
   it('tests if small requests that respond with 500 retry correctly', async () => {
     clock = useFakeTimers();
     // Kick off upload
-    const promise = handleStateChange(fakeOneShot503ServerHandler(), smallBlob);
+    const { promise } = handleStateChange(
+      fakeOneShot503ServerHandler(),
+      smallBlob
+    );
     // Run all timers
     await clock.runAllAsync();
     const { events, progress } = await promise;
diff --git a/packages/storage/test/unit/testshared.ts b/packages/storage/test/unit/testshared.ts
@@ -356,7 +356,7 @@ export function fakeServerHandler(
 
 /**
  * Responds with a 503 for finalize.
- * @param fakeMetadata metadata to respond with for query
+ * @param fakeMetadata metadata to respond with for finalize
  * @returns a handler for requests
  */
 export function fake503ForFinalizeServerHandler(
@@ -459,7 +459,8 @@ export function fake503ForFinalizeServerHandler(
  * @returns a handler for requests
  */
 export function fake503ForUploadServerHandler(
-  fakeMetadata: Partial<Metadata> = defaultFakeMetadata
+  fakeMetadata: Partial<Metadata> = defaultFakeMetadata,
+  cb?: any
 ): RequestHandler {
   const stats: {
     [num: number]: {
@@ -536,6 +537,9 @@ export function fake503ForUploadServerHandler(
     const isUpload = commands.indexOf('upload') !== -1;
 
     if (isUpload) {
+      if (cb) {
+        cb();
+      }
       return {
         status: 503,
         body: JSON.stringify(fakeMetadata),
