fix(storage): Set the idempotency token for async rewrites (googleapis#16114)

Author: kalragauri

google/cloud/storage/internal/async/rewriter_connection_impl.cc

@@ -49,11 +49,12 @@ RewriterConnectionImpl::Iterate() {
              current_->get<storage::RetryPolicyOption>()->clone(),
              current_->get<storage::BackoffPolicyOption>()->clone(),
              policy->RewriteObject(request_), cq_,
-             [stub = stub_](
+             [stub = stub_, id = invocation_id_generator_.MakeInvocationId()](
                  CompletionQueue& cq,
                  std::shared_ptr<grpc::ClientContext> context,
                  google::cloud::internal::ImmutableOptions options,
                  google::storage::v2::RewriteObjectRequest const& proto) {
+               AddIdempotencyToken(*context, id);
                return stub->AsyncRewriteObject(cq, std::move(context),
                                                std::move(options), proto);
              },

google/cloud/storage/internal/async/rewriter_connection_impl.h

@@ -20,6 +20,7 @@
 #include "google/cloud/storage/options.h"
 #include "google/cloud/completion_queue.h"
 #include "google/cloud/future.h"
+#include "google/cloud/internal/invocation_id_generator.h"
 #include "google/cloud/status_or.h"
 #include "google/cloud/version.h"
 #include "google/storage/v2/storage.pb.h"
@@ -53,6 +54,7 @@ class RewriterConnectionImpl
   std::shared_ptr<StorageStub> stub_;
   google::cloud::internal::ImmutableOptions current_;
   google::storage::v2::RewriteObjectRequest request_;
+  google::cloud::internal::InvocationIdGenerator invocation_id_generator_;
 };
 
 GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END

google/cloud/storage/internal/async/rewriter_connection_impl_test.cc

@@ -23,6 +23,7 @@
 #include "google/cloud/internal/background_threads_impl.h"
 #include "google/cloud/testing_util/async_sequencer.h"
 #include "google/cloud/testing_util/status_matchers.h"
+#include "google/cloud/testing_util/validate_metadata.h"
 #include <gmock/gmock.h>
 
 namespace google {
@@ -238,6 +239,65 @@ TEST(RewriterConnectionImplTest, TooManyTransients) {
   EXPECT_THAT(r1.get(), StatusIs(TransientError().status().code()));
 }
 
+TEST(RewriterConnectionImplTest, IterateReusesIdempotencyTokenOnRetry) {
+  google::cloud::testing_util::ValidateMetadataFixture
+      validate_metadata_fixture;
+  AsyncSequencer<bool> sequencer;
+  auto mock = std::make_shared<MockStorageStub>();
+  std::string first_token;
+
+  EXPECT_CALL(*mock, AsyncRewriteObject)
+      .WillOnce([&](auto&, auto const& context, auto const&,
+                    google::storage::v2::RewriteObjectRequest const&) {
+        auto metadata = validate_metadata_fixture.GetMetadata(*context);
+        auto l = metadata.find("x-goog-gcs-idempotency-token");
+        EXPECT_NE(l, metadata.end());
+        if (l != metadata.end()) {
+          first_token = l->second;
+          EXPECT_FALSE(first_token.empty());
+        }
+
+        return sequencer.PushBack("RewriteObject(1)").then([](auto) {
+          return TransientError();
+        });
+      })
+      .WillOnce([&](auto&, auto const& context, auto const&,
+                    google::storage::v2::RewriteObjectRequest const&) {
+        auto metadata = validate_metadata_fixture.GetMetadata(*context);
+        auto l = metadata.find("x-goog-gcs-idempotency-token");
+        EXPECT_NE(l, metadata.end());
+        if (l != metadata.end()) {
+          EXPECT_EQ(l->second, first_token);
+        }
+
+        return sequencer.PushBack("RewriteObject(2)").then([](auto) {
+          google::storage::v2::RewriteResponse response;
+          response.set_total_bytes_rewritten(1000);
+          response.set_object_size(3000);
+          response.set_rewrite_token("test-rewrite-token");
+          return make_status_or(response);
+        });
+      });
+
+  internal::AutomaticallyCreatedBackgroundThreads pool(1);
+  auto connection = std::make_shared<RewriterConnectionImpl>(
+      pool.cq(), std::move(mock), TestOptions(), MakeRequest());
+
+  auto r1 = connection->Iterate();
+  auto next = sequencer.PopFrontWithName();
+  EXPECT_EQ(next.second, "RewriteObject(1)");
+  next.first.set_value(true);
+  next = sequencer.PopFrontWithName();
+  EXPECT_EQ(next.second, "RewriteObject(2)");
+  next.first.set_value(true);
+  EXPECT_THAT(
+      r1.get(),
+      IsOkAndHolds(ResultOf(
+          "total bytes",
+          [](RewriteResponse const& v) { return v.total_bytes_rewritten(); },
+          1000)));
+}
+
 GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END
 }  // namespace storage_internal
 }  // namespace cloud