fix: ensure attempts increment without exception (#7688)

googleapis · Sep 20, 2024 · 95dff6d · 95dff6d
1 parent 89f982a
commit 95dff6d
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 13 deletions.
diff --git a/src/ResumableStream.php b/src/ResumableStream.php
@@ -163,7 +163,6 @@ public function readAll()
                     $headers['bigtable-attempt'] = [(string) $totalAttempt];
                     ($this->delayFunction)($currentAttempt);
                 }
-
                 $stream = call_user_func_array(
                     [$this->gapicClient, $this->method],
                     [$this->request, ['headers' => $headers] + $this->callOptions]
@@ -175,11 +174,14 @@ public function readAll()
                         $currentAttempt = 0; // reset delay and attempt on successful read.
                     }
                 } catch (\Exception $ex) {
-                    $totalAttempt++;
-                    $currentAttempt++;
                 }
+                // It's possible for the retry function to retry even when `$ex` is null
+                // (see Table::mutateRowsWithEntries). For this reason, we increment the attemts
+                // outside the try/catch block.
+                $totalAttempt++;
+                $currentAttempt++;
             }
-        } while ((!$this->retryFunction || ($this->retryFunction)($ex)) && $currentAttempt <= $this->retries);
+        } while (($this->retryFunction)($ex) && $currentAttempt <= $this->retries);
         if ($ex !== null) {
             throw $ex;
         }

diff --git a/tests/Unit/SmartRetriesTest.php b/tests/Unit/SmartRetriesTest.php
@@ -676,20 +676,22 @@ public function testMutateRowsOnlyRetriesFailedEntries()
                 return iterator_to_array($request->getEntries()) == $entries;
             }),
             Argument::type('array')
-        )->shouldBeCalledTimes(1)
-        ->willReturn(
-            $this->serverStream->reveal()
-        );
+        )
+            ->shouldBeCalledTimes(1)
+            ->willReturn($this->serverStream->reveal());
+
         $entries = $this->generateEntries(2, 3);
+
+        // ensure the bigtable-attempt header is sent in for the next retry.
         $this->bigtableClient->mutateRows(
             Argument::that(function ($request) use ($entries) {
                 return iterator_to_array($request->getEntries()) == $entries;
             }),
-            Argument::type('array')
-        )->shouldBeCalled()
-        ->willReturn(
-            $this->serverStream->reveal()
-        );
+            Argument::withEntry('headers', ['bigtable-attempt' => ['1']] + $this->options['headers'])
+        )
+            ->shouldBeCalledTimes(1)
+            ->willReturn($this->serverStream->reveal());
+
         $mutations = $this->generateMutations(0, 5);
         $this->table->mutateRows($mutations);
     }