Skip to content

Commit

Permalink
feat(Bigtable): add bigtable-attempt header (#7414)
Browse files Browse the repository at this point in the history
  • Loading branch information
bshaffer authored Jun 20, 2024
1 parent 7b9054c commit f7f6aec
Show file tree
Hide file tree
Showing 3 changed files with 46 additions and 9 deletions.
2 changes: 1 addition & 1 deletion Bigtable/src/Exception/BigtableDataOperationException.php
Original file line number Diff line number Diff line change
Expand Up @@ -41,7 +41,7 @@ public function __construct(
) {
$this->metadata = $metadata;

parent::__construct($message, $code);
parent::__construct($message, $code ?? 0);
}

/**
Expand Down
19 changes: 12 additions & 7 deletions Bigtable/src/ResumableStream.php
Original file line number Diff line number Diff line change
Expand Up @@ -126,19 +126,25 @@ public function __construct(
*/
public function readAll()
{
$tries = 0;
$argumentFunction = $this->argumentFunction;
$retryFunction = $this->retryFunction;
$attempt = 0;
do {
$ex = null;
list($this->request, $this->callOptions) = $argumentFunction($this->request, $this->callOptions);
list($this->request, $this->callOptions) =
($this->argumentFunction)($this->request, $this->callOptions);

$completed = $this->pluck('requestCompleted', $this->callOptions, false);

if ($completed !== true) {
// Send in "bigtable-attempt" header on retry request
$headers = $this->callOptions['headers'] ?? [];
if ($attempt > 0) {
$headers['bigtable-attempt'] = [(string) $attempt];
}
$attempt++;

$stream = call_user_func_array(
[$this->gapicClient, $this->method],
[$this->request, $this->callOptions]
[$this->request, ['headers' => $headers] + $this->callOptions]
);

try {
Expand All @@ -148,8 +154,7 @@ public function readAll()
} catch (\Exception $ex) {
}
}
$tries++;
} while ((!$this->retryFunction || $retryFunction($ex)) && $tries <= $this->retries);
} while ((!$this->retryFunction || ($this->retryFunction)($ex)) && $attempt <= $this->retries);
if ($ex !== null) {
throw $ex;
}
Expand Down
34 changes: 33 additions & 1 deletion Bigtable/tests/Unit/SmartRetriesTest.php
Original file line number Diff line number Diff line change
Expand Up @@ -136,6 +136,38 @@ public function testReadRowsShouldRetryForProvidedAttempts()
$iterator->getIterator()->current();
}

public function testReadRowsContainsAttemptHeader()
{
$attempt = 0;
$expectedArgs = $this->options;
$retryingApiException = $this->retryingApiException;
$this->serverStream->readAll()
->shouldBeCalledTimes(2)
->will(function () use (&$attempt, $retryingApiException) {
// throw a retriable exception on the first call
return 0 === $attempt++ ? throw $retryingApiException : [];
});
$this->bigtableClient->readRows(
Argument::type(ReadRowsRequest::class),
Argument::that(function ($callOptions) use (&$attempt) {
$attemptHeader = $callOptions['headers']['bigtable-attempt'][0] ?? null;
if ($attempt === 0) {
$this->assertNull($attemptHeader);
} else {
$this->assertSame((string) $attempt, $attemptHeader);
}

return true;
})
)->shouldBeCalledTimes(2)
->willReturn(
$this->serverStream->reveal()
);

$iterator = $this->table->readRows();
$iterator->getIterator()->current();
}

public function testReadRowsPartialSuccess()
{
$expectedArgs = $this->options;
Expand Down Expand Up @@ -186,7 +218,7 @@ public function testReadRowsWithRowsLimit()
$this->generateRowsResponse(3, 4)
)
);

$allowedRowsLimit = ['5' => 1, '3' => 1];
$this->bigtableClient->readRows(
Argument::that(function ($request) use (&$allowedRowsLimit) {
Expand Down

0 comments on commit f7f6aec

Please sign in to comment.