Implemented exponential backoff with query

Author: maneesht

packages/storage/src/implementation/error.ts

@@ -35,7 +35,7 @@ export class StorageError extends FirebaseError {
    *  added to the end of the message.
    * @param message  - Error message.
    */
-  constructor(code: StorageErrorCode, message: string) {
+  constructor(code: StorageErrorCode, message: string, private status_ = 0) {
     super(
       prependCode(code),
       `Firebase Storage: ${message} (${prependCode(code)})`
@@ -46,6 +46,14 @@ export class StorageError extends FirebaseError {
     Object.setPrototypeOf(this, StorageError.prototype);
   }
 
+  get status(): number {
+    return this.status_;
+  }
+
+  set status(status: number) {
+    this.status_ = status;
+  }
+
   /**
    * Compares a StorageErrorCode against this error's code, filtering out the prefix.
    */

packages/storage/src/implementation/request.ts

@@ -26,6 +26,7 @@ import { ErrorHandler, RequestHandler, RequestInfo } from './requestinfo';
 import { isJustDef } from './type';
 import { makeQueryString } from './url';
 import { Connection, ErrorCode, Headers, ConnectionType } from './connection';
+import { isRetryStatusCode_ } from './utils';
 
 export interface Request<T> {
   getPromise(): Promise<T>;
@@ -69,7 +70,8 @@ class NetworkRequest<I extends ConnectionType, O> implements Request<O> {
     private errorCallback_: ErrorHandler | null,
     private timeout_: number,
     private progressCallback_: ((p1: number, p2: number) => void) | null,
-    private connectionFactory_: () => Connection<I>
+    private connectionFactory_: () => Connection<I>,
+    private retry = true
   ) {
     this.promise_ = new Promise((resolve, reject) => {
       this.resolve_ = resolve as (value?: O | PromiseLike<O>) => void;
@@ -93,16 +95,15 @@ class NetworkRequest<I extends ConnectionType, O> implements Request<O> {
       const connection = this.connectionFactory_();
       this.pendingConnection_ = connection;
 
-      const progressListener: (progressEvent: ProgressEvent) => void =
-        progressEvent => {
-          const loaded = progressEvent.loaded;
-          const total = progressEvent.lengthComputable
-            ? progressEvent.total
-            : -1;
-          if (this.progressCallback_ !== null) {
-            this.progressCallback_(loaded, total);
-          }
-        };
+      const progressListener: (
+        progressEvent: ProgressEvent
+      ) => void = progressEvent => {
+        const loaded = progressEvent.loaded;
+        const total = progressEvent.lengthComputable ? progressEvent.total : -1;
+        if (this.progressCallback_ !== null) {
+          this.progressCallback_(loaded, total);
+        }
+      };
       if (this.progressCallback_ !== null) {
         connection.addUploadProgressListener(progressListener);
       }
@@ -118,7 +119,11 @@ class NetworkRequest<I extends ConnectionType, O> implements Request<O> {
           this.pendingConnection_ = null;
           const hitServer = connection.getErrorCode() === ErrorCode.NO_ERROR;
           const status = connection.getStatus();
-          if (!hitServer || this.isRetryStatusCode_(status)) {
+          if (
+            (!hitServer ||
+              isRetryStatusCode_(status, this.additionalRetryCodes_)) &&
+            this.retry
+          ) {
             const wasCanceled = connection.getErrorCode() === ErrorCode.ABORT;
             backoffCallback(
               false,
@@ -196,22 +201,6 @@ class NetworkRequest<I extends ConnectionType, O> implements Request<O> {
       this.pendingConnection_.abort();
     }
   }
-
-  private isRetryStatusCode_(status: number): boolean {
-    // The codes for which to retry came from this page:
-    // https://cloud.google.com/storage/docs/exponential-backoff
-    const isFiveHundredCode = status >= 500 && status < 600;
-    const extraRetryCodes = [
-      // Request Timeout: web server didn't receive full request in time.
-      408,
-      // Too Many Requests: you're getting rate-limited, basically.
-      429
-    ];
-    const isExtraRetryCode = extraRetryCodes.indexOf(status) !== -1;
-    const isRequestSpecificRetryCode =
-      this.additionalRetryCodes_.indexOf(status) !== -1;
-    return isFiveHundredCode || isExtraRetryCode || isRequestSpecificRetryCode;
-  }
 }
 
 /**
@@ -271,7 +260,8 @@ export function makeRequest<I extends ConnectionType, O>(
   authToken: string | null,
   appCheckToken: string | null,
   requestFactory: () => Connection<I>,
-  firebaseVersion?: string
+  firebaseVersion?: string,
+  retry = true
 ): Request<O> {
   const queryPart = makeQueryString(requestInfo.urlParams);
   const url = requestInfo.url + queryPart;
@@ -291,6 +281,7 @@ export function makeRequest<I extends ConnectionType, O>(
     requestInfo.errorHandler,
     requestInfo.timeout,
     requestInfo.progressCallback,
-    requestFactory
+    requestFactory,
+    retry
   );
 }

packages/storage/src/implementation/requests.ts

@@ -104,7 +104,7 @@ export function sharedErrorHandler(
     xhr: Connection<ConnectionType>,
     err: StorageError
   ): StorageError {
-    let newErr;
+    let newErr: StorageError;
     if (xhr.getStatus() === 401) {
       if (
         // This exact message string is the only consistent part of the
@@ -126,6 +126,7 @@ export function sharedErrorHandler(
         }
       }
     }
+    newErr.status = xhr.getStatus();
     newErr.serverResponse = err.serverResponse;
     return newErr;
   }
@@ -534,8 +535,15 @@ export function continueResumableUpload(
   }
   const startByte = status_.current;
   const endByte = startByte + bytesToUpload;
-  const uploadCommand =
-    bytesToUpload === bytesLeft ? 'upload, finalize' : 'upload';
+  let uploadCommand = '';
+  if (bytesToUpload === 0) {
+    // TODO(mtewani): Maybe we should extract this out.
+    uploadCommand = 'finalize';
+  } else if (bytesLeft === bytesToUpload) {
+    uploadCommand = 'upload, finalize';
+  } else {
+    uploadCommand = 'upload';
+  }
   const headers = {
     'X-Goog-Upload-Command': uploadCommand,
     'X-Goog-Upload-Offset': `${status_.current}`

packages/storage/src/implementation/utils.ts

@@ -0,0 +1,35 @@
+/**
+ * @license
+ * Copyright 2022 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+export function isRetryStatusCode_(
+  status: number,
+  additionalRetryCodes: number[]
+): boolean {
+  // The codes for which to retry came from this page:
+  // https://cloud.google.com/storage/docs/exponential-backoff
+  const isFiveHundredCode = status >= 500 && status < 600;
+  const extraRetryCodes = [
+    // Request Timeout: web server didn't receive full request in time.
+    408,
+    // Too Many Requests: you're getting rate-limited, basically.
+    429
+  ];
+  const isExtraRetryCode = extraRetryCodes.indexOf(status) !== -1;
+  const isRequestSpecificRetryCode =
+    additionalRetryCodes.indexOf(status) !== -1;
+  return isFiveHundredCode || isExtraRetryCode || isRequestSpecificRetryCode;
+}

packages/storage/src/service.ts

@@ -302,7 +302,8 @@ export class FirebaseStorageImpl implements FirebaseStorage {
     requestInfo: RequestInfo<I, O>,
     requestFactory: () => Connection<I>,
     authToken: string | null,
-    appCheckToken: string | null
+    appCheckToken: string | null,
+    retry = true
   ): Request<O> {
     if (!this._deleted) {
       const request = makeRequest(
@@ -311,7 +312,8 @@ export class FirebaseStorageImpl implements FirebaseStorage {
         authToken,
         appCheckToken,
         requestFactory,
-        this._firebaseVersion
+        this._firebaseVersion,
+        retry
       );
       this._requests.add(request);
       // Request removes itself from set when complete.

packages/storage/src/task.ts

@@ -22,7 +22,8 @@ import { FbsBlob } from './implementation/blob';
 import {
   canceled,
   StorageErrorCode,
-  StorageError
+  StorageError,
+  retryLimitExceeded
 } from './implementation/error';
 import {
   InternalTaskState,
@@ -53,6 +54,7 @@ import {
 } from './implementation/requests';
 import { Reference } from './reference';
 import { newTextConnection } from './platform/connection';
+import { isRetryStatusCode_ } from './implementation/utils';
 
 /**
  * Represents a blob being uploaded. Can be used to pause/resume/cancel the
@@ -92,6 +94,15 @@ export class UploadTask {
   private _reject?: (p1: StorageError) => void = undefined;
   private _promise: Promise<UploadTaskSnapshot>;
 
+  // TODO(mtewani): Update these to use predefined constants
+  private sleepTime = 0;
+
+  private maxSleepTime = 10000;
+
+  isExponentialBackoffExpired(): boolean {
+    return this.sleepTime > this.maxSleepTime;
+  }
+
   /**
    * @param ref - The firebaseStorage.Reference object this task came
    *     from, untyped to avoid cyclic dependencies.
@@ -111,6 +122,17 @@ export class UploadTask {
         this._needToFetchStatus = true;
         this.completeTransitions_();
       } else {
+        const backoffExpired = this.isExponentialBackoffExpired();
+        if (isRetryStatusCode_(error.status, [])) {
+          if (backoffExpired) {
+            error = retryLimitExceeded();
+          } else {
+            this.sleepTime = Math.max(this.sleepTime * 2, 1000);
+            this._needToFetchStatus = true;
+            this.completeTransitions_();
+            return;
+          }
+        }
         this._error = error;
         this._transition(InternalTaskState.ERROR);
       }
@@ -163,7 +185,9 @@ export class UploadTask {
             // Happens if we miss the metadata on upload completion.
             this._fetchMetadata();
           } else {
-            this._continueUpload();
+            setTimeout(() => {
+              this._continueUpload();
+            }, this.sleepTime);
           }
         }
       }
@@ -283,7 +307,8 @@ export class UploadTask {
         requestInfo,
         newTextConnection,
         authToken,
-        appCheckToken
+        appCheckToken,
+        false
       );
       this._request = uploadRequest;
       uploadRequest.getPromise().then((newStatus: ResumableUploadStatus) => {
@@ -344,7 +369,8 @@ export class UploadTask {
         requestInfo,
         newTextConnection,
         authToken,
-        appCheckToken
+        appCheckToken,
+        false
       );
       this._request = multipartRequest;
       multipartRequest.getPromise().then(metadata => {