fix(storage): avoid premature flush completion and concurrent writes in AsyncWriterConnectionBuffered (#16169)

Author: kalragauri

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

@@ -350,6 +350,19 @@ class AsyncWriterConnectionBufferedState
     }
     // If the buffer is small enough, collect all the handlers to notify them.
     auto const handlers = ClearHandlersIfEmpty(lk);
+    if (is_resume) {
+      // We are resuming. The pending flush promises (if any) should not be
+      // satisfied yet, because we haven't actually flushed the data on the new
+      // connection. The `WriteLoop` will trigger a flush (potentially empty)
+      // if `flush_` is still true, which will satisfy the promises when it
+      // completes. However, we still need to notify any handlers waiting for
+      // the buffer to shrink, and we need to restart the write loop.
+      resuming_ = false;
+      lk.unlock();
+      for (auto const& h : handlers) h->Execute(Status{});
+      WriteLoop(std::unique_lock<std::mutex>(mu_));
+      return;
+    }
     // SetFlushed will release the lock before returning.
     SetFlushed(std::move(lk), Status{});
     // Re-acquire the lock to re-enter the write loop.
@@ -388,6 +401,9 @@ class AsyncWriterConnectionBufferedState
       if (!s.ok() && cancelled_) {
         return SetError(std::move(lk), std::move(s));
       }
+      // Guard against concurrent resume attempts.
+      if (resuming_) return;
+      resuming_ = true;
     }
     // Pass the original status `s`, `was_finalizing`, and `was_closing` to the
     // callback.
@@ -463,6 +479,7 @@ class AsyncWriterConnectionBufferedState
     finalize_ = false;
     finalizing_ = false;  // Reset finalizing flag
     flush_ = false;
+    resuming_ = false;  // Reset resuming flag
     // Check if the promise has already been completed.
     if (finalized_promise_completed_) {
       // Since the lock is passed by value, no explicit unlock is needed.
@@ -487,6 +504,7 @@ class AsyncWriterConnectionBufferedState
     close_ = false;
     closing_ = false;
     flush_ = false;
+    resuming_ = false;
     // Check if the promise has already been completed.
     if (closed_promise_completed_) {
       return;
@@ -526,10 +544,6 @@ class AsyncWriterConnectionBufferedState
     // lock.
     for (auto& h : handlers) h->Execute(Status{});
     flushed.set_value(result);
-    // Restart the write loop ONLY if we are not already finalizing.
-    // If finalizing_ is true, the completion will be handled by OnFinalize.
-    std::unique_lock<std::mutex> loop_lk(mu_);
-    if (!finalizing_) WriteLoop(std::move(loop_lk));
   }
 
   void SetError(std::unique_lock<std::mutex> lk, Status const& status) {
@@ -540,6 +554,7 @@ class AsyncWriterConnectionBufferedState
     close_ = false;
     closing_ = false;  // Reset closing flag
     flush_ = false;
+    resuming_ = false;  // Reset resuming flag
 
     // Always clear handlers and pending flushes on error.
     auto handlers = ClearHandlers(lk);
@@ -685,6 +700,9 @@ class AsyncWriterConnectionBufferedState
 
   // Tracks if the final promise (`finalized_`) has been completed.
   bool finalized_promise_completed_ = false;
+
+  // True if the resume loop is running. Prevents re-entry.
+  bool resuming_ = false;
 };
 
 /**

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

@@ -771,6 +771,84 @@ TEST(WriteConnectionBuffered, ErrorFailsPendingFlushes) {
   EXPECT_THAT(f2.get(), StatusIs(resume_error.code()));
 }
 
+TEST(WriteConnectionBuffered, FlushResumesAndDoesNotCompletePrematurely) {
+  AsyncSequencer<bool> sequencer;
+
+  auto expected_write_size = [](std::size_t n) {
+    return ResultOf(
+        "payload size", [](auto payload) { return payload.size(); }, Eq(n));
+  };
+
+  auto mock1 = std::make_unique<MockAsyncWriterConnection>();
+  EXPECT_CALL(*mock1, UploadId).WillRepeatedly(Return("test-upload-id"));
+  EXPECT_CALL(*mock1, PersistedState)
+      .WillRepeatedly(Return(MakePersistedState(0)));
+  // The first flush fails, triggering resume.
+  EXPECT_CALL(*mock1, Flush(expected_write_size(8 * 1024))).WillOnce([&](auto) {
+    return sequencer.PushBack("Flush1").then(
+        [](auto) { return Status(StatusCode::kUnavailable, "try again"); });
+  });
+
+  auto mock2 = std::make_unique<MockAsyncWriterConnection>();
+  EXPECT_CALL(*mock2, UploadId).WillRepeatedly(Return("test-upload-id"));
+  // OnResume will query persisted state. We return 0, meaning the data was
+  // lost.
+  EXPECT_CALL(*mock2, PersistedState)
+      .WillRepeatedly(Return(MakePersistedState(0)));
+  // The resumed connection should receive the Flush call again.
+  EXPECT_CALL(*mock2, Flush(expected_write_size(8 * 1024))).WillOnce([&](auto) {
+    return sequencer.PushBack("Flush2").then([](auto) { return Status{}; });
+  });
+  // After Flush2 succeeds, it will query the status.
+  EXPECT_CALL(*mock2, Query).WillOnce([&]() {
+    return sequencer.PushBack("Query").then([](auto) {
+      return make_status_or(static_cast<std::int64_t>(8 * 1024));
+    });
+  });
+
+  MockFactory mock_factory;
+  EXPECT_CALL(mock_factory, Call).WillOnce([&]() {
+    return sequencer.PushBack("Resume").then([&](auto) {
+      return make_status_or(
+          std::unique_ptr<storage::AsyncWriterConnection>(std::move(mock2)));
+    });
+  });
+
+  auto connection = MakeWriterConnectionBuffered(
+      mock_factory.AsStdFunction(), std::move(mock1), TestOptions());
+
+  auto f = connection->Flush(TestPayload(8 * 1024));
+  ASSERT_FALSE(f.is_ready());
+
+  // Let the first flush fail.
+  auto next = sequencer.PopFrontWithName();
+  EXPECT_EQ(next.second, "Flush1");
+  next.first.set_value(true);
+
+  // Trigger resume.
+  next = sequencer.PopFrontWithName();
+  EXPECT_EQ(next.second, "Resume");
+  next.first.set_value(true);
+
+  // The flush promise must not be ready yet, because the data has not been
+  // flushed on the new connection.
+  ASSERT_FALSE(f.is_ready());
+
+  // Let the second flush succeed.
+  next = sequencer.PopFrontWithName();
+  EXPECT_EQ(next.second, "Flush2");
+  next.first.set_value(true);
+
+  // Let the query complete.
+  next = sequencer.PopFrontWithName();
+  EXPECT_EQ(next.second, "Query");
+  next.first.set_value(true);
+
+  // Now the flush promise should be completed.
+  ASSERT_TRUE(f.is_ready());
+  EXPECT_STATUS_OK(f.get());
+}
+
 TEST(WriteConnectionBuffered, CloseEmpty) {
   AsyncSequencer<bool> sequencer;
   auto mock = std::make_unique<MockAsyncWriterConnection>();
@@ -950,29 +1028,53 @@ TEST(WriteConnectionBuffered, CloseFailsAndResumeSucceedsButNotClosed) {
 TEST(WriteConnectionBuffered, CloseFailsAndResumeSucceedsAndFinalized) {
   AsyncSequencer<bool> sequencer;
   auto close_error = TransientError();
+  auto write_error = TransientError();
+  auto resume2_error = PermanentError();
 
-  auto mock = std::make_unique<MockAsyncWriterConnection>();
-  EXPECT_CALL(*mock, UploadId).WillRepeatedly(Return("test-upload-id"));
-  EXPECT_CALL(*mock, PersistedState)
+  auto mock1 = std::make_unique<MockAsyncWriterConnection>();
+  EXPECT_CALL(*mock1, UploadId).WillRepeatedly(Return("test-upload-id"));
+  EXPECT_CALL(*mock1, PersistedState)
       .WillRepeatedly(Return(MakePersistedState(0)));
-  EXPECT_CALL(*mock, Close).WillOnce([&](auto) {
+  EXPECT_CALL(*mock1, Close).WillOnce([&](auto) {
     return sequencer.PushBack("Close").then(
         [close_error](auto) { return close_error; });
   });
 
+  auto resumed_mock = std::make_unique<MockAsyncWriterConnection>();
+  auto* resumed_mock_ptr = resumed_mock.get();
+
+  // Mock Write on the resumed connection to fail, to trigger a second resume.
+  EXPECT_CALL(*resumed_mock_ptr, Write).WillOnce([&](auto) {
+    return sequencer.PushBack("Write2").then(
+        [write_error](auto) { return write_error; });
+  });
+
   MockFactory mock_factory;
-  EXPECT_CALL(mock_factory, Call).WillOnce([&]() {
-    return sequencer.PushBack("Resume").then([](auto) {
-      auto resumed_mock = std::make_unique<MockAsyncWriterConnection>();
-      EXPECT_CALL(*resumed_mock, PersistedState)
-          .WillRepeatedly(Return(TestObject()));
-      return make_status_or(std::unique_ptr<storage::AsyncWriterConnection>(
-          std::move(resumed_mock)));
+  {
+    InSequence seq;
+    // The resume will succeed and return a new mock that reports the object is
+    // already finalized (which implies closed).
+    EXPECT_CALL(mock_factory, Call)
+        .WillOnce([&, rm = std::move(resumed_mock)]() mutable {
+          return sequencer.PushBack("Resume").then([rm = std::move(rm)](
+                                                       auto) mutable {
+            EXPECT_CALL(*rm, PersistedState)
+                .WillRepeatedly(Return(TestObject()));
+            return make_status_or(
+                std::unique_ptr<storage::AsyncWriterConnection>(std::move(rm)));
+          });
+        });
+    // The second resume (after write failure) will fail.
+    EXPECT_CALL(mock_factory, Call).WillOnce([&]() {
+      return sequencer.PushBack("Resume2").then([resume2_error](auto) {
+        return StatusOr<std::unique_ptr<storage::AsyncWriterConnection>>(
+            resume2_error);
+      });
     });
-  });
+  }
 
   auto connection = MakeWriterConnectionBuffered(
-      mock_factory.AsStdFunction(), std::move(mock), TestOptions());
+      mock_factory.AsStdFunction(), std::move(mock1), TestOptions());
 
   auto close = connection->Close({});
   ASSERT_FALSE(close.is_ready());
@@ -986,6 +1088,30 @@ TEST(WriteConnectionBuffered, CloseFailsAndResumeSucceedsAndFinalized) {
   next.first.set_value(true);
 
   EXPECT_STATUS_OK(close.get());
+
+  // Write to the closed connection. Since HWM is 32KB, writing 16 bytes
+  // returns OK immediately because it is buffered.
+  auto w = connection->Write(TestPayload(16));
+  ASSERT_TRUE(w.is_ready());
+  EXPECT_STATUS_OK(w.get());
+
+  // However, the background write loop will call resumed_mock->Write.
+  // Let that background Write fail.
+  next = sequencer.PopFrontWithName();
+  EXPECT_EQ(next.second, "Write2");
+  next.first.set_value(true);
+
+  // This triggers a second Resume. If resuming_ was not reset in SetClosed,
+  // this would hang because Resume would return early and not call
+  // mock_factory.
+  next = sequencer.PopFrontWithName();
+  EXPECT_EQ(next.second, "Resume2");
+  next.first.set_value(true);
+
+  // A subsequent write should fail immediately with the second resume error.
+  auto w2 = connection->Write(TestPayload(16));
+  ASSERT_TRUE(w2.is_ready());
+  EXPECT_THAT(w2.get(), StatusIs(resume2_error.code()));
 }
 
 TEST(WriteConnectionBuffered, Query) {
@@ -1250,11 +1376,9 @@ TEST(WriteConnectionBuffered, MultipleConcurrentFlushesAreQueued) {
       return sequencer.PushBack("Query1").then(
           [](auto) { return make_status_or(static_cast<std::int64_t>(4096)); });
     });
-    // The race causes Write() to be called twice for the remaining
-    // 8k.
     EXPECT_CALL(*mock, Write(expected_write_size(8192)))
-        .Times(2)
-        .WillRepeatedly([&](auto) {
+        .Times(1)
+        .WillOnce([&](auto) {
           return sequencer.PushBack("Write").then(
               [](auto) { return Status{}; });
         });
@@ -1285,10 +1409,7 @@ TEST(WriteConnectionBuffered, MultipleConcurrentFlushesAreQueued) {
   // After the Query, the first Flush future should be completed.
   EXPECT_STATUS_OK(f1.get());
 
-  // Satisfy the two racy Write calls for the remaining data.
-  next = sequencer.PopFrontWithName();
-  EXPECT_EQ(next.second, "Write");
-  next.first.set_value(true);
+  // Satisfy the Write call for the remaining data.
   next = sequencer.PopFrontWithName();
   EXPECT_EQ(next.second, "Write");
   next.first.set_value(true);