fix(bigtable): use the configured retry policy for bulk mutation failures

google/cloud/bigtable/internal/async_bulk_apply.cc

@@ -94,7 +94,7 @@ void AsyncBulkApplier::MakeRequest() {
 void AsyncBulkApplier::OnRead(
     google::bigtable::v2::MutateRowsResponse response) {
   limiter_->Update(response);
-  state_.OnRead(std::move(response));
+  state_.OnRead(std::move(response), *options_);
 }
 
 void AsyncBulkApplier::OnFinish(Status const& status) {

google/cloud/bigtable/internal/bulk_mutator.cc

@@ -13,6 +13,7 @@
 // limitations under the License.
 
 #include "google/cloud/bigtable/internal/bulk_mutator.h"
+#include "google/cloud/bigtable/options.h"
 #include "google/cloud/bigtable/rpc_retry_policy.h"
 #include "google/cloud/bigtable/table.h"
 #include "google/cloud/grpc_options.h"
@@ -80,7 +81,19 @@ google::bigtable::v2::MutateRowsRequest const& BulkMutatorState::BeforeStart() {
 }
 
 void BulkMutatorState::OnRead(
-    google::bigtable::v2::MutateRowsResponse response) {
+    google::bigtable::v2::MutateRowsResponse response,
+    Options const& options) {
+  // Use the retry policy configured by the application to decide if a failed
+  // mutation is worth retrying, like we do for the stream failures. The policy
+  // is only asked to classify the error, it should not count it: this is not a
+  // new attempt of the operation, and the retry loop driving the RPC does the
+  // accounting. We fall back to the default classification when the caller did
+  // not configure a policy.
+  auto const& retry_policy = options.get<bigtable::DataRetryPolicyOption>();
+  auto is_transient = [&retry_policy](Status const& status) {
+    if (retry_policy) return !retry_policy->IsPermanentFailure(status);
+    return SafeGrpcRetry::IsTransientFailure(status);
+  };
   for (auto& entry : *response.mutable_entries()) {
     // The type of `entry.index()` is a 64-bit int. But we can never create more
     // than std::numeric_limits<std::size_t>::max() entries in the request
@@ -108,7 +121,7 @@ void BulkMutatorState::OnRead(
     if (status.ok()) continue;
     auto& original = *mutations_.mutable_entries(static_cast<int>(index));
     // Failed responses are handled according to the current policies.
-    if (SafeGrpcRetry::IsTransientFailure(status) &&
+    if (is_transient(status) &&
         (annotation.idempotency == Idempotency::kIdempotent)) {
       // Retryable requests are saved in the pending mutations, along with the
       // mapping from their index in pending_mutations_ to the original
@@ -222,7 +235,7 @@ Status BulkMutator::MakeOneRequest(BigtableStub& stub,
       break;
     }
     limiter.Update(response);
-    state_.OnRead(std::move(response));
+    state_.OnRead(std::move(response), options);
   }
   state_.OnFinish(*std::move(status), enable_server_retries);
   operation_context_->PostCall(*client_context, state_.last_status());

google/cloud/bigtable/internal/bulk_mutator.h

@@ -46,8 +46,14 @@ class BulkMutatorState {
   /// Returns the Request parameter for the next MutateRows() RPC.
   google::bigtable::v2::MutateRowsRequest const& BeforeStart();
 
-  /// Handle the result of a `Read()` operation on the MutateRows RPC.
-  void OnRead(google::bigtable::v2::MutateRowsResponse response);
+  /**
+   * Handle the result of a `Read()` operation on the MutateRows RPC.
+   *
+   * @param options the options in effect for the request that produced
+   *     @p response. They determine if a failed mutation is worth retrying.
+   */
+  void OnRead(google::bigtable::v2::MutateRowsResponse response,
+              Options const& options);
 
   /// Handle the result of a `Finish()` operation on the MutateRows() RPC.
   void OnFinish(Status finish_status, bool enable_server_retries = false);

google/cloud/bigtable/internal/bulk_mutator_test.cc

@@ -14,6 +14,7 @@
 
 #include "google/cloud/bigtable/internal/bulk_mutator.h"
 #include "google/cloud/bigtable/internal/operation_context.h"
+#include "google/cloud/bigtable/options.h"
 #include "google/cloud/grpc_options.h"
 #ifdef GOOGLE_CLOUD_CPP_BIGTABLE_WITH_OTEL_METRICS
 #include "google/cloud/bigtable/internal/metrics.h"
@@ -22,6 +23,7 @@
 #endif
 #include "google/cloud/bigtable/testing/mock_bigtable_stub.h"
 #include "google/cloud/bigtable/testing/mock_mutate_rows_limiter.h"
+#include "google/cloud/bigtable/testing/mock_policies.h"
 #include "google/cloud/internal/make_status.h"
 #include "google/cloud/testing_util/chrono_literals.h"
 #include "google/cloud/testing_util/status_matchers.h"
@@ -36,6 +38,7 @@ namespace {
 namespace v2 = ::google::bigtable::v2;
 using ::google::cloud::testing_util::chrono_literals::operator""_ms;
 using ::google::cloud::bigtable::testing::MockBigtableStub;
+using ::google::cloud::bigtable::testing::MockDataRetryPolicy;
 using ::google::cloud::bigtable::testing::MockMutateRowsLimiter;
 using ::google::cloud::bigtable::testing::MockMutateRowsStream;
 using ::google::cloud::testing_util::StatusIs;
@@ -243,6 +246,91 @@ TEST_F(BulkMutatorTest, RetryPartialFailure) {
   EXPECT_TRUE(failures.empty());
 }
 
+TEST_F(BulkMutatorTest, RetryPolicyMakesMutationFailurePermanent) {
+  // UNAVAILABLE is worth retrying by default, but the application configured a
+  // policy that disagrees, and the mutation failure should be reported instead
+  // of retried.
+  BulkMutation mut(IdempotentMutation("r0"));
+
+  auto mock = std::make_shared<MockBigtableStub>();
+  EXPECT_CALL(*mock, MutateRows)
+      .WillOnce([this](auto context, auto const&,
+                       google::bigtable::v2::MutateRowsRequest const& request,
+                       auto const&) {
+        metadata_fixture_.SetServerMetadata(*context, {});
+        EXPECT_THAT(request, HasCorrectResourceNames());
+        auto stream = std::make_unique<MockMutateRowsStream>();
+        EXPECT_CALL(*stream, Read)
+            .WillOnce([](google::bigtable::v2::MutateRowsResponse* r) {
+              *r = MakeResponse({{0, grpc::StatusCode::UNAVAILABLE}});
+              return std::nullopt;
+            })
+            .WillOnce(Return(Status()));
+        return stream;
+      });
+
+  auto retry_policy = std::make_shared<MockDataRetryPolicy>();
+  EXPECT_CALL(*retry_policy, IsPermanentFailure)
+      .WillRepeatedly(Return(true));
+  auto policy = DefaultIdempotentMutationPolicy();
+  bigtable_internal::BulkMutator mutator(
+      kAppProfile, kTableName, *policy, std::move(mut),
+      std::make_shared<bigtable_internal::OperationContext>());
+
+  EXPECT_TRUE(mutator.HasPendingMutations());
+  bigtable_internal::NoopMutateRowsLimiter limiter;
+  auto status = mutator.MakeOneRequest(
+      *mock, limiter, Options{}.set<DataRetryPolicyOption>(retry_policy));
+  EXPECT_STATUS_OK(status);
+  EXPECT_FALSE(mutator.HasPendingMutations());
+  auto failures = std::move(mutator).OnRetryDone();
+  ASSERT_EQ(1UL, failures.size());
+  EXPECT_EQ(0, failures[0].original_index());
+  EXPECT_THAT(failures[0].status(), StatusIs(StatusCode::kUnavailable));
+}
+
+TEST_F(BulkMutatorTest, RetryPolicyMakesMutationFailureRetryable) {
+  // The mirror image of the previous test: PERMISSION_DENIED is permanent by
+  // default, but the policy configured by the application says otherwise, so
+  // the mutation stays pending for the next attempt.
+  BulkMutation mut(IdempotentMutation("r0"));
+
+  auto mock = std::make_shared<MockBigtableStub>();
+  EXPECT_CALL(*mock, MutateRows)
+      .WillOnce([this](auto context, auto const&,
+                       google::bigtable::v2::MutateRowsRequest const& request,
+                       auto const&) {
+        metadata_fixture_.SetServerMetadata(*context, {});
+        EXPECT_THAT(request, HasCorrectResourceNames());
+        auto stream = std::make_unique<MockMutateRowsStream>();
+        EXPECT_CALL(*stream, Read)
+            .WillOnce([](google::bigtable::v2::MutateRowsResponse* r) {
+              *r = MakeResponse({{0, grpc::StatusCode::PERMISSION_DENIED}});
+              return std::nullopt;
+            })
+            .WillOnce(Return(Status()));
+        return stream;
+      });
+
+  auto retry_policy = std::make_shared<MockDataRetryPolicy>();
+  EXPECT_CALL(*retry_policy, IsPermanentFailure)
+      .WillRepeatedly(Return(false));
+  auto policy = DefaultIdempotentMutationPolicy();
+  bigtable_internal::BulkMutator mutator(
+      kAppProfile, kTableName, *policy, std::move(mut),
+      std::make_shared<bigtable_internal::OperationContext>());
+
+  EXPECT_TRUE(mutator.HasPendingMutations());
+  bigtable_internal::NoopMutateRowsLimiter limiter;
+  auto status = mutator.MakeOneRequest(
+      *mock, limiter, Options{}.set<DataRetryPolicyOption>(retry_policy));
+  EXPECT_STATUS_OK(status);
+  EXPECT_TRUE(mutator.HasPendingMutations());
+  auto failures = std::move(mutator).OnRetryDone();
+  ASSERT_EQ(1UL, failures.size());
+  EXPECT_THAT(failures[0].status(), StatusIs(StatusCode::kPermissionDenied));
+}
+
 TEST_F(BulkMutatorTest, PermanentFailure) {
   // In this test we handle a recoverable and one unrecoverable failures.
   // Create a bulk mutation with two SetCell() mutations.