From ef0789b73ef28d79a08c354d1361a9ccc6206088 Mon Sep 17 00:00:00 2001 From: Saransh Dhingra Date: Tue, 31 Oct 2023 17:33:18 +0530 Subject: [PATCH] feat: Add custom retries in gax (#489) Co-authored-by: Brent Shaffer --- src/Middleware/RetryMiddleware.php | 50 +++- src/RetrySettings.php | 80 ++++-- .../Unit/Middleware/RetryMiddlewareTest.php | 239 ++++++++++++++++++ tests/Tests/Unit/RetrySettingsTest.php | 45 +++- 4 files changed, 391 insertions(+), 23 deletions(-) diff --git a/src/Middleware/RetryMiddleware.php b/src/Middleware/RetryMiddleware.php index 84b33013d..df94244fc 100644 --- a/src/Middleware/RetryMiddleware.php +++ b/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; + }; + } } diff --git a/src/RetrySettings.php b/src/RetrySettings.php index c8af73292..94bb6f189 100644 --- a/src/RetrySettings.php +++ b/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 = []; diff --git a/tests/Tests/Unit/Middleware/RetryMiddlewareTest.php b/tests/Tests/Unit/Middleware/RetryMiddlewareTest.php index 2dbdfb224..a31f118e5 100644 --- a/tests/Tests/Unit/Middleware/RetryMiddlewareTest.php +++ b/tests/Tests/Unit/Middleware/RetryMiddlewareTest.php @@ -275,4 +275,243 @@ public function testNoRetryLogicalTimeout() $this->assertSame('Ok!', $response); $this->assertEquals($observedTimeout, $timeout); } + + public function testCustomRetry() + { + $call = $this->getMockBuilder(Call::class) + ->disableOriginalConstructor() + ->getMock(); + + $maxAttempts = 3; + $currentAttempt = 0; + $retrySettings = RetrySettings::constructDefault() + ->with([ + 'retriesEnabled' => true, + 'retryFunction' => function ($ex, $options) use ($maxAttempts, &$currentAttempt) { + $currentAttempt++; + if($currentAttempt < $maxAttempts) { + return true; + } + + return false; + } + ]); + $callCount = 0; + $handler = function(Call $call, $options) use (&$callCount) { + return new Promise(function () use (&$callCount) { + ++$callCount; + throw new ApiException('Call Count: ' . $callCount, 0, ''); + }); + }; + $middleware = new RetryMiddleware($handler, $retrySettings); + + $this->expectException(ApiException::class); + // test if the custom retry func threw an exception after $maxAttempts + $this->expectExceptionMessage('Call Count: ' . ($maxAttempts)); + + $middleware($call, [])->wait(); + } + + public function testCustomRetryRespectsRetriesEnabled() + { + $call = $this->getMockBuilder(Call::class) + ->disableOriginalConstructor() + ->getMock(); + + $retrySettings = RetrySettings::constructDefault() + ->with([ + 'retriesEnabled' => false, + 'retryFunction' => function ($ex, $options) { + // This should not run as retriesEnabled is false + $this->fail('Custom retry function shouldn\'t have run.'); + return true; + } + ]); + + $handler = function (Call $call, $options) { + return new Promise(function () { + throw new ApiException('Exception msg', 0, ''); + }); + }; + $middleware = new RetryMiddleware($handler, $retrySettings); + + $this->expectException(ApiException::class); + $middleware($call, [])->wait(); + } + + public function testCustomRetryRespectsTimeout() + { + $call = $this->getMockBuilder(Call::class) + ->disableOriginalConstructor() + ->getMock(); + + $retrySettings = RetrySettings::constructDefault() + ->with([ + 'retriesEnabled' => true, + 'totalTimeoutMillis' => 1, + 'retryFunction' => function ($ex, $options) { + usleep(900); + return true; + } + ]); + + $callCount = 0; + $handler = function (Call $call, $options) use (&$callCount) { + return new Promise(function () use(&$callCount) { + ++$callCount; + throw new ApiException('Call count: ' . $callCount, 0, ''); + }); + }; + $middleware = new RetryMiddleware($handler, $retrySettings); + + try { + $middleware($call, [])->wait(); + $this->fail('Expected an exception, but didn\'t receive any'); + } catch(ApiException $e) { + $this->assertEquals('Retry total timeout exceeded.', $e->getMessage()); + // we used a total timeout of 1 ms and every retry sleeps for .9 ms + // This means that the call count should be 2(original call and 1 retry) + $this->assertEquals(2, $callCount); + } + } + + public function testMaxRetries() + { + $call = $this->getMockBuilder(Call::class) + ->disableOriginalConstructor() + ->getMock(); + + $maxRetries = 2; + $retrySettings = RetrySettings::constructDefault() + ->with([ + 'retriesEnabled' => true, + 'retryableCodes' => [ApiStatus::CANCELLED], + 'maxRetries' => $maxRetries + ]); + $callCount = 0; + $handler = function(Call $call, $options) use (&$callCount) { + return new Promise(function () use (&$callCount) { + ++$callCount; + throw new ApiException('Call Count: ' . $callCount, 0, ApiStatus::CANCELLED); + }); + }; + $middleware = new RetryMiddleware($handler, $retrySettings); + + $this->expectException(ApiException::class); + // test if the custom retry func threw an exception after $maxRetries + 1 calls + $this->expectExceptionMessage('Call Count: ' . ($maxRetries + 1)); + + $middleware($call, [])->wait(); + } + + /** + * Tests for maxRetries to be evaluated before the retry function. + */ + public function testMaxRetriesOverridesCustomRetryFunc() + { + $call = $this->getMockBuilder(Call::class) + ->disableOriginalConstructor() + ->getMock(); + + $maxRetries = 2; + $callCount = 0; + $retrySettings = RetrySettings::constructDefault() + ->with([ + 'retriesEnabled' => true, + 'retryableCodes' => [ApiStatus::CANCELLED], + 'maxRetries' => $maxRetries, + 'retryFunction' => function ($ex, $options) use (&$callCount) { + // The retryFunction will signal a retry until the total call count reaches 5. + return $callCount < 5 ? true : false; + } + ]); + $handler = function(Call $call, $options) use (&$callCount) { + return new Promise(function () use (&$callCount) { + ++$callCount; + throw new ApiException('Call Count: ' . $callCount, 0, ApiStatus::CANCELLED); + }); + }; + $middleware = new RetryMiddleware($handler, $retrySettings); + + $this->expectException(ApiException::class); + // Even though our custom retry function wants 4 retries + // the exception should be thrown after $maxRetries + 1 calls + $this->expectExceptionMessage('Call Count: ' . ($maxRetries + 1)); + + $middleware($call, [])->wait(); + } + + /** + * Tests for custom retry function returning false, before we reach maxRetries. + */ + public function testCustomRetryThrowsExceptionBeforeMaxRetries() + { + $call = $this->getMockBuilder(Call::class) + ->disableOriginalConstructor() + ->getMock(); + + $maxRetries = 10; + $customRetryMaxCalls = 4; + $callCount = 0; + $retrySettings = RetrySettings::constructDefault() + ->with([ + 'retriesEnabled' => true, + 'retryableCodes' => [ApiStatus::CANCELLED], + 'maxRetries' => $maxRetries, + 'retryFunction' => function ($ex, $options) use (&$callCount, $customRetryMaxCalls) { + // The retryFunction will signal a retry until the total call count reaches $customRetryMaxCalls. + return $callCount < $customRetryMaxCalls ? true : false; + } + ]); + $handler = function(Call $call, $options) use (&$callCount) { + return new Promise(function () use (&$callCount) { + ++$callCount; + throw new ApiException('Call Count: ' . $callCount, 0, ApiStatus::CANCELLED); + }); + }; + $middleware = new RetryMiddleware($handler, $retrySettings); + + $this->expectException(ApiException::class); + // Even though our maxRetries hasn't reached + // the exception should be thrown after $customRetryMaxCalls + // because the custom retry function would return false. + $this->expectExceptionMessage('Call Count: ' . ($customRetryMaxCalls)); + + $middleware($call, [])->wait(); + } + + public function testUnlimitedMaxRetries() + { + $call = $this->getMockBuilder(Call::class) + ->disableOriginalConstructor() + ->getMock(); + + $customRetryMaxCalls = 4; + $callCount = 0; + $retrySettings = RetrySettings::constructDefault() + ->with([ + 'retriesEnabled' => true, + 'retryableCodes' => [ApiStatus::CANCELLED], + 'maxRetries' => 0, + 'retryFunction' => function ($ex, $options) use (&$callCount, $customRetryMaxCalls) { + // The retryFunction will signal a retry until the total call count reaches $customRetryMaxCalls. + return $callCount < $customRetryMaxCalls ? true : false; + } + ]); + $handler = function(Call $call, $options) use (&$callCount) { + return new Promise(function () use (&$callCount) { + ++$callCount; + throw new ApiException('Call Count: ' . $callCount, 0, ApiStatus::CANCELLED); + }); + }; + $middleware = new RetryMiddleware($handler, $retrySettings); + + $this->expectException(ApiException::class); + // Since the maxRetries is set to 0(unlimited), + // the exception should be thrown after $customRetryMaxCalls + // because then the custom retry function would return false. + $this->expectExceptionMessage('Call Count: ' . ($customRetryMaxCalls)); + + $middleware($call, [])->wait(); + } } diff --git a/tests/Tests/Unit/RetrySettingsTest.php b/tests/Tests/Unit/RetrySettingsTest.php index 48c1e25d9..6bd428354 100644 --- a/tests/Tests/Unit/RetrySettingsTest.php +++ b/tests/Tests/Unit/RetrySettingsTest.php @@ -70,6 +70,7 @@ public function testConstructSettings() $timeoutOnlyMethod = $defaultRetrySettings['TimeoutOnlyMethod']; $this->assertFalse($timeoutOnlyMethod->retriesEnabled()); $this->assertEquals(40000, $timeoutOnlyMethod->getNoRetriesRpcTimeoutMillis()); + $this->assertEquals(RetrySettings::DEFAULT_MAX_RETRIES, $simpleMethod->getMaxRetries()); } public function testLoadInvalid() @@ -222,7 +223,9 @@ public function retrySettingsProvider() 'totalTimeoutMillis' => 2000, 'retryableCodes' => [1], 'noRetriesRpcTimeoutMillis' => 150, - 'retriesEnabled' => true + 'retriesEnabled' => true, + 'maxRetries' => RetrySettings::DEFAULT_MAX_RETRIES, + 'retryFunction' => null, ]; return [ [ @@ -268,6 +271,24 @@ public function retrySettingsProvider() [ 'noRetriesRpcTimeoutMillis' => 151, ] + $defaultExpectedValues + ], + [ + // Test with a custom retry function + [ + 'retryFunction' => function($ex, $options) {return true;} + ] + $defaultSettings, + [ + 'retryFunction' => function($ex, $options) {return true;} + ] + $defaultExpectedValues + ], + [ + // Test with a maxRetries value + [ + 'maxRetries' => 2 + ] + $defaultSettings, + [ + 'maxRetries' => 2 + ] + $defaultExpectedValues ] ]; } @@ -297,6 +318,8 @@ public function withRetrySettingsProvider() 'retryableCodes' => [1], 'noRetriesRpcTimeoutMillis' => 1, 'retriesEnabled' => true, + 'maxRetries' => RetrySettings::DEFAULT_MAX_RETRIES, + 'retryFunction' => null, ]; return [ [ @@ -342,6 +365,26 @@ public function withRetrySettingsProvider() 'noRetriesRpcTimeoutMillis' => 10, 'retriesEnabled' => false, ] + ], + [ + // Test with a custom retry function + $defaultSettings, + [ + 'retryFunction' => function($ex, $options) {return true;} + ], + [ + 'retryFunction' => function($ex, $options) {return true;} + ] + $defaultExpectedValues + ], + [ + // Test with a maxRetries value + $defaultSettings, + [ + 'maxRetries' => 2 + ], + [ + 'maxRetries' => 2 + ] + $defaultExpectedValues ] ]; }