fix(Bigtable): prevent race condition (#7644)

googleapis · Sep 5, 2024 · f94f1d4 · f94f1d4
1 parent 7c699f0
commit f94f1d4
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 1 deletion.
diff --git a/Bigtable/src/ChunkFormatter.php b/Bigtable/src/ChunkFormatter.php
@@ -142,7 +142,7 @@ public function __construct(
         $this->gapicClient = $gapicClient;
         $this->request = $request;
         $this->options = $options;
-        
+
         if ($request->getRowsLimit()) {
             $this->originalRowsLimit = $request->getRowsLimit();
         }
@@ -228,6 +228,9 @@ private function updateReadRowsRequest($request, $options, $prevRowKey)
     {
         if ($this->originalRowsLimit) {
             $request->setRowsLimit($this->originalRowsLimit - $this->numberOfRowsRead);
+            if ($this->numberOfRowsRead === $this->originalRowsLimit) {
+                $options['requestCompleted'] = true;
+            }
         }
 
         if ($request->hasRows()) {

diff --git a/Bigtable/tests/Unit/SmartRetriesTest.php b/Bigtable/tests/Unit/SmartRetriesTest.php
@@ -247,6 +247,34 @@ public function testReadRowsWithRowsLimit()
         $this->assertEquals($expectedRows, $rows);
     }
 
+    public function testReadRowsWithRowsLimitAndExceptionThrownAfterSuccess()
+    {
+        $args = ['rowsLimit' => 1];
+        $expectedArgs = $args + $this->options;
+        $this->serverStream->readAll()
+            ->shouldBeCalledTimes(1)
+            ->willReturn(
+                $this->arrayAsGeneratorWithException(
+                    $this->generateRowsResponse(1, 1),
+                    $this->retryingApiException
+                )
+            );
+        $this->bigtableClient->readRows(
+            Argument::that(function ($request) use (&$allowedRowsLimit) {
+                return $request->getRowsLimit() === 1;
+            }),
+            Argument::type('array')
+        )
+            ->shouldBeCalledTimes(1)
+            ->willReturn(
+                $this->serverStream->reveal()
+            );
+
+        $rows = $this->table->readRows($args);
+        $expectedRows = $this->generateExpectedRows(1, 1);
+        $this->assertEquals($expectedRows, iterator_to_array($rows));
+    }
+
     public function testReadRowsWithRowKeys()
     {
         $args = ['rowKeys' => ['rk1', 'rk2', 'rk3', 'rk4']];