feat: Add custom retries in gax (#489)

Co-authored-by: Brent Shaffer <[email protected]>

Author: saranshdhingra

src/Middleware/RetryMiddleware.php

@@ -51,14 +51,22 @@ class RetryMiddleware
     /** @var float|null */
     private $deadlineMs;
 
+    /*
+     * The number of retries that have already been attempted.
+     * The original API call will have $retryAttempts set to 0.
+     */
+    private int $retryAttempts;
+
     public function __construct(
         callable $nextHandler,
         RetrySettings $retrySettings,
-        $deadlineMs = null
+        $deadlineMs = null,
+        $retryAttempts = 0
     ) {
         $this->nextHandler = $nextHandler;
         $this->retrySettings = $retrySettings;
         $this->deadlineMs = $deadlineMs;
+        $this->retryAttempts = $retryAttempts;
     }
 
     /**
@@ -86,14 +94,23 @@ public function __invoke(Call $call, array $options)
         }
 
         return $nextHandler($call, $options)->then(null, function ($e) use ($call, $options) {
-            if (!$e instanceof ApiException) {
+            $retryFunction = $this->getRetryFunction();
+
+            // If the number of retries has surpassed the max allowed retries
+            // then throw the exception as we normally would.
+            // If the maxRetries is set to 0, then we don't check this condition.
+            if (0 !== $this->retrySettings->getMaxRetries()
+                && $this->retryAttempts >= $this->retrySettings->getMaxRetries()
+            ) {
                 throw $e;
             }
-
-            if (!in_array($e->getStatus(), $this->retrySettings->getRetryableCodes())) {
+            // If the retry function returns false then throw the
+            // exception as we normally would.
+            if (!$retryFunction($e, $options)) {
                 throw $e;
             }
 
+            // Retry function returned true, so we attempt another retry
             return $this->retry($call, $options, $e->getStatus());
         });
     }
@@ -139,7 +156,8 @@ private function retry(Call $call, array $options, string $status)
             $this->retrySettings->with([
                 'initialRetryDelayMillis' => $delayMs,
             ]),
-            $deadlineMs
+            $deadlineMs,
+            $this->retryAttempts + 1
         );
 
         // Set the timeout for the call
@@ -155,4 +173,26 @@ protected function getCurrentTimeMs()
     {
         return microtime(true) * 1000.0;
     }
+
+    /**
+     * This is the default retry behaviour.
+     */
+    private function getRetryFunction()
+    {
+        return $this->retrySettings->getRetryFunction() ??
+            function (\Exception $e, array $options): bool {
+                // This is the default retry behaviour, i.e. we don't retry an ApiException
+                // and for other exception types, we only retry when the error code is in
+                // the list of retryable error codes.
+                if (!$e instanceof ApiException) {
+                    return false;
+                }
+
+                if (!in_array($e->getStatus(), $this->retrySettings->getRetryableCodes())) {
+                    return false;
+                }
+
+                return true;
+            };
+    }
 }

src/RetrySettings.php

@@ -31,6 +31,8 @@
  */
 namespace Google\ApiCore;
 
+use Closure;
+
 /**
  * The RetrySettings class is used to configure retrying and timeouts for RPCs.
  * This class can be passed as an optional parameter to RPC methods, or as part
@@ -203,6 +205,8 @@ class RetrySettings
 {
     use ValidationTrait;
 
+    const DEFAULT_MAX_RETRIES = 0;
+
     private $retriesEnabled;
 
     private $retryableCodes;
@@ -217,6 +221,20 @@ class RetrySettings
 
     private $noRetriesRpcTimeoutMillis;
 
+    /**
+     * The number of maximum retries an operation can do.
+     * This doesn't include the original API call.
+     * Setting this to 0 means no limit.
+     */
+    private int $maxRetries;
+
+    /**
+     * When set, this function will be used to evaluate if the retry should
+     * take place or not. The callable will have the following signature:
+     * function (Exception $e, array $options): bool
+     */
+    private ?Closure $retryFunction;
+
     /**
      * Constructs an instance.
      *
@@ -225,22 +243,28 @@ class RetrySettings
      *     $retriesEnabled and $noRetriesRpcTimeoutMillis, which are optional and have defaults
      *     determined based on the other settings provided.
      *
-     *     @type bool    $retriesEnabled Optional. Enables retries. If not specified, the value is
-     *                   determined using the $retryableCodes setting. If $retryableCodes is empty,
-     *                   then $retriesEnabled is set to false; otherwise, it is set to true.
-     *     @type int     $noRetriesRpcTimeoutMillis Optional. The timeout of the rpc call to be used
-     *                   if $retriesEnabled is false, in milliseconds. It not specified, the value
-     *                   of $initialRpcTimeoutMillis is used.
-     *     @type array   $retryableCodes The Status codes that are retryable. Each status should be
-     *                   either one of the string constants defined on {@see \Google\ApiCore\ApiStatus}
-     *                   or an integer constant defined on {@see \Google\Rpc\Code}.
-     *     @type int     $initialRetryDelayMillis The initial delay of retry in milliseconds.
-     *     @type int     $retryDelayMultiplier The exponential multiplier of retry delay.
-     *     @type int     $maxRetryDelayMillis The max delay of retry in milliseconds.
-     *     @type int     $initialRpcTimeoutMillis The initial timeout of rpc call in milliseconds.
-     *     @type int     $rpcTimeoutMultiplier The exponential multiplier of rpc timeout.
-     *     @type int     $maxRpcTimeoutMillis The max timeout of rpc call in milliseconds.
-     *     @type int     $totalTimeoutMillis The max accumulative timeout in total.
+     *     @type bool     $retriesEnabled Optional. Enables retries. If not specified, the value is
+     *                    determined using the $retryableCodes setting. If $retryableCodes is empty,
+     *                    then $retriesEnabled is set to false; otherwise, it is set to true.
+     *     @type int      $noRetriesRpcTimeoutMillis Optional. The timeout of the rpc call to be used
+     *                    if $retriesEnabled is false, in milliseconds. It not specified, the value
+     *                    of $initialRpcTimeoutMillis is used.
+     *     @type array    $retryableCodes The Status codes that are retryable. Each status should be
+     *                    either one of the string constants defined on {@see \Google\ApiCore\ApiStatus}
+     *                    or an integer constant defined on {@see \Google\Rpc\Code}.
+     *     @type int      $initialRetryDelayMillis The initial delay of retry in milliseconds.
+     *     @type int      $retryDelayMultiplier The exponential multiplier of retry delay.
+     *     @type int      $maxRetryDelayMillis The max delay of retry in milliseconds.
+     *     @type int      $initialRpcTimeoutMillis The initial timeout of rpc call in milliseconds.
+     *     @type int      $rpcTimeoutMultiplier The exponential multiplier of rpc timeout.
+     *     @type int      $maxRpcTimeoutMillis The max timeout of rpc call in milliseconds.
+     *     @type int      $totalTimeoutMillis The max accumulative timeout in total.
+     *     @type int      $maxRetries The max retries allowed for an operation.
+     *                    Defaults to the value of the DEFAULT_MAX_RETRIES constant.
+     *                    This option is experimental.
+     *     @type callable $retryFunction This function will be used to decide if we should retry or not.
+     *                    Callable signature: `function (Exception $e, array $options): bool`
+     *                    This option is experimental.
      * }
      */
     public function __construct(array $settings)
@@ -269,6 +293,8 @@ public function __construct(array $settings)
         $this->noRetriesRpcTimeoutMillis = array_key_exists('noRetriesRpcTimeoutMillis', $settings)
             ? $settings['noRetriesRpcTimeoutMillis']
             : $this->initialRpcTimeoutMillis;
+        $this->maxRetries = $settings['maxRetries'] ?? self::DEFAULT_MAX_RETRIES;
+        $this->retryFunction = $settings['retryFunction'] ?? null;
     }
 
     /**
@@ -348,7 +374,9 @@ public static function constructDefault()
             'rpcTimeoutMultiplier' => 1,
             'maxRpcTimeoutMillis' => 20000,
             'totalTimeoutMillis' => 600000,
-            'retryableCodes' => []]);
+            'retryableCodes' => [],
+            'maxRetries' => self::DEFAULT_MAX_RETRIES,
+            'retryFunction' => null]);
     }
 
     /**
@@ -375,6 +403,8 @@ public function with(array $settings)
             'retryableCodes' => $this->getRetryableCodes(),
             'retriesEnabled' => $this->retriesEnabled(),
             'noRetriesRpcTimeoutMillis' => $this->getNoRetriesRpcTimeoutMillis(),
+            'maxRetries' => $this->getMaxRetries(),
+            'retryFunction' => $this->getRetryFunction(),
         ];
         return new RetrySettings($settings + $existingSettings);
     }
@@ -489,6 +519,22 @@ public function getTotalTimeoutMillis()
         return $this->totalTimeoutMillis;
     }
 
+    /**
+     * @experimental
+     */
+    public function getMaxRetries()
+    {
+        return $this->maxRetries;
+    }
+
+    /**
+     * @experimental
+     */
+    public function getRetryFunction()
+    {
+        return $this->retryFunction;
+    }
+
     private static function convertArrayFromSnakeCase(array $settings)
     {
         $camelCaseSettings = [];