perf(pubsub): move batch flush timer from Dispatcher into batch actors by panimidi · Pull Request #6050 · googleapis/google-cloud-rust

panimidi · 2026-07-14T18:48:01Z

perf(pubsub): move batch flush timer from Dispatcher into batch actors

The Dispatcher runs a single batch-flush timer that is armed unconditionally
and reset on every fire, regardless of whether any batch actor has buffered
messages. Because each Publisher owns its own Dispatcher, an application with
many publishers performs continuous timer work even when completely idle. In
Dispatcher::run the timer is created once, and on every delay_threshold
interval it flushes all batch actors and then re-arms itself unconditionally,
so an idle publisher keeps scheduling and servicing a timer with nothing to
send. At high publisher cardinality this becomes a constant source of
time-wheel work; we observed it dominating idle CPU in a production service
that fans a single logical publisher out across several thousand topics.

For reference, the Java client's Publisher.setupAlarm only schedules the
alarm while a batch is non-empty and cancels it once the batch drains, so it
holds no timer at rest. This change brings the Rust client in line with that
behavior.

This change removes the shared timer from Dispatcher::run, leaving the
Dispatcher as a pure router of Publish, Flush, and ResumePublish commands, and
adds a lazy flush timer to both ConcurrentBatchActor and
SequentialBatchActor. Each actor arms its timer when the first message enters
an otherwise-empty batch, re-arms it after a timer flush that leaves work
pending, and disarms it once a flush drains the batch. An idle actor then
blocks in recv with no scheduled timer and therefore no time-wheel entry.

Batching semantics are unchanged: a batch is still flushed after at most
delay_threshold. There is no public API change and no change to item
visibility. Existing tests cover the behavior, including
batch_sends_on_delay_threshold, which disables the count and byte thresholds
and asserts that batches flush at exactly delay_threshold across repeated arm
and disarm cycles, for both keyed and unkeyed messages. The package test suite
passes.

google-cla · 2026-07-14T18:48:12Z

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

gemini-code-assist

Code Review

This pull request refactors the Pub/Sub publisher's batching mechanism by moving the flush timer from the central Dispatcher into individual batch actors (ConcurrentBatchActor and SequentialBatchActor) to achieve zero idle overhead. The review feedback highlights several opportunities to optimize these lazy timers. Specifically, when a timer fires or a batch is flushed due to reaching its threshold, the timer should be unconditionally disarmed (timer = None) to prevent unnecessary wakeups and avoid an early flush bug where subsequent messages are flushed prematurely.

gemini-code-assist · 2026-07-14T18:51:10Z

+                // Flush on timer — only active when messages are pending.
+                _ = async { timer.as_mut().unwrap().await }, if timer.is_some() => {
+                    self.flush(&mut inflight, &mut batch);
+                    if batch.is_empty() && inflight.is_empty() {
+                        timer = None;
+                    } else {
+                        timer.as_mut().unwrap().as_mut()
+                            .reset(tokio::time::Instant::now() + delay);
+                    }
+                }


Timer Efficiency & Early Flush Bug

When the timer fires, self.flush is called, which drains batch completely. Thus, batch is guaranteed to be empty. Keeping the timer active when batch is empty (just because inflight is not empty) keeps an unnecessary timer on the time-wheel, which goes against the PR's goal of zero idle overhead. Furthermore, if a new message arrives while inflight is active, the timer is already running and will fire early, flushing the new message before it has been buffered for the full delay_threshold. Therefore, we should always set timer = None immediately after flushing, rather than resetting it or checking inflight.is_empty().

// Flush on timer — only active when messages are pending. _ = async { timer.as_mut().unwrap().await; }, if timer.is_some() => { self.flush(&mut inflight, &mut batch); timer = None; }

References

Using expect() or unwrap() is acceptable in production Rust code when it asserts an invariant or handles an unreachable state that is guaranteed by the system or backend logic.

This is correct, the concurrent batch actor flushes all messages. Please fix.

gemini-code-assist · 2026-07-14T18:51:10Z

+                // Flush on timer — only active when messages are pending.
+                _ = async { timer.as_mut().unwrap().await }, if timer.is_some() => {
+                    self.flush(&mut inflight, &mut batch).await;
+                    inflight = JoinSet::new();
+                    if self.pending_msgs.is_empty() && batch.is_empty() {
+                        timer = None;
+                    } else {
+                        timer.as_mut().unwrap().as_mut()
+                            .reset(tokio::time::Instant::now() + delay);
+                    }
+                }


Timer Efficiency

self.flush flushes all pending messages and awaits them, so both self.pending_msgs and batch are guaranteed to be empty when it returns. Therefore, we can safely and unconditionally set timer = None to avoid keeping an unnecessary timer active.

// Flush on timer — only active when messages are pending. _ = async { timer.as_mut().unwrap().await; }, if timer.is_some() => { self.flush(&mut inflight, &mut batch).await; inflight = JoinSet::new(); timer = None; }

References

Using expect() or unwrap() is acceptable in production Rust code when it asserts an invariant or handles an unreachable state that is guaranteed by the system or backend logic.

This is also correct. We block until all messages are sent, so it will be empty at the time we reach this line.

gemini-code-assist · 2026-07-14T18:51:10Z

                        Some(ToBatchActor::Publish(msg)) => {
                            self.add_msg_and_flush(&mut inflight, &mut batch, msg);
+                            if timer.is_none() {
+                                timer = Some(Box::pin(tokio::time::sleep(delay)));
+                            }
                        },


Disarm Timer on Threshold Flush

If add_msg_and_flush flushes the batch because it reached the threshold, batch becomes empty. In this case, we should disarm the timer (timer = None) immediately so that it doesn't fire unnecessarily for an empty batch. If the batch is not empty and the timer is not armed, we arm it.

Some(ToBatchActor::Publish(msg)) => { self.add_msg_and_flush(&mut inflight, &mut batch, msg); if batch.is_empty() { timer = None; } else if timer.is_none() { timer = Some(Box::pin(tokio::time::sleep(delay))); } },

This is correct. The batch may be empty here. Either checking the batch to see if it is empty or modifying add_msg_and_flush to return whether it left messages in the batch would be reasonable options.

gemini-code-assist · 2026-07-14T18:51:10Z

                        Some(ToBatchActor::Publish(msg)) => {
                            self.pending_msgs.push_back(msg);
                            if inflight.is_empty() {
                                self.move_to_batch_and_flush(&mut inflight, &mut batch);
                            }
+                            if timer.is_none() {
+                                timer = Some(Box::pin(tokio::time::sleep(delay)));
+                            }
                        },


Disarm Timer on Threshold Flush

If move_to_batch_and_flush flushes the batch because it reached the threshold, and both self.pending_msgs and batch become empty, we should disarm the timer (timer = None) immediately to avoid unnecessary timer wakeups.

Some(ToBatchActor::Publish(msg)) => { self.pending_msgs.push_back(msg); if inflight.is_empty() { self.move_to_batch_and_flush(&mut inflight, &mut batch); } if self.pending_msgs.is_empty() && batch.is_empty() { timer = None; } else if timer.is_none() { timer = Some(Box::pin(tokio::time::sleep(delay))); } },

This is correct as well. A flush may have just occurred so there may be no messages.

panimidi · 2026-07-14T18:52:18Z

For full disclosure: the branch, perf/pubsub-lazy-timers-full, contains this timer change plus an additional optimization that isn't part of this PR: a fast path where Publisher::publish sends messages without an ordering key directly to a pre-created default batch actor, bypassing the Dispatcher entirely.

This is the exact code we ran our tests against, leaving it here for reference but I believe the "fast-path" adds a lot of noise like visibility changes and wasn't the main driver of performance improvement anyway.

suzmue · 2026-07-14T18:54:11Z

Thank you for the PR! In addition to signing the CLA, would you please file an issue corresponding to this PR? (See Contributing).

Edit: Filed #6052, no need to file an issue.

suzmue · 2026-07-14T18:54:18Z

/gcbrun

codecov · 2026-07-14T20:32:26Z

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 97.01%. Comparing base (a69f6d5) to head (4ecdf66).
⚠️ Report is 16 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #6050      +/-   ##
==========================================
- Coverage   97.03%   97.01%   -0.02%     
==========================================
  Files         252      252              
  Lines       63526    63539      +13     
==========================================
+ Hits        61640    61645       +5     
- Misses       1886     1894       +8

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

suzmue · 2026-07-14T22:47:59Z

+                // Flush on timer — only active when messages are pending.
+                _ = async { timer.as_mut().unwrap().await }, if timer.is_some() => {
+                    self.flush(&mut inflight, &mut batch);
+                    if batch.is_empty() && inflight.is_empty() {
+                        timer = None;
+                    } else {
+                        timer.as_mut().unwrap().as_mut()
+                            .reset(tokio::time::Instant::now() + delay);
+                    }
+                }


This is correct, the concurrent batch actor flushes all messages. Please fix.

suzmue · 2026-07-14T22:50:43Z

                        Some(ToBatchActor::Publish(msg)) => {
                            self.add_msg_and_flush(&mut inflight, &mut batch, msg);
+                            if timer.is_none() {
+                                timer = Some(Box::pin(tokio::time::sleep(delay)));
+                            }
                        },


This is correct. The batch may be empty here. Either checking the batch to see if it is empty or modifying add_msg_and_flush to return whether it left messages in the batch would be reasonable options.

suzmue · 2026-07-14T22:52:04Z

        // We have multiple inflight batches concurrently.
+        let delay = self.context.batching_options.delay_threshold;
+        // Lazy timer: armed on the first message, disarmed after a flush that
+        // leaves the batch empty. An idle actor blocks in recv() with no time-
+        // wheel entry — zero scheduler overhead at rest.
+        let mut timer: Option<std::pin::Pin<Box<tokio::time::Sleep>>> = None;
        let mut inflight = JoinSet::new();


The comment got split from the line of code it was referencing.

Suggested change

// We have multiple inflight batches concurrently.

let delay = self.context.batching_options.delay_threshold;

// Lazy timer: armed on the first message, disarmed after a flush that

// leaves the batch empty. An idle actor blocks in recv() with no time-

// wheel entry — zero scheduler overhead at rest.

let mut timer: Option<std::pin::Pin<Box<tokio::time::Sleep>>> = None;

let mut inflight = JoinSet::new();

let delay = self.context.batching_options.delay_threshold;

// Lazy timer: armed on the first message, disarmed after a flush that

// leaves the batch empty. An idle actor blocks in recv() with no time-

// wheel entry — zero scheduler overhead at rest.

let mut timer: Option<std::pin::Pin<Box<tokio::time::Sleep>>> = None;

// We have multiple inflight batches concurrently.

let mut inflight = JoinSet::new();

suzmue · 2026-07-14T23:03:39Z

+                // Flush on timer — only active when messages are pending.
+                _ = async { timer.as_mut().unwrap().await }, if timer.is_some() => {
+                    self.flush(&mut inflight, &mut batch).await;
+                    inflight = JoinSet::new();
+                    if self.pending_msgs.is_empty() && batch.is_empty() {
+                        timer = None;
+                    } else {
+                        timer.as_mut().unwrap().as_mut()
+                            .reset(tokio::time::Instant::now() + delay);
+                    }
+                }


This is also correct. We block until all messages are sent, so it will be empty at the time we reach this line.

suzmue · 2026-07-14T23:08:52Z

                        Some(ToBatchActor::Publish(msg)) => {
                            self.pending_msgs.push_back(msg);
                            if inflight.is_empty() {
                                self.move_to_batch_and_flush(&mut inflight, &mut batch);
                            }
+                            if timer.is_none() {
+                                timer = Some(Box::pin(tokio::time::sleep(delay)));
+                            }
                        },


This is correct as well. A flush may have just occurred so there may be no messages.

suzmue · 2026-07-14T23:14:21Z

-                            // By dropping the BatchActorHandles, they will individually handle the
-                            // shutdown procedures.
                            drop(batch_actors);
-                            // When we drop actor_tasks, some batch actors may not have started execution


Extraneous diff, please restore comment.

suzmue · 2026-07-14T23:14:27Z

                        },
                        Some(ToDispatcher::ResumePublish(ordering_key)) => {
                            if let Some(batch_actor) = batch_actors.get_mut(&ordering_key) {
-                                // Send down the same tx for the BatchActors to directly signal completion


Extraneous diff, please restore comment.

suzmue · 2026-07-14T23:14:33Z

                                flush_set.spawn(rx);
                            }
                            tokio::spawn(async move {
-                                // Wait on all the flush operations.


Extraneous diff, please restore comment.

suzmue · 2026-07-14T23:32:55Z

+    /// Each batch actor owns its own lazy flush timer (see `ConcurrentBatchActor`
+    /// and `SequentialBatchActor`), so the Dispatcher no longer runs a timer of
+    /// its own.


nit: this will live in our code forever and "no longer" isn't necessary to document the state.

Suggested change

/// Each batch actor owns its own lazy flush timer (see `ConcurrentBatchActor`

/// and `SequentialBatchActor`), so the Dispatcher no longer runs a timer of

/// its own.

/// Each batch actor owns its own lazy flush timer (see `ConcurrentBatchActor`

/// and `SequentialBatchActor`).

suzmue · 2026-07-14T23:43:55Z

+        // Lazy timer: armed on the first message, disarmed after a flush that
+        // leaves the batch empty. An idle actor blocks in recv() with no time-
+        // wheel entry — zero scheduler overhead at rest.


nit: When reviewing it was initially unclear to me why the tradeoff of using a heap allocation on (potentially) every batch was worth it compared to adding a boolean to guard the tokio::pin!. I believe from your comment it is because we want to fully get rid of the timer to get rid of the entries in the time wheel. I have included a suggestion for updating the comment that makes it a bit more explicit.

I am curious, did you also test this with a single timer that gets reset, as opposed to one that gets dropped? Was there a performance difference?

Suggested change

// Lazy timer: armed on the first message, disarmed after a flush that

// leaves the batch empty. An idle actor blocks in recv() with no time-

// wheel entry — zero scheduler overhead at rest.

// Lazy timer: Armed when the first message enters a batch and dropped (`None`)

// when a flush drains the batch. Dropping the timer unregisters it from Tokio's

// timer driver, ensuring idle actors hold zero time-wheel entries and incur

// zero CPU wakeups at rest.

I did not, I wanted to mirror the Java client library's behavior (canceling the scheduled task -> dropping it here), but it's a good point. Will test and report back later.

Measured this with a small demo project (https://github.com/panimidi/pubsub-demo) for quick repro — 2000 publishers on one topic, delay_threshold = 1ms, against a local emulator — with a 10s idle phase and a 20s load phase at ~500 msg/s.

Results reading the client process's own CPU via getrusage:

client idle %CPU load %CPU

unmodified (before this PR) 154.6% 228.6%

this PR — drop timer 0.0% 12.7%

alternative — reset timer 0.0% 12.6%

(The 1ms delay and 2k publishers is a stress case to show the issue clearly)

The reset variant shown on the table is on branch perf/pubsub-lazy-timers-reset if you'd like to review, but I'd lean on keeping the drop version.

panimidi · 2026-07-15T14:41:49Z

Thank you for the comments, I'll be reviewing them and sending fixes later.

I am also looking into the CLA, my company (MercadoLibre) doesn't currently have a CCLA signed with Google so it may take a little while.

The Dispatcher armed a single delay-threshold timer unconditionally and reset it on every fire, regardless of whether any batch actor had pending messages. Because each Publisher owns its own Dispatcher, an application with many publishers performs continuous timer work even when completely idle. This change moves the flush timer into the ConcurrentBatchActor and the SequentialBatchActor. Each actor arms a timer lazily when the first message enters an empty batch and disarms it once a flush drains the batch, so an idle actor blocks in recv with no scheduled timer. The Dispatcher becomes a pure router. This mirrors the Java client, whose setupAlarm only schedules the alarm while a batch is non-empty and cancels it once the batch empties. Batching semantics are unchanged: a batch is still flushed after at most the configured delay threshold. Existing tests cover the behavior, including batch_sends_on_delay_threshold, which asserts flush timing across repeated arm and disarm cycles for both keyed and unkeyed messages. There is no public API change and no change to item visibility.

Address review feedback on the lazy batch-actor timers. Both batch actors now drop the timer (set it to None) after a timer-driven flush, rather than resetting it when in-flight sends remain. flush fully drains the batch, so keeping the timer armed left an unnecessary time-wheel entry and could flush a later message early, before it had waited the full delay threshold. The Publish arms now also disarm the timer when a threshold flush leaves the batch empty, and only arm it while messages remain buffered. The concurrent timer arm keeps a minimal flush: the send is spawned into the inflight set and reaped by the existing join_next arm, so it does not block on join_all the way the Flush command does. The sequential timer arm mirrors its Flush branch, whose flush is async. Also restore the surrounding comments removed incidentally and reword the Dispatcher and lazy-timer comments.

panimidi · 2026-07-16T18:49:14Z

Force-pushed to re-author the commits for the corporate CLA. No content changes.

suzmue · 2026-07-16T20:37:51Z

/gcbrun

product-auto-label Bot added the api: pubsub Issues related to the Pub/Sub API. label Jul 14, 2026

gemini-code-assist Bot reviewed Jul 14, 2026

View reviewed changes

suzmue mentioned this pull request Jul 14, 2026

Improve Publish client performance when there are many idle clients and the batch delay is short #6052

Open

suzmue reviewed Jul 14, 2026

View reviewed changes

panimidi added 2 commits July 16, 2026 15:26

panimidi force-pushed the perf/pubsub-lazy-timers branch from 2f502ba to 4ecdf66 Compare July 16, 2026 18:44

-        // Lazy timer: armed on the first message, disarmed after a flush that
-        // leaves the batch empty. An idle actor blocks in recv() with no time-
-        // wheel entry — zero scheduler overhead at rest.
+// Lazy timer: Armed when the first message enters a batch and dropped (`None`)
+// when a flush drains the batch. Dropping the timer unregisters it from Tokio's
+// timer driver, ensuring idle actors hold zero time-wheel entries and incur
+// zero CPU wakeups at rest.

client	idle %CPU	load %CPU
unmodified (before this PR)	154.6%	228.6%
this PR — drop timer	0.0%	12.7%
alternative — reset timer	0.0%	12.6%

Uh oh!

Conversation

panimidi commented Jul 14, 2026

Uh oh!

google-cla Bot commented Jul 14, 2026

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist Bot Jul 14, 2026

Choose a reason for hiding this comment

Timer Efficiency & Early Flush Bug

Uh oh!

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jul 14, 2026

Choose a reason for hiding this comment

Timer Efficiency

Uh oh!

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jul 14, 2026

Choose a reason for hiding this comment

Disarm Timer on Threshold Flush

Uh oh!

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jul 14, 2026

Choose a reason for hiding this comment

Disarm Timer on Threshold Flush

Uh oh!

Choose a reason for hiding this comment

Uh oh!

panimidi commented Jul 14, 2026

Uh oh!

suzmue commented Jul 14, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

suzmue commented Jul 14, 2026 • edited by dbolduc Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

codecov Bot commented Jul 14, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

panimidi commented Jul 15, 2026

Uh oh!

panimidi commented Jul 16, 2026

Uh oh!

suzmue commented Jul 16, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

suzmue commented Jul 14, 2026 •

edited

Loading

suzmue commented Jul 14, 2026 •

edited by dbolduc

Loading

codecov Bot commented Jul 14, 2026 •

edited

Loading