Skip to content

Commit f70bcf5

Browse files
kalragaurikalragauri
authored
fix(storage): enforce mutual exclusion between Close() and Finalize() in async writers (googleapis#16211)
1 parent a52ae45 commit f70bcf5

4 files changed

Lines changed: 192 additions & 2 deletions

File tree

‎google/cloud/storage/internal/async/writer_connection_buffered.cc‎

Lines changed: 30 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -115,6 +115,17 @@ class AsyncWriterConnectionBufferedState
115115
future<StatusOr<google::storage::v2::Object>> Finalize(
116116
storage::WritePayload const& p) {
117117
std::unique_lock<std::mutex> lk(mu_);
118+
if (finalize_called_) {
119+
return make_ready_future(StatusOr<google::storage::v2::Object>(
120+
internal::FailedPreconditionError("Finalize() already called",
121+
GCP_ERROR_INFO())));
122+
}
123+
if (close_called_) {
124+
return make_ready_future(StatusOr<google::storage::v2::Object>(
125+
internal::FailedPreconditionError(
126+
"Finalize() cannot be called after Close()", GCP_ERROR_INFO())));
127+
}
128+
finalize_called_ = true;
118129
resend_buffer_.Append(WritePayloadImpl::GetImpl(p));
119130
finalize_ = true;
120131
HandleNewData(std::move(lk), true);
@@ -124,10 +135,15 @@ class AsyncWriterConnectionBufferedState
124135

125136
future<Status> Close(storage::WritePayload const& p) {
126137
std::unique_lock<std::mutex> lk(mu_);
127-
if (close_ || closed_promise_completed_) {
138+
if (close_called_) {
128139
return make_ready_future(internal::FailedPreconditionError(
129140
"Close() already called", GCP_ERROR_INFO()));
130141
}
142+
if (finalize_called_) {
143+
return make_ready_future(internal::FailedPreconditionError(
144+
"Close() cannot be called after Finalize()", GCP_ERROR_INFO()));
145+
}
146+
close_called_ = true;
131147
resend_buffer_.Append(WritePayloadImpl::GetImpl(p));
132148
close_ = true;
133149
// Force flush to drain the buffer first.
@@ -703,6 +719,19 @@ class AsyncWriterConnectionBufferedState
703719

704720
// True if the resume loop is running. Prevents re-entry.
705721
bool resuming_ = false;
722+
723+
// Tracks if the `Finalize()` method has been called.
724+
// This is distinct from `finalize_` (which tracks if a finalize operation is
725+
// pending in the write loop) and `finalized_promise_completed_` (which tracks
726+
// if the upload is finalized).
727+
// We need this separate field to support the case where the upload is
728+
// finalized on construction; in that case, the promise is completed
729+
// immediately, but the application is still allowed to call `Finalize()` once
730+
// to retrieve the object.
731+
bool finalize_called_ = false;
732+
733+
// Tracks if the `Close()` method has been called.
734+
bool close_called_ = false;
706735
};
707736

708737
/**

‎google/cloud/storage/internal/async/writer_connection_buffered_test.cc‎

Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1665,6 +1665,66 @@ TEST(WriteConnectionBuffered, ResetWriteOffsetOnResume) {
16651665
EXPECT_STATUS_OK(write.get());
16661666
}
16671667

1668+
TEST(WriteConnectionBuffered, CloseAfterFinalizeFails) {
1669+
auto mock = std::make_unique<MockAsyncWriterConnection>();
1670+
EXPECT_CALL(*mock, UploadId).WillRepeatedly(Return("test-upload-id"));
1671+
EXPECT_CALL(*mock, PersistedState)
1672+
.WillRepeatedly(Return(MakePersistedState(0)));
1673+
EXPECT_CALL(*mock, Finalize).WillOnce([](auto) {
1674+
return make_ready_future(make_status_or(google::storage::v2::Object{}));
1675+
});
1676+
1677+
MockFactory mock_factory;
1678+
auto connection = MakeWriterConnectionBuffered(
1679+
mock_factory.AsStdFunction(), std::move(mock), TestOptions());
1680+
1681+
auto finalize = connection->Finalize({});
1682+
auto close = connection->Close({});
1683+
1684+
EXPECT_STATUS_OK(finalize.get());
1685+
EXPECT_THAT(close.get(), StatusIs(StatusCode::kFailedPrecondition));
1686+
}
1687+
1688+
TEST(WriteConnectionBuffered, FinalizeAfterCloseFails) {
1689+
auto mock = std::make_unique<MockAsyncWriterConnection>();
1690+
EXPECT_CALL(*mock, UploadId).WillRepeatedly(Return("test-upload-id"));
1691+
EXPECT_CALL(*mock, PersistedState)
1692+
.WillRepeatedly(Return(MakePersistedState(0)));
1693+
EXPECT_CALL(*mock, Close).WillOnce([](auto) {
1694+
return make_ready_future(Status{});
1695+
});
1696+
1697+
MockFactory mock_factory;
1698+
auto connection = MakeWriterConnectionBuffered(
1699+
mock_factory.AsStdFunction(), std::move(mock), TestOptions());
1700+
1701+
auto close = connection->Close({});
1702+
auto finalize = connection->Finalize({});
1703+
1704+
EXPECT_STATUS_OK(close.get());
1705+
EXPECT_THAT(finalize.get(), StatusIs(StatusCode::kFailedPrecondition));
1706+
}
1707+
1708+
TEST(WriteConnectionBuffered, DuplicateFinalizeFails) {
1709+
auto mock = std::make_unique<MockAsyncWriterConnection>();
1710+
EXPECT_CALL(*mock, UploadId).WillRepeatedly(Return("test-upload-id"));
1711+
EXPECT_CALL(*mock, PersistedState)
1712+
.WillRepeatedly(Return(MakePersistedState(0)));
1713+
EXPECT_CALL(*mock, Finalize).WillOnce([](auto) {
1714+
return make_ready_future(make_status_or(google::storage::v2::Object{}));
1715+
});
1716+
1717+
MockFactory mock_factory;
1718+
auto connection = MakeWriterConnectionBuffered(
1719+
mock_factory.AsStdFunction(), std::move(mock), TestOptions());
1720+
1721+
auto finalize1 = connection->Finalize({});
1722+
auto finalize2 = connection->Finalize({});
1723+
1724+
EXPECT_STATUS_OK(finalize1.get());
1725+
EXPECT_THAT(finalize2.get(), StatusIs(StatusCode::kFailedPrecondition));
1726+
}
1727+
16681728
} // namespace
16691729
GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END
16701730
} // namespace storage_internal

‎google/cloud/storage/internal/async/writer_connection_resumed.cc‎

Lines changed: 30 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -125,6 +125,17 @@ class AsyncWriterConnectionResumedState
125125
future<StatusOr<google::storage::v2::Object>> Finalize(
126126
storage::WritePayload const& p) {
127127
std::unique_lock<std::mutex> lk(mu_);
128+
if (finalize_called_) {
129+
return make_ready_future(StatusOr<google::storage::v2::Object>(
130+
internal::FailedPreconditionError("Finalize() already called",
131+
GCP_ERROR_INFO())));
132+
}
133+
if (close_called_) {
134+
return make_ready_future(StatusOr<google::storage::v2::Object>(
135+
internal::FailedPreconditionError(
136+
"Finalize() cannot be called after Close()", GCP_ERROR_INFO())));
137+
}
138+
finalize_called_ = true;
128139
resend_buffer_.Append(WritePayloadImpl::GetImpl(p));
129140
finalize_ = true;
130141
HandleNewData(std::move(lk));
@@ -148,10 +159,15 @@ class AsyncWriterConnectionResumedState
148159

149160
future<Status> Close(storage::WritePayload const& p) {
150161
std::unique_lock<std::mutex> lk(mu_);
151-
if (close_ || closed_promise_completed_) {
162+
if (close_called_) {
152163
return make_ready_future(internal::FailedPreconditionError(
153164
"Close() already called", GCP_ERROR_INFO()));
154165
}
166+
if (finalize_called_) {
167+
return make_ready_future(internal::FailedPreconditionError(
168+
"Close() cannot be called after Finalize()", GCP_ERROR_INFO()));
169+
}
170+
close_called_ = true;
155171
resend_buffer_.Append(WritePayloadImpl::GetImpl(p));
156172
close_ = true;
157173
// Force flush to drain the buffer first.
@@ -771,6 +787,19 @@ class AsyncWriterConnectionResumedState
771787

772788
// Track the latest write handle seen in responses.
773789
absl::optional<google::storage::v2::BidiWriteHandle> latest_write_handle_;
790+
791+
// Tracks if the `Finalize()` method has been called.
792+
// This is distinct from `finalize_` (which tracks if a finalize operation is
793+
// pending in the write loop) and `finalized_promise_completed_` (which tracks
794+
// if the upload is finalized).
795+
// We need this separate field to support the case where the upload is
796+
// finalized on construction; in that case, the promise is completed
797+
// immediately, but the application is still allowed to call `Finalize()` once
798+
// to retrieve the object.
799+
bool finalize_called_ = false;
800+
801+
// Tracks if the `Close()` method has been called.
802+
bool close_called_ = false;
774803
};
775804

776805
/**

‎google/cloud/storage/internal/async/writer_connection_resumed_test.cc‎

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1198,6 +1198,78 @@ TEST(WriterConnectionResumed, CloseFailsAndResumeSucceedsAndFinalized) {
11981198
EXPECT_STATUS_OK(close.get());
11991199
}
12001200

1201+
TEST(WriterConnectionResumed, CloseAfterFinalizeFails) {
1202+
auto mock = std::make_unique<MockAsyncWriterConnection>();
1203+
auto initial_request = google::storage::v2::BidiWriteObjectRequest{};
1204+
auto first_response = google::storage::v2::BidiWriteObjectResponse{};
1205+
1206+
EXPECT_CALL(*mock, UploadId).WillRepeatedly(Return("test-upload-id"));
1207+
EXPECT_CALL(*mock, PersistedState)
1208+
.WillRepeatedly(Return(MakePersistedState(0)));
1209+
EXPECT_CALL(*mock, Finalize).WillOnce([](auto) {
1210+
return make_ready_future(make_status_or(google::storage::v2::Object{}));
1211+
});
1212+
1213+
MockFactory mock_factory;
1214+
auto connection = MakeWriterConnectionResumed(
1215+
mock_factory.AsStdFunction(), std::move(mock), initial_request, nullptr,
1216+
first_response, Options{});
1217+
1218+
auto finalize = connection->Finalize({});
1219+
auto close = connection->Close({});
1220+
1221+
EXPECT_STATUS_OK(finalize.get());
1222+
EXPECT_THAT(close.get(), StatusIs(StatusCode::kFailedPrecondition));
1223+
}
1224+
1225+
TEST(WriterConnectionResumed, FinalizeAfterCloseFails) {
1226+
auto mock = std::make_unique<MockAsyncWriterConnection>();
1227+
auto initial_request = google::storage::v2::BidiWriteObjectRequest{};
1228+
auto first_response = google::storage::v2::BidiWriteObjectResponse{};
1229+
1230+
EXPECT_CALL(*mock, UploadId).WillRepeatedly(Return("test-upload-id"));
1231+
EXPECT_CALL(*mock, PersistedState)
1232+
.WillRepeatedly(Return(MakePersistedState(0)));
1233+
EXPECT_CALL(*mock, Close).WillOnce([](auto) {
1234+
return make_ready_future(Status{});
1235+
});
1236+
1237+
MockFactory mock_factory;
1238+
auto connection = MakeWriterConnectionResumed(
1239+
mock_factory.AsStdFunction(), std::move(mock), initial_request, nullptr,
1240+
first_response, Options{});
1241+
1242+
auto close = connection->Close({});
1243+
auto finalize = connection->Finalize({});
1244+
1245+
EXPECT_STATUS_OK(close.get());
1246+
EXPECT_THAT(finalize.get(), StatusIs(StatusCode::kFailedPrecondition));
1247+
}
1248+
1249+
TEST(WriterConnectionResumed, DuplicateFinalizeFails) {
1250+
auto mock = std::make_unique<MockAsyncWriterConnection>();
1251+
auto initial_request = google::storage::v2::BidiWriteObjectRequest{};
1252+
auto first_response = google::storage::v2::BidiWriteObjectResponse{};
1253+
1254+
EXPECT_CALL(*mock, UploadId).WillRepeatedly(Return("test-upload-id"));
1255+
EXPECT_CALL(*mock, PersistedState)
1256+
.WillRepeatedly(Return(MakePersistedState(0)));
1257+
EXPECT_CALL(*mock, Finalize).WillOnce([](auto) {
1258+
return make_ready_future(make_status_or(google::storage::v2::Object{}));
1259+
});
1260+
1261+
MockFactory mock_factory;
1262+
auto connection = MakeWriterConnectionResumed(
1263+
mock_factory.AsStdFunction(), std::move(mock), initial_request, nullptr,
1264+
first_response, Options{});
1265+
1266+
auto finalize1 = connection->Finalize({});
1267+
auto finalize2 = connection->Finalize({});
1268+
1269+
EXPECT_STATUS_OK(finalize1.get());
1270+
EXPECT_THAT(finalize2.get(), StatusIs(StatusCode::kFailedPrecondition));
1271+
}
1272+
12011273
} // namespace
12021274
GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END
12031275
} // namespace storage_internal

0 commit comments

Comments
 (0)