From 1586e2558ab9b57b9bbc12ad37eb9dd6d7d2eb61 Mon Sep 17 00:00:00 2001
From: Gauri Kalra <kalragauri@google.com>
Date: Fri, 24 Jul 2026 09:20:27 +0000
Subject: [PATCH] fix(storage): preserve flush state across concurrent flushes
 in async writers

---
 .../async/writer_connection_buffered.cc       |  6 ++++--
 .../async/writer_connection_buffered_test.cc  | 20 +++++++++++++++----
 .../async/writer_connection_resumed.cc        |  6 ++++--
 3 files changed, 24 insertions(+), 8 deletions(-)

diff --git a/google/cloud/storage/internal/async/writer_connection_buffered.cc b/google/cloud/storage/internal/async/writer_connection_buffered.cc
index 1da25199a35da..c047cece556ea 100644
--- a/google/cloud/storage/internal/async/writer_connection_buffered.cc
+++ b/google/cloud/storage/internal/async/writer_connection_buffered.cc
@@ -540,7 +540,6 @@ class AsyncWriterConnectionBufferedState
 
   void SetFlushed(std::unique_lock<std::mutex> lk, Status const& result) {
     if (!result.ok()) return SetError(std::move(lk), std::move(result));
-    flush_ = false;  // Reset flush flag; WriteLoop may set it again.
     // Do NOT reset finalize_ or finalizing_ here.
     auto handlers = ClearHandlers(lk);
     // Dequeue the promise corresponding to an explicit Flush() call, if any.
@@ -548,13 +547,16 @@ class AsyncWriterConnectionBufferedState
       // This can happen if SetError cleared the queue first, or if this
       // flush was triggered internally by buffer size (not by an explicit
       // Flush() call) and thus has no promise in the queue.
+      flush_ = false;
       lk.unlock();
       for (auto& h : handlers) h->Execute(Status{});
       return;
     }
     auto flushed = std::move(pending_flush_promises_.front());
     pending_flush_promises_.pop_front();
-
+    if (pending_flush_promises_.empty()) {
+      flush_ = false;
+    }
     lk.unlock();  // Unlock only once before notifying
     // Notify handlers and the specific flush promise *after* releasing the
     // lock.
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 3a37d8ae15a61..a179e4b1e6c6d 100644
--- a/google/cloud/storage/internal/async/writer_connection_buffered_test.cc
+++ b/google/cloud/storage/internal/async/writer_connection_buffered_test.cc
@@ -1376,12 +1376,17 @@ TEST(WriteConnectionBuffered, MultipleConcurrentFlushesAreQueued) {
       return sequencer.PushBack("Query1").then(
           [](auto) { return make_status_or(static_cast<std::int64_t>(4096)); });
     });
-    EXPECT_CALL(*mock, Write(expected_write_size(8192)))
+    EXPECT_CALL(*mock, Flush(expected_write_size(8192)))
         .Times(1)
         .WillOnce([&](auto) {
-          return sequencer.PushBack("Write").then(
+          return sequencer.PushBack("Flush2").then(
               [](auto) { return Status{}; });
         });
+    EXPECT_CALL(*mock, Query).WillOnce([&]() {
+      return sequencer.PushBack("Query2").then([](auto) {
+        return make_status_or(static_cast<std::int64_t>(4096 + 8192));
+      });
+    });
   }
 
   MockFactory mock_factory;
@@ -1409,10 +1414,17 @@ TEST(WriteConnectionBuffered, MultipleConcurrentFlushesAreQueued) {
   // After the Query, the first Flush future should be completed.
   EXPECT_STATUS_OK(f1.get());
 
-  // Satisfy the Write call for the remaining data.
+  // Satisfy the second Flush call.
   next = sequencer.PopFrontWithName();
-  EXPECT_EQ(next.second, "Write");
+  EXPECT_EQ(next.second, "Flush2");
+  next.first.set_value(true);
+
+  // The second Flush is also followed by a Query.
+  next = sequencer.PopFrontWithName();
+  EXPECT_EQ(next.second, "Query2");
   next.first.set_value(true);
+
+  EXPECT_STATUS_OK(f2.get());
 }
 
 TEST(WriteConnectionBuffered, FinalizeWithEmptyPayloadIsImmediate) {
diff --git a/google/cloud/storage/internal/async/writer_connection_resumed.cc b/google/cloud/storage/internal/async/writer_connection_resumed.cc
index a1bc93b656110..37faf9983817e 100644
--- a/google/cloud/storage/internal/async/writer_connection_resumed.cc
+++ b/google/cloud/storage/internal/async/writer_connection_resumed.cc
@@ -578,7 +578,6 @@ class AsyncWriterConnectionResumedState
 
   void SetFlushed(std::unique_lock<std::mutex> lk, Status const& result) {
     if (!result.ok()) return SetError(std::move(lk), std::move(result));
-    flush_ = false;  // Reset flush flag; WriteLoop may set it again.
     // Do NOT reset finalize_ or finalizing_ here.
     auto handlers = ClearHandlers(lk);
     // Dequeue the promise corresponding to an explicit Flush() call, if any.
@@ -586,13 +585,16 @@ class AsyncWriterConnectionResumedState
       // This can happen if SetError cleared the queue first, or if this
       // flush was triggered internally by buffer size (not by an explicit
       // Flush() call) and thus has no promise in the queue.
+      flush_ = false;
       lk.unlock();
       for (auto& h : handlers) h->Execute(Status{});
       return;
     }
     auto flushed = std::move(pending_flush_promises_.front());
     pending_flush_promises_.pop_front();
-
+    if (pending_flush_promises_.empty()) {
+      flush_ = false;
+    }
     lk.unlock();  // Unlock only once before notifying
     // Notify handlers and the specific flush promise *after* releasing the
     // lock.