writer_connection_resumed_test.cc

// Copyright 2025 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include "google/cloud/storage/internal/async/writer_connection_resumed.h"
#include "google/cloud/storage/async/connection.h"
#include "google/cloud/storage/mocks/mock_async_writer_connection.h"
#include "google/cloud/storage/testing/canonical_errors.h"
#include "google/cloud/storage/testing/mock_hash_function.h"
#include "google/cloud/testing_util/async_sequencer.h"
#include "google/cloud/testing_util/is_proto_equal.h"
#include "google/cloud/testing_util/status_matchers.h"
#include <google/storage/v2/storage.pb.h>
#include <gmock/gmock.h>
#include <chrono>
#include <thread>

namespace google {
namespace cloud {
namespace storage_internal {
GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_BEGIN
namespace {

using ::google::cloud::storage::testing::canonical_errors::TransientError;
using ::google::cloud::storage_mocks::MockAsyncWriterConnection;
using ::google::cloud::testing_util::AsyncSequencer;
using ::google::cloud::testing_util::IsOkAndHolds;
using ::google::cloud::testing_util::IsProtoEqual;
using ::google::cloud::testing_util::StatusIs;
using ::testing::_;
using ::testing::Return;
using ::testing::VariantWith;

using MockFactory =
    ::testing::MockFunction<future<StatusOr<WriteObject::WriteResult>>(
        google::storage::v2::BidiWriteObjectRequest)>;

using MockStreamingRpc =
    ::testing::MockFunction<std::unique_ptr<WriteObject::StreamingRpc>()>;

absl::variant<std::int64_t, google::storage::v2::Object> MakePersistedState(
    std::int64_t persisted_size) {
  return persisted_size;
}

storage_experimental::WritePayload TestPayload(std::size_t n) {
  return storage_experimental::WritePayload(std::string(n, 'A'));
}

auto TestObject() {
  auto object = google::storage::v2::Object{};
  object.set_bucket("projects/_/buckets/test-bucket");
  object.set_name("test-object");
  return object;
}

TEST(WriteConnectionResumed, FinalizeEmpty) {
  AsyncSequencer<bool> sequencer;
  auto mock = std::make_unique<MockAsyncWriterConnection>();
  auto initial_request = google::storage::v2::BidiWriteObjectRequest{};
  auto first_response = google::storage::v2::BidiWriteObjectResponse{};

  EXPECT_CALL(*mock, UploadId).WillOnce(Return("test-upload-id"));
  EXPECT_CALL(*mock, PersistedState)
      .WillRepeatedly(Return(MakePersistedState(0)));
  EXPECT_CALL(*mock, Finalize).WillRepeatedly([&](auto) {
    return sequencer.PushBack("Finalize")
        .then([](auto f) -> StatusOr<google::storage::v2::Object> {
          if (!f.get()) return TransientError();
          return TestObject();
        });
  });
  MockFactory mock_factory;
  EXPECT_CALL(mock_factory, Call).Times(0);

  auto connection = MakeWriterConnectionResumed(
      mock_factory.AsStdFunction(), std::move(mock), initial_request, nullptr,
      first_response, Options{});
  EXPECT_EQ(connection->UploadId(), "test-upload-id");
  EXPECT_THAT(connection->PersistedState(), VariantWith<std::int64_t>(0));

  auto finalize = connection->Finalize({});
  auto next = sequencer.PopFrontWithName();
  EXPECT_EQ(next.second, "Finalize");
  next.first.set_value(true);
  EXPECT_THAT(finalize.get(), IsOkAndHolds(IsProtoEqual(TestObject())));
}

TEST(WriteConnectionResumed, FinalizedOnConstruction) {
  AsyncSequencer<bool> sequencer;
  auto initial_request = google::storage::v2::BidiWriteObjectRequest{};
  auto first_response = google::storage::v2::BidiWriteObjectResponse{};
  auto mock = std::make_unique<MockAsyncWriterConnection>();
  EXPECT_CALL(*mock, UploadId).WillRepeatedly(Return("test-upload-id"));
  EXPECT_CALL(*mock, PersistedState).WillRepeatedly(Return(TestObject()));
  EXPECT_CALL(*mock, Finalize(_)).WillOnce([&](auto) {
    return sequencer.PushBack("Finalize").then([](auto) {
      return StatusOr<google::storage::v2::Object>(TestObject());
    });
  });

  MockFactory mock_factory;
  EXPECT_CALL(mock_factory, Call).Times(0);

  auto connection = MakeWriterConnectionResumed(
      mock_factory.AsStdFunction(), std::move(mock), initial_request, nullptr,
      first_response, Options{});
  EXPECT_EQ(connection->UploadId(), "test-upload-id");
  EXPECT_THAT(
      connection->PersistedState(),
      VariantWith<google::storage::v2::Object>(IsProtoEqual(TestObject())));

  auto finalize = connection->Finalize({});

  EXPECT_FALSE(finalize.is_ready());
  auto next = sequencer.PopFrontWithName();
  EXPECT_EQ(next.second, "Finalize");
  next.first.set_value(true);
  EXPECT_TRUE(finalize.is_ready());
  EXPECT_THAT(finalize.get(), IsOkAndHolds(IsProtoEqual(TestObject())));
}

TEST(WriteConnectionResumed, Cancel) {
  AsyncSequencer<bool> sequencer;
  auto initial_request = google::storage::v2::BidiWriteObjectRequest{};
  auto first_response = google::storage::v2::BidiWriteObjectResponse{};
  auto mock = std::make_unique<MockAsyncWriterConnection>();
  EXPECT_CALL(*mock, UploadId).WillRepeatedly(Return("test-upload-id"));
  EXPECT_CALL(*mock, PersistedState)
      .WillRepeatedly(Return(MakePersistedState(0)));
  EXPECT_CALL(*mock, Flush).WillOnce([&](auto) {
    return sequencer.PushBack("Flush").then(
        [](auto) { return Status(StatusCode::kCancelled, "cancel"); });
  });
  EXPECT_CALL(*mock, Cancel).Times(1);

  MockFactory mock_factory;
  EXPECT_CALL(mock_factory, Call).Times(0);

  auto connection = MakeWriterConnectionResumed(
      mock_factory.AsStdFunction(), std::move(mock), initial_request, nullptr,
      first_response, Options{});

  auto write = connection->Write(TestPayload(64 * 1024));
  ASSERT_FALSE(write.is_ready());

  auto next = sequencer.PopFrontWithName();
  EXPECT_EQ(next.second, "Flush");
  connection->Cancel();
  next.first.set_value(true);

  ASSERT_TRUE(write.is_ready());
  auto status = write.get();
  EXPECT_THAT(status, StatusIs(StatusCode::kCancelled));
}

TEST(WriterConnectionResumed, FlushEmpty) {
  AsyncSequencer<bool> sequencer;
  auto initial_request = google::storage::v2::BidiWriteObjectRequest{};
  auto first_response = google::storage::v2::BidiWriteObjectResponse{};

  auto mock = std::make_unique<MockAsyncWriterConnection>();
  EXPECT_CALL(*mock, PersistedState)
      .WillRepeatedly(Return(MakePersistedState(0)));
  EXPECT_CALL(*mock, Flush).WillRepeatedly([&](auto const& p) {
    EXPECT_TRUE(p.payload().empty());
    return sequencer.PushBack("Flush").then([](auto f) {
      if (!f.get()) return TransientError();
      return Status{};
    });
  });
  EXPECT_CALL(*mock, Query).WillOnce([&]() {
    return sequencer.PushBack("Query").then(
        [](auto f) -> StatusOr<std::int64_t> {
          if (!f.get()) return TransientError();
          return 0;
        });
  });

  MockFactory mock_factory;
  EXPECT_CALL(mock_factory, Call).Times(0);

  auto connection = MakeWriterConnectionResumed(
      mock_factory.AsStdFunction(), std::move(mock), initial_request, nullptr,
      first_response, Options{});
  EXPECT_THAT(connection->PersistedState(), VariantWith<std::int64_t>(0));

  auto flush = connection->Flush({});
  auto next = sequencer.PopFrontWithName();
  EXPECT_EQ(next.second, "Flush");
  next.first.set_value(true);

  next = sequencer.PopFrontWithName();
  EXPECT_EQ(next.second, "Query");
  next.first.set_value(true);

  EXPECT_THAT(flush.get(), StatusIs(StatusCode::kOk));
}

TEST(WriteConnectionResumed, FlushNonEmpty) {
  AsyncSequencer<bool> sequencer;
  auto mock = std::make_unique<MockAsyncWriterConnection>();
  auto initial_request = google::storage::v2::BidiWriteObjectRequest{};
  auto first_response = google::storage::v2::BidiWriteObjectResponse{};
  auto const payload = TestPayload(1024);

  EXPECT_CALL(*mock, PersistedState)
      .WillRepeatedly(Return(MakePersistedState(0)));
  EXPECT_CALL(*mock, Flush)
      .WillOnce([&](auto const& p) {
        EXPECT_EQ(p.payload(), payload.payload());
        return sequencer.PushBack("Flush").then([](auto f) {
          if (!f.get()) return TransientError();
          return Status{};
        });
      })
      .WillOnce([&](auto const& p) {
        EXPECT_TRUE(p.payload().empty());
        return sequencer.PushBack("Flush").then([](auto f) {
          if (!f.get()) return TransientError();
          return Status{};
        });
      });
  EXPECT_CALL(*mock, Query).WillOnce([&]() {
    return sequencer.PushBack("Query").then(
        [](auto f) -> StatusOr<std::int64_t> {
          if (!f.get()) return TransientError();
          return 1024;
        });
  });

  MockFactory mock_factory;
  EXPECT_CALL(mock_factory, Call).Times(0);

  auto connection = MakeWriterConnectionResumed(
      mock_factory.AsStdFunction(), std::move(mock), initial_request, nullptr,
      first_response, Options{});
  EXPECT_THAT(connection->PersistedState(), VariantWith<std::int64_t>(0));

  auto write = connection->Write(payload);
  ASSERT_FALSE(write.is_ready());

  auto flush = connection->Flush({});
  ASSERT_FALSE(flush.is_ready());

  auto next = sequencer.PopFrontWithName();
  EXPECT_EQ(next.second, "Flush");
  next.first.set_value(true);

  next = sequencer.PopFrontWithName();
  EXPECT_EQ(next.second, "Query");
  next.first.set_value(true);

  EXPECT_TRUE(flush.is_ready());
  EXPECT_THAT(flush.get(), StatusIs(StatusCode::kOk));

  EXPECT_TRUE(write.is_ready());
  EXPECT_THAT(write.get(), StatusIs(StatusCode::kOk));
}

TEST(WriteConnectionResumed, ResumeUsesWriteObjectSpecFromInitialRequest) {
  AsyncSequencer<bool> sequencer;
  auto mock = std::make_unique<MockAsyncWriterConnection>();
  auto initial_request = google::storage::v2::BidiWriteObjectRequest{};
  initial_request.mutable_write_object_spec()->mutable_resource()->set_bucket(
      "projects/_/buckets/test-bucket");
  initial_request.mutable_write_object_spec()->mutable_resource()->set_name(
      "test-object");

  google::storage::v2::BidiWriteObjectResponse first_response;
  first_response.mutable_write_handle()->set_handle("test-handle");
  first_response.mutable_resource()->set_generation(12345);

  EXPECT_CALL(*mock, PersistedState)
      .WillRepeatedly(Return(MakePersistedState(0)));

  // Expect Flush to be called because flush_ is true by default.
  // Make it fail to trigger Resume().
  EXPECT_CALL(*mock, Flush(_)).WillOnce([&](auto) {
    return sequencer.PushBack("Flush").then([](auto f) {
      if (f.get()) return google::cloud::Status{};  // Should not be true
      return TransientError();
    });
  });

  MockFactory mock_factory;
  google::storage::v2::BidiWriteObjectRequest captured_request;
  EXPECT_CALL(mock_factory, Call(_))
      .WillOnce([&](google::storage::v2::BidiWriteObjectRequest request) {
        captured_request = std::move(request);
        return sequencer.PushBack("Factory").then([](auto) {
          return StatusOr<WriteObject::WriteResult>(
              internal::AbortedError("stop test", GCP_ERROR_INFO()));
        });
      });

  auto connection = MakeWriterConnectionResumed(
      mock_factory.AsStdFunction(), std::move(mock), initial_request, nullptr,
      first_response, Options{});

  // This will call FlushStep -> mock->Flush()
  auto write = connection->Write(TestPayload(1));
  ASSERT_FALSE(write.is_ready());

  // Trigger the Flush error in the mock
  auto next = sequencer.PopFrontWithName();
  EXPECT_EQ(next.second, "Flush");
  next.first.set_value(false);  // This makes the lambda return TransientError

  // The error in OnFlush triggers Resume(), which calls the mock_factory.
  // Allow the factory callback to proceed.
  next = sequencer.PopFrontWithName();
  EXPECT_EQ(next.second, "Factory");
  next.first.set_value(true);

  // The write future should now be ready, containing the error from the
  // factory.
  EXPECT_THAT(write.get(), StatusIs(StatusCode::kAborted));

  EXPECT_FALSE(captured_request.has_write_object_spec());
  EXPECT_TRUE(captured_request.has_append_object_spec());
  EXPECT_TRUE(captured_request.append_object_spec().has_write_handle());
  EXPECT_EQ(captured_request.append_object_spec().write_handle().handle(),
            "test-handle");
  EXPECT_EQ(captured_request.append_object_spec().generation(), 12345);
  EXPECT_EQ(captured_request.append_object_spec().object(), "test-object");
  EXPECT_EQ(captured_request.append_object_spec().bucket(),
            "projects/_/buckets/test-bucket");
}

TEST(WriteConnectionResumed, ResumeUsesAppendObjectSpecFromInitialRequest) {
  AsyncSequencer<bool> sequencer;
  auto mock = std::make_unique<MockAsyncWriterConnection>();
  auto initial_request = google::storage::v2::BidiWriteObjectRequest{};
  initial_request.mutable_append_object_spec()->set_bucket(
      "projects/_/buckets/test-bucket");
  initial_request.mutable_append_object_spec()->set_object("test-object");

  google::storage::v2::BidiWriteObjectResponse first_response;
  first_response.mutable_resource()->set_generation(12345);

  EXPECT_CALL(*mock, PersistedState)
      .WillRepeatedly(Return(MakePersistedState(0)));

  // Expect Flush to be called because flush_ is true by default.
  // Make it fail to trigger Resume().
  EXPECT_CALL(*mock, Flush(_)).WillOnce([&](auto) {
    return sequencer.PushBack("Flush").then([](auto f) {
      if (f.get()) return google::cloud::Status{};  // Should not be true
      return TransientError();
    });
  });

  MockFactory mock_factory;
  google::storage::v2::BidiWriteObjectRequest captured_request;
  EXPECT_CALL(mock_factory, Call(_))
      .WillOnce([&](google::storage::v2::BidiWriteObjectRequest request) {
        captured_request = std::move(request);
        return sequencer.PushBack("Factory").then([](auto) {
          return StatusOr<WriteObject::WriteResult>(
              internal::AbortedError("stop test", GCP_ERROR_INFO()));
        });
      });

  auto connection = MakeWriterConnectionResumed(
      mock_factory.AsStdFunction(), std::move(mock), initial_request, nullptr,
      first_response, Options{});

  // This will call FlushStep -> mock->Flush()
  auto write = connection->Write(TestPayload(1));
  ASSERT_FALSE(write.is_ready());

  // Trigger the Flush error in the mock
  auto next = sequencer.PopFrontWithName();
  EXPECT_EQ(next.second, "Flush");
  next.first.set_value(false);  // This makes the lambda return TransientError

  // The error in OnFlush triggers Resume(), which calls the mock_factory.
  // Allow the factory callback to proceed.
  next = sequencer.PopFrontWithName();
  EXPECT_EQ(next.second, "Factory");
  next.first.set_value(true);

  // The write future should now be ready, containing the error from the
  // factory.
  EXPECT_THAT(write.get(), StatusIs(StatusCode::kAborted));

  EXPECT_FALSE(captured_request.has_write_object_spec());
  EXPECT_TRUE(captured_request.has_append_object_spec());
  EXPECT_FALSE(captured_request.append_object_spec().has_write_handle());
  EXPECT_EQ(captured_request.append_object_spec().generation(), 12345);
  EXPECT_EQ(captured_request.append_object_spec().object(), "test-object");
  EXPECT_EQ(captured_request.append_object_spec().bucket(),
            "projects/_/buckets/test-bucket");
}

TEST(WriteConnectionResumed, NoConcurrentWritesWhenFlushAndWriteRace) {
  AsyncSequencer<bool> sequencer;
  auto mock = std::make_unique<MockAsyncWriterConnection>();
  auto initial_request = google::storage::v2::BidiWriteObjectRequest{};
  auto first_response = google::storage::v2::BidiWriteObjectResponse{};

  EXPECT_CALL(*mock, PersistedState)
      .WillRepeatedly(Return(MakePersistedState(0)));
  EXPECT_CALL(*mock, Flush(_)).WillRepeatedly([&](auto) {
    return sequencer.PushBack("Flush").then([](auto f) {
      if (!f.get()) return TransientError();
      return Status{};
    });
  });
  EXPECT_CALL(*mock, Query).WillOnce([&]() {
    return sequencer.PushBack("Query").then([](auto f) -> StatusOr<std::int64_t> {
      if (!f.get()) return TransientError();
      return 0;
    });
  });

  // Make Write detect concurrent invocations. If two writes run concurrently
  // the compare_exchange will fail and the test will fail.
  std::atomic<bool> in_write{false};
  EXPECT_CALL(*mock, Write(_))
      .WillRepeatedly([&](auto) {
        bool expected = false;
        EXPECT_TRUE(in_write.compare_exchange_strong(expected, true));
        // Simulate some work that allows a concurrent write to attempt to run.
        std::this_thread::sleep_for(std::chrono::milliseconds(50));
        in_write.store(false);
        return make_ready_future(Status{});
      });

  MockFactory mock_factory;
  EXPECT_CALL(mock_factory, Call).Times(0);

  auto connection = MakeWriterConnectionResumed(
      mock_factory.AsStdFunction(), std::move(mock), initial_request, nullptr,
      first_response, Options{});

  // Start a flush which will call impl->Flush() and block.
  auto flush_future = connection->Flush({});
  // Allow the Flush to complete, this will schedule a Query (but Query will
  // remain blocked until we pop it).
  auto next = sequencer.PopFrontWithName();
  EXPECT_EQ(next.second, "Flush");
  next.first.set_value(true);

  // Immediately perform a user Write after the flush completed but before
  // Query completes. This can race with the OnQuery-driven write.
  auto write_future = connection->Write(TestPayload(1024));

  // Now allow the Query to complete; OnQuery may schedule a write.
  next = sequencer.PopFrontWithName();
  EXPECT_EQ(next.second, "Query");
  next.first.set_value(true);

  // Wait for both futures to complete with a timeout to avoid indefinite hang.
  auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(2);
  while (!write_future.is_ready() && std::chrono::steady_clock::now() < deadline) {
    std::this_thread::sleep_for(std::chrono::milliseconds(5));
  }
  deadline = std::chrono::steady_clock::now() + std::chrono::seconds(2);
  while (!flush_future.is_ready() && std::chrono::steady_clock::now() < deadline) {
    std::this_thread::sleep_for(std::chrono::milliseconds(5));
  }

  ASSERT_TRUE(write_future.is_ready());
  ASSERT_TRUE(flush_future.is_ready());

  // Both futures should complete successfully.
  EXPECT_THAT(write_future.get(), StatusIs(StatusCode::kOk));
  EXPECT_THAT(flush_future.get(), StatusIs(StatusCode::kOk));
}


TEST(WriteConnectionResumed, WriteHandleAssignmentAfterResume) {
  struct {
    bool use_write_object_spec;
    std::string bucket, object, handle;
    std::int64_t generation;
  } cases[] = {
      {false, "projects/_/buckets/test-bucket", "test-object",
       "expected-handle", 12345},
      {true, "bucket1", "object1", "handle1", 111},
      {false, "bucket2", "object2", "handle2", 222},
  };

  for (auto const& tc : cases) {
    AsyncSequencer<bool> sequencer;
    auto mock = std::make_unique<MockAsyncWriterConnection>();
    google::storage::v2::BidiWriteObjectRequest req;
    if (tc.use_write_object_spec) {
      req.mutable_write_object_spec()->mutable_resource()->set_bucket(
          tc.bucket);
      req.mutable_write_object_spec()->mutable_resource()->set_name(tc.object);
    } else {
      req.mutable_append_object_spec()->set_bucket(tc.bucket);
      req.mutable_append_object_spec()->set_object(tc.object);
    }
    google::storage::v2::BidiWriteObjectResponse resp;
    resp.mutable_write_handle()->set_handle(tc.handle);
    resp.mutable_resource()->set_generation(tc.generation);

    EXPECT_CALL(*mock, PersistedState)
        .WillRepeatedly(Return(MakePersistedState(0)));
    EXPECT_CALL(*mock, Flush(_)).WillOnce([&](auto) {
      return sequencer.PushBack("Flush").then([](auto f) {
        if (f.get()) return google::cloud::Status{};
        return TransientError();
      });
    });

    MockFactory mock_factory;
    google::storage::v2::BidiWriteObjectRequest captured;
    EXPECT_CALL(mock_factory, Call(_))
        .WillOnce([&](google::storage::v2::BidiWriteObjectRequest r) {
          captured = std::move(r);
          return sequencer.PushBack("Factory").then([](auto) {
            return StatusOr<WriteObject::WriteResult>(
                internal::AbortedError("stop test", GCP_ERROR_INFO()));
          });
        });

    auto conn = MakeWriterConnectionResumed(mock_factory.AsStdFunction(),
                                            std::move(mock), req, nullptr, resp,
                                            Options{});
    auto write = conn->Write(TestPayload(1));
    sequencer.PopFrontWithName().first.set_value(false);
    sequencer.PopFrontWithName().first.set_value(true);

    EXPECT_THAT(write.get(), StatusIs(StatusCode::kAborted));
    EXPECT_TRUE(captured.has_append_object_spec());
    EXPECT_EQ(captured.append_object_spec().bucket(), tc.bucket);
    EXPECT_EQ(captured.append_object_spec().object(), tc.object);
    EXPECT_EQ(captured.append_object_spec().generation(), tc.generation);
    EXPECT_TRUE(captured.append_object_spec().has_write_handle());
    EXPECT_EQ(captured.append_object_spec().write_handle().handle(), tc.handle);
  }
}

}  // namespace
GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END
}  // namespace storage_internal
}  // namespace cloud
}  // namespace google