diff --git a/google/cloud/storage/internal/async/writer_connection_buffered.cc b/google/cloud/storage/internal/async/writer_connection_buffered.cc index d00ead1f2cf40..5d13e51526f80 100644 --- a/google/cloud/storage/internal/async/writer_connection_buffered.cc +++ b/google/cloud/storage/internal/async/writer_connection_buffered.cc @@ -115,6 +115,16 @@ class AsyncWriterConnectionBufferedState future> Finalize( storage::WritePayload const& p) { std::unique_lock lk(mu_); + if (!finalized_future_.valid()) { + return make_ready_future(StatusOr( + internal::FailedPreconditionError("Finalize() already called", + GCP_ERROR_INFO()))); + } + if (!closed_future_.valid()) { + return make_ready_future(StatusOr( + internal::FailedPreconditionError( + "Finalize() cannot be called after Close()", GCP_ERROR_INFO()))); + } resend_buffer_.Append(WritePayloadImpl::GetImpl(p)); finalize_ = true; HandleNewData(std::move(lk), true); @@ -124,10 +134,14 @@ class AsyncWriterConnectionBufferedState future Close(storage::WritePayload const& p) { std::unique_lock lk(mu_); - if (close_ || closed_promise_completed_) { + if (!closed_future_.valid()) { return make_ready_future(internal::FailedPreconditionError( "Close() already called", GCP_ERROR_INFO())); } + if (!finalized_future_.valid()) { + return make_ready_future(internal::FailedPreconditionError( + "Close() cannot be called after Finalize()", GCP_ERROR_INFO())); + } resend_buffer_.Append(WritePayloadImpl::GetImpl(p)); close_ = true; // Force flush to drain the buffer first. diff --git a/google/cloud/storage/internal/async/writer_connection_buffered_test.cc b/google/cloud/storage/internal/async/writer_connection_buffered_test.cc index 384e685863e7a..3a37d8ae15a61 100644 --- a/google/cloud/storage/internal/async/writer_connection_buffered_test.cc +++ b/google/cloud/storage/internal/async/writer_connection_buffered_test.cc @@ -1665,6 +1665,66 @@ TEST(WriteConnectionBuffered, ResetWriteOffsetOnResume) { EXPECT_STATUS_OK(write.get()); } +TEST(WriteConnectionBuffered, CloseAfterFinalizeFails) { + auto mock = std::make_unique(); + EXPECT_CALL(*mock, UploadId).WillRepeatedly(Return("test-upload-id")); + EXPECT_CALL(*mock, PersistedState) + .WillRepeatedly(Return(MakePersistedState(0))); + EXPECT_CALL(*mock, Finalize).WillOnce([](auto) { + return make_ready_future(make_status_or(google::storage::v2::Object{})); + }); + + MockFactory mock_factory; + auto connection = MakeWriterConnectionBuffered( + mock_factory.AsStdFunction(), std::move(mock), TestOptions()); + + auto finalize = connection->Finalize({}); + auto close = connection->Close({}); + + EXPECT_STATUS_OK(finalize.get()); + EXPECT_THAT(close.get(), StatusIs(StatusCode::kFailedPrecondition)); +} + +TEST(WriteConnectionBuffered, FinalizeAfterCloseFails) { + auto mock = std::make_unique(); + EXPECT_CALL(*mock, UploadId).WillRepeatedly(Return("test-upload-id")); + EXPECT_CALL(*mock, PersistedState) + .WillRepeatedly(Return(MakePersistedState(0))); + EXPECT_CALL(*mock, Close).WillOnce([](auto) { + return make_ready_future(Status{}); + }); + + MockFactory mock_factory; + auto connection = MakeWriterConnectionBuffered( + mock_factory.AsStdFunction(), std::move(mock), TestOptions()); + + auto close = connection->Close({}); + auto finalize = connection->Finalize({}); + + EXPECT_STATUS_OK(close.get()); + EXPECT_THAT(finalize.get(), StatusIs(StatusCode::kFailedPrecondition)); +} + +TEST(WriteConnectionBuffered, DuplicateFinalizeFails) { + auto mock = std::make_unique(); + EXPECT_CALL(*mock, UploadId).WillRepeatedly(Return("test-upload-id")); + EXPECT_CALL(*mock, PersistedState) + .WillRepeatedly(Return(MakePersistedState(0))); + EXPECT_CALL(*mock, Finalize).WillOnce([](auto) { + return make_ready_future(make_status_or(google::storage::v2::Object{})); + }); + + MockFactory mock_factory; + auto connection = MakeWriterConnectionBuffered( + mock_factory.AsStdFunction(), std::move(mock), TestOptions()); + + auto finalize1 = connection->Finalize({}); + auto finalize2 = connection->Finalize({}); + + EXPECT_STATUS_OK(finalize1.get()); + EXPECT_THAT(finalize2.get(), StatusIs(StatusCode::kFailedPrecondition)); +} + } // namespace GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END } // namespace storage_internal diff --git a/google/cloud/storage/internal/async/writer_connection_resumed.cc b/google/cloud/storage/internal/async/writer_connection_resumed.cc index 6d96e523a25d8..27a3e21620f17 100644 --- a/google/cloud/storage/internal/async/writer_connection_resumed.cc +++ b/google/cloud/storage/internal/async/writer_connection_resumed.cc @@ -125,6 +125,16 @@ class AsyncWriterConnectionResumedState future> Finalize( storage::WritePayload const& p) { std::unique_lock lk(mu_); + if (!finalized_future_.valid()) { + return make_ready_future(StatusOr( + internal::FailedPreconditionError("Finalize() already called", + GCP_ERROR_INFO()))); + } + if (!closed_future_.valid()) { + return make_ready_future(StatusOr( + internal::FailedPreconditionError( + "Finalize() cannot be called after Close()", GCP_ERROR_INFO()))); + } resend_buffer_.Append(WritePayloadImpl::GetImpl(p)); finalize_ = true; HandleNewData(std::move(lk)); @@ -148,10 +158,14 @@ class AsyncWriterConnectionResumedState future Close(storage::WritePayload const& p) { std::unique_lock lk(mu_); - if (close_ || closed_promise_completed_) { + if (!closed_future_.valid()) { return make_ready_future(internal::FailedPreconditionError( "Close() already called", GCP_ERROR_INFO())); } + if (!finalized_future_.valid()) { + return make_ready_future(internal::FailedPreconditionError( + "Close() cannot be called after Finalize()", GCP_ERROR_INFO())); + } resend_buffer_.Append(WritePayloadImpl::GetImpl(p)); close_ = true; // Force flush to drain the buffer first. diff --git a/google/cloud/storage/internal/async/writer_connection_resumed_test.cc b/google/cloud/storage/internal/async/writer_connection_resumed_test.cc index e510ddad65e3c..3dd44a1a80736 100644 --- a/google/cloud/storage/internal/async/writer_connection_resumed_test.cc +++ b/google/cloud/storage/internal/async/writer_connection_resumed_test.cc @@ -1198,6 +1198,78 @@ TEST(WriterConnectionResumed, CloseFailsAndResumeSucceedsAndFinalized) { EXPECT_STATUS_OK(close.get()); } +TEST(WriterConnectionResumed, CloseAfterFinalizeFails) { + auto mock = std::make_unique(); + auto initial_request = google::storage::v2::BidiWriteObjectRequest{}; + auto first_response = google::storage::v2::BidiWriteObjectResponse{}; + + EXPECT_CALL(*mock, UploadId).WillRepeatedly(Return("test-upload-id")); + EXPECT_CALL(*mock, PersistedState) + .WillRepeatedly(Return(MakePersistedState(0))); + EXPECT_CALL(*mock, Finalize).WillOnce([](auto) { + return make_ready_future(make_status_or(google::storage::v2::Object{})); + }); + + MockFactory mock_factory; + auto connection = MakeWriterConnectionResumed( + mock_factory.AsStdFunction(), std::move(mock), initial_request, nullptr, + first_response, Options{}); + + auto finalize = connection->Finalize({}); + auto close = connection->Close({}); + + EXPECT_STATUS_OK(finalize.get()); + EXPECT_THAT(close.get(), StatusIs(StatusCode::kFailedPrecondition)); +} + +TEST(WriterConnectionResumed, FinalizeAfterCloseFails) { + auto mock = std::make_unique(); + auto initial_request = google::storage::v2::BidiWriteObjectRequest{}; + auto first_response = google::storage::v2::BidiWriteObjectResponse{}; + + EXPECT_CALL(*mock, UploadId).WillRepeatedly(Return("test-upload-id")); + EXPECT_CALL(*mock, PersistedState) + .WillRepeatedly(Return(MakePersistedState(0))); + EXPECT_CALL(*mock, Close).WillOnce([](auto) { + return make_ready_future(Status{}); + }); + + MockFactory mock_factory; + auto connection = MakeWriterConnectionResumed( + mock_factory.AsStdFunction(), std::move(mock), initial_request, nullptr, + first_response, Options{}); + + auto close = connection->Close({}); + auto finalize = connection->Finalize({}); + + EXPECT_STATUS_OK(close.get()); + EXPECT_THAT(finalize.get(), StatusIs(StatusCode::kFailedPrecondition)); +} + +TEST(WriterConnectionResumed, DuplicateFinalizeFails) { + auto mock = std::make_unique(); + auto initial_request = google::storage::v2::BidiWriteObjectRequest{}; + auto first_response = google::storage::v2::BidiWriteObjectResponse{}; + + EXPECT_CALL(*mock, UploadId).WillRepeatedly(Return("test-upload-id")); + EXPECT_CALL(*mock, PersistedState) + .WillRepeatedly(Return(MakePersistedState(0))); + EXPECT_CALL(*mock, Finalize).WillOnce([](auto) { + return make_ready_future(make_status_or(google::storage::v2::Object{})); + }); + + MockFactory mock_factory; + auto connection = MakeWriterConnectionResumed( + mock_factory.AsStdFunction(), std::move(mock), initial_request, nullptr, + first_response, Options{}); + + auto finalize1 = connection->Finalize({}); + auto finalize2 = connection->Finalize({}); + + EXPECT_STATUS_OK(finalize1.get()); + EXPECT_THAT(finalize2.get(), StatusIs(StatusCode::kFailedPrecondition)); +} + } // namespace GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END } // namespace storage_internal