feat(Bigtable): add attempts headers on retry and exception logging (#…

…7428)
googleapis · Jun 21, 2024 · 48fd79b · 48fd79b
1 parent f59ff28
commit 48fd79b
Show file tree

Hide file tree

Showing 6 changed files with 183 additions and 11 deletions.
diff --git a/Bigtable/src/BigtableClient.php b/Bigtable/src/BigtableClient.php
@@ -23,6 +23,7 @@
 use Google\Auth\FetchAuthTokenInterface;
 use Google\Cloud\Bigtable\V2\BigtableClient as GapicClient;
 use Google\Cloud\Core\ArrayTrait;
+use Psr\Log\LoggerInterface;
 
 /**
  * Google Cloud Bigtable is Google's NoSQL Big Data database service.
@@ -144,6 +145,7 @@ public function __construct(array $config = [])
      *     @type string $appProfileId This value specifies routing for
      *           replication. **Defaults to** the "default" application profile.
      *     @type array $headers Headers to be passed with each request.
+     *     @type LoggerInterface $logger
      * }
      * @return Table
      */

diff --git a/Bigtable/src/ChunkFormatter.php b/Bigtable/src/ChunkFormatter.php
@@ -140,7 +140,8 @@ function ($ex) {
                 }
                 return false;
             },
-            $this->pluck('retries', $this->options, false)
+            $this->pluck('retries', $this->options, false),
+            $this->pluck('logger', $this->options, false)
         );
     }
 

diff --git a/Bigtable/src/ResumableStream.php b/Bigtable/src/ResumableStream.php
@@ -20,6 +20,7 @@
 use Google\ApiCore\ApiException;
 use Google\Cloud\Core\ExponentialBackoff;
 use Google\Rpc\Code;
+use Psr\Log\LoggerInterface;
 
 /**
  * User stream which handles failure from upstream, retries if necessary and
@@ -58,6 +59,11 @@ class ResumableStream implements \IteratorAggregate
      */
     private $retryFunction;
 
+    /**
+     * @var LoggerInterface
+     */
+    private $logger;
+
     /**
      * Constructs a resumable stream.
      *
@@ -72,12 +78,14 @@ public function __construct(
         callable $apiFunction,
         callable $argumentFunction,
         callable $retryFunction,
-        $retries = self::DEFAULT_MAX_RETRIES
+        $retries = self::DEFAULT_MAX_RETRIES,
+        LoggerInterface $logger = null
     ) {
         $this->retries = $retries ?: self::DEFAULT_MAX_RETRIES;
         $this->apiFunction = $apiFunction;
         $this->argumentFunction = $argumentFunction;
         $this->retryFunction = $retryFunction;
+        $this->logger = $logger;
     }
 
     /**
@@ -95,12 +103,23 @@ public function readAll()
             $ex = null;
             $args = $argumentFunction();
             if (!isset($args[1]['requestCompleted']) || $args[1]['requestCompleted'] !== true) {
+                if ($tries > 0) {
+                    // Send in "bigtable-attempt" header on retry
+                    $optionalArgs = array_pop($args);
+                    $headers = $optionalArgs['headers'] ?? [];
+                    $headers['bigtable-attempt'] = [(string) $tries];
+                    $optionalArgs['headers'] = $headers;
+                    $args[] = $optionalArgs;
+                }
                 $stream = $this->createExponentialBackoff()->execute($this->apiFunction, $args);
                 try {
                     foreach ($stream->readAll() as $item) {
                         yield $item;
                     }
                 } catch (\Exception $ex) {
+                    if ($this->logger) {
+                        $this->logger->error($ex->getMessage());
+                    }
                 }
             }
             $tries++;

diff --git a/Bigtable/src/Table.php b/Bigtable/src/Table.php
@@ -30,6 +30,7 @@
 use Google\Cloud\Bigtable\V2\RowSet;
 use Google\Cloud\Core\ArrayTrait;
 use Google\Rpc\Code;
+use Psr\Log\LoggerInterface;
 
 /**
  * A table instance can be used to read rows and to perform insert, update, and
@@ -83,6 +84,7 @@ class Table
      *           This settings only applies to {@see \Google\Cloud\Bigtable\Table::mutateRows()},
      *           {@see \Google\Cloud\Bigtable\Table::upsert()} and
      *           {@see \Google\Cloud\Bigtable\Table::readRows()}.
+     *     @type LoggerInterface $logger
      * }
      */
     public function __construct(
@@ -519,7 +521,8 @@ private function mutateRowsWithEntries(array $entries, array $options = [])
             [$this->gapicClient, 'mutateRows'],
             $argumentFunction,
             $retryFunction,
-            $this->pluck('retries', $options, false)
+            $this->pluck('retries', $options, false),
+            $this->pluck('logger', $options, false)
         );
         $message = 'partial failure';
         try {

diff --git a/Bigtable/tests/Unit/SmartRetriesTest.php b/Bigtable/tests/Unit/SmartRetriesTest.php
@@ -37,6 +37,7 @@
 use PHPUnit\Framework\TestCase;
 use Prophecy\PhpUnit\ProphecyTrait;
 use Prophecy\Argument;
+use Psr\Log\LoggerInterface;
 
 /**
  * @group bigtable
@@ -91,12 +92,22 @@ public function testReadRowsShouldRetryDefaultTimes()
         $this->expectExceptionMessage('DEADLINE_EXCEEDED');
 
         $expectedArgs = $this->options;
+        $attempt = 0;
         $this->serverStream->readAll()
             ->shouldBeCalledTimes(4)
             ->willThrow(
                 $this->retryingApiException
             );
-        $this->bigtableClient->readRows(self::TABLE_NAME, $expectedArgs)
+        $this->bigtableClient->readRows(
+            self::TABLE_NAME,
+            Argument::that(function ($optionalArgs) use ($expectedArgs, &$attempt) {
+                $attemptHeader = $optionalArgs['headers']['bigtable-attempt'][0] ?? null;
+                unset($optionalArgs['headers']['bigtable-attempt']);
+
+                return $optionalArgs === $expectedArgs
+                    && ($attemptHeader === (string) $attempt++ || $attempt === 1);
+            })
+        )
             ->shouldBeCalledTimes(4)
             ->willReturn(
                 $this->serverStream->reveal()
@@ -112,12 +123,22 @@ public function testReadRowsShouldRetryForProvidedAttempts()
         $this->expectExceptionMessage('DEADLINE_EXCEEDED');
 
         $expectedArgs = $this->options;
+        $attempt = 0;
         $this->serverStream->readAll()
             ->shouldBeCalledTimes(6)
             ->willThrow(
                 $this->retryingApiException
             );
-        $this->bigtableClient->readRows(self::TABLE_NAME, $expectedArgs)
+        $this->bigtableClient->readRows(
+            self::TABLE_NAME,
+            Argument::that(function ($optionalArgs) use ($expectedArgs, &$attempt) {
+                $attemptHeader = $optionalArgs['headers']['bigtable-attempt'][0] ?? null;
+                unset($optionalArgs['headers']['bigtable-attempt']);
+
+                return $optionalArgs === $expectedArgs
+                    && ($attemptHeader === (string) $attempt++ || $attempt === 1);
+            })
+        )
             ->shouldBeCalledTimes(6)
             ->willReturn(
                 $this->serverStream->reveal()
@@ -147,6 +168,7 @@ public function testReadRowsPartialSuccess()
             ->willReturn(
                 $this->serverStream->reveal()
             );
+        $expectedArgs['headers']['bigtable-attempt'] = ['1'];
         $secondCallArgument = [
             'rows' => (new RowSet)->setRowRanges([(new RowRange)->setStartKeyOpen('rk2')])
         ] + $expectedArgs;
@@ -169,6 +191,39 @@ public function testReadRowsPartialSuccess()
         }
     }
 
+    public function testReadRowsContainsAttemptHeader()
+    {
+        $attempt = 0;
+        $expectedArgs = $this->options;
+        $retryingApiException = $this->retryingApiException;
+        $this->bigtableClient->readRows(
+            self::TABLE_NAME,
+            Argument::that(function ($callOptions) use (&$attempt) {
+                $attemptHeader = $callOptions['headers']['bigtable-attempt'][0] ?? null;
+                ($attempt === 0)
+                    ? $this->assertNull($attemptHeader)
+                    : $this->assertSame('1', $attemptHeader);
+
+                return true;
+            })
+        )->shouldBeCalledTimes(2)
+            ->willReturn(
+                $this->serverStream->reveal()
+            );
+
+        $this->serverStream->readAll()
+            ->will(function () use (&$attempt, $retryingApiException) {
+                // throw a retriable exception on the first call
+                if (0 === $attempt++) {
+                    throw $retryingApiException;
+                }
+                return [];
+            });
+
+        $iterator = $this->table->readRows();
+        $iterator->getIterator()->current();
+    }
+
     public function testReadRowsWithRowsLimit()
     {
         $args = ['rowsLimit' => 5];
@@ -189,6 +244,7 @@ public function testReadRowsWithRowsLimit()
             ->willReturn(
                 $this->serverStream->reveal()
             );
+        $expectedArgs['headers']['bigtable-attempt'] = ['1'];
         $secondCallArgument = [
             'rows' => (new RowSet)->setRowRanges([(new RowRange)->setStartKeyOpen('rk2')]),
             'rowsLimit' => 3
@@ -519,7 +575,18 @@ public function testMutateRowsShouldRetryDefaultNumberOfTimes()
             );
         $mutations = $this->generateMutations(1, 5);
         $entries = $this->generateEntries(1, 5);
-        $this->bigtableClient->mutateRows(self::TABLE_NAME, $entries, $this->options)
+        $attempt = 0;
+        $this->bigtableClient->mutateRows(
+            self::TABLE_NAME,
+            $entries,
+            Argument::that(function ($optionalArgs) use (&$attempt) {
+                $attemptHeader = $optionalArgs['headers']['bigtable-attempt'][0] ?? null;
+                unset($optionalArgs['headers']['bigtable-attempt']);
+
+                return $optionalArgs === $this->options
+                    && ($attemptHeader === (string) $attempt++ || $attempt === 1);
+            })
+        )
             ->shouldBeCalledTimes(4)
             ->willReturn(
                 $this->serverStream->reveal()
@@ -539,7 +606,19 @@ public function testMutateRowsRespectRetriesAttempt()
             );
         $mutations = $this->generateMutations(1, 5);
         $entries = $this->generateEntries(1, 5);
-        $this->bigtableClient->mutateRows(self::TABLE_NAME, $entries, ['retries' => 5] + $this->options)
+        $expectedArgs = ['retries' => 5] + $this->options;
+        $attempt = 0;
+        $this->bigtableClient->mutateRows(
+            self::TABLE_NAME,
+            $entries,
+            Argument::that(function ($optionalArgs) use ($expectedArgs, &$attempt) {
+                $attemptHeader = $optionalArgs['headers']['bigtable-attempt'][0] ?? null;
+                unset($optionalArgs['headers']['bigtable-attempt']);
+
+                return $optionalArgs === $expectedArgs
+                    && ($attemptHeader === (string) $attempt++ || $attempt === 1);
+            })
+        )
             ->shouldBeCalledTimes(6)
             ->willReturn(
                 $this->serverStream->reveal()
@@ -549,6 +628,7 @@ public function testMutateRowsRespectRetriesAttempt()
 
     public function testMutateRowsOnlyRetriesFailedEntries()
     {
+        $expectedArgs = $this->options;
         $this->serverStream->readAll()
             ->shouldBeCalledTimes(2)
             ->willReturn(
@@ -560,13 +640,14 @@ public function testMutateRowsOnlyRetriesFailedEntries()
                 )
             );
         $entries = $this->generateEntries(0, 5);
-        $this->bigtableClient->mutateRows(self::TABLE_NAME, $entries, $this->options)
+        $this->bigtableClient->mutateRows(self::TABLE_NAME, $entries, $expectedArgs)
             ->shouldBeCalledTimes(1)
             ->willReturn(
                 $this->serverStream->reveal()
             );
+        $expectedArgs['headers']['bigtable-attempt'] = ['1'];
         $entries = $this->generateEntries(2, 3);
-        $this->bigtableClient->mutateRows(self::TABLE_NAME, $entries, $this->options)
+        $this->bigtableClient->mutateRows(self::TABLE_NAME, $entries, $expectedArgs)
             ->shouldBeCalled()
             ->willReturn(
                 $this->serverStream->reveal()
@@ -577,6 +658,7 @@ public function testMutateRowsOnlyRetriesFailedEntries()
 
     public function testMutateRowsExceptionShouldAddEntryToPendingMutations()
     {
+        $expectedArgs = $this->options;
         $this->serverStream->readAll()
             ->shouldBeCalledTimes(2)
             ->willReturn(
@@ -589,13 +671,14 @@ public function testMutateRowsExceptionShouldAddEntryToPendingMutations()
                 )
             );
         $entries = $this->generateEntries(0, 5);
-        $this->bigtableClient->mutateRows(self::TABLE_NAME, $entries, $this->options)
+        $this->bigtableClient->mutateRows(self::TABLE_NAME, $entries, $expectedArgs)
             ->shouldBeCalledTimes(1)
             ->willReturn(
                 $this->serverStream->reveal()
             );
         $entries = array_merge($this->generateEntries(1, 2), $this->generateEntries(4, 5));
-        $this->bigtableClient->mutateRows(self::TABLE_NAME, $entries, $this->options)
+        $expectedArgs['headers']['bigtable-attempt'] = ['1'];
+        $this->bigtableClient->mutateRows(self::TABLE_NAME, $entries, $expectedArgs)
             ->shouldBeCalled()
             ->willReturn(
                 $this->serverStream->reveal()
@@ -667,6 +750,66 @@ public function testMutateRowsShouldNotRetryIfAnyMutationIsNotRetryable()
         }
     }
 
+    public function testReadRowsShouldLogRetryableExeception()
+    {
+        $attempt = 0;
+        $expectedArgs = $this->options;
+        $retryingApiException = $this->retryingApiException;
+        $logger = $this->prophesize(LoggerInterface::class);
+        $logger->error('DEADLINE_EXCEEDED')
+            ->shouldBeCalledOnce();
+        $this->bigtableClient->readRows(
+            self::TABLE_NAME,
+            Argument::that(function ($callOptions) use (&$attempt) {
+                $attemptHeader = $callOptions['headers']['bigtable-attempt'][0] ?? null;
+                ($attempt === 0)
+                    ? $this->assertNull($attemptHeader)
+                    : $this->assertSame('1', $attemptHeader);
+
+                return true;
+            })
+        )->shouldBeCalledTimes(2)
+            ->willReturn(
+                $this->serverStream->reveal()
+            );
+
+        $this->serverStream->readAll()
+            ->will(function () use (&$attempt, $retryingApiException) {
+                // throw a retriable exception on the first call
+                if (0 === $attempt++) {
+                    throw $retryingApiException;
+                }
+                return [];
+            });
+
+        $iterator = $this->table->readRows(['logger' => $logger->reveal()]);
+        $iterator->getIterator()->current();
+    }
+
+    public function testReadRowsShouldLogNonRetryableExeception()
+    {
+        $this->expectException(ApiException::class);
+        $this->expectExceptionMessage('UNAUTHENTICATED');
+
+        $logger = $this->prophesize(LoggerInterface::class);
+        $logger->error('UNAUTHENTICATED')
+            ->shouldBeCalledOnce();
+
+        $this->bigtableClient->readRows(
+            self::TABLE_NAME,
+            $this->options
+        )->shouldBeCalledOnce()
+            ->willReturn(
+                $this->serverStream->reveal()
+            );
+
+        $this->serverStream->readAll()
+            ->willThrow($this->nonRetryingApiException);
+
+        $iterator = $this->table->readRows(['logger' => $logger->reveal()]);
+        $iterator->getIterator()->current();
+    }
+
     private function generateRowsResponse($from, $to)
     {
         $rows = [];

diff --git a/Core/src/InsecureCredentialsWrapper.php b/Core/src/InsecureCredentialsWrapper.php
@@ -33,4 +33,8 @@ public function getAuthorizationHeaderCallback($audience = null)
     {
         return null;
     }
+
+    public function checkUniverseDomain()
+    {
+    }
 }
-Original file line number
+Diff line change
@@ Expand Up @@
         {
             return null;
         }
+        public function checkUniverseDomain()
+        {
+        }
     }