[worker] Funnel all command proposals through LeaderState::run by MohamedBassem · Pull Request #5055 · restatedev/restate

MohamedBassem · 2026-07-17T11:50:34Z

This PR changes the handling of the PP RPCs and ingestion requests from happening inside the main PP loop and into the leader state run loop. This is done by introducing a new channel that gets consumed in the large stream_select of LeaderState run. With that all main callsites to SelfProposer are encapsulated inside LeaderState and in a lter PR will be dispatched by a new Self proposer scheduler (*).

A couple of caveats to be aware of:

The new channel is of size one to maintain the semantics of keeping only one inflight RPC not enqueued in self proposer. I don't forsee any problems in increasing it though if needed.
The network service arm is now guarded by having a permit into this new channel. This means that although we could service stuff like DedupQquery (which doesn't propose commands), it'll currenly be blocked until SelfProposer is unblocked. This is not very different from now though.
ActionEffects is no longer effects, and I renamed it to LeaderEvents now that it carries also network service actions. Now the new NetworkServiceEvent is huge (~0.5KB) and it made the entire LeaderEvent wider. Luckily, we only keep 10 of those in memory at a time. In a later PR, I'm getting rid of this enum altogether.

(*) Two other callsites to self proposer remain: the one callsite when candidate to send the AnnounceLeader command, and after becoming the leader to do any VersionBarrier command. Those two are rare and will be excluded from the SelfProposer backpressure.

Stack created with Sapling. Best reviewed with ReviewStack.

github-actions · 2026-07-17T12:18:08Z

Test Results

8 files ±0 8 suites ±0 4m 50s ⏱️ -24s
61 tests ±0 61 ✅ ±0 0 💤 ±0 0 ❌ ±0
268 runs ±0 268 ✅ ±0 0 💤 ±0 0 ❌ ±0

Results for commit 05a3446. ± Comparison against base commit 8c8d8de.

♻️ This comment has been updated with latest results.

chatgpt-codex-connector

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: fbb5177d4d

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

Open a pull request for review
Mark a draft as ready
Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

chatgpt-codex-connector · 2026-07-20T15:07:38Z

+        let decision = rpc::RpcHandler::handle(context, body).await;
+
+        match decision {
+            rpc::Decision::Propose(proposal) => permit.buffer_rpc_proposal(proposal, response_tx),


Preserve pause fencing when buffering RPC proposals

When the selected network branch handles a pause RPC for a VQueue-owned invocation, this now only enqueues the proposal. Before that queued NetworkService event is drained, LeaderState::handle_events can process a ready LeaderEvent::Invoker first (crates/worker/src/partition/leadership/leader_state.rs:451-469), and the fencing token is not removed until propose_pause_and_fence appends the pause (crates/worker/src/partition/leadership/leader_state.rs:677-687). In that race, an invoker effect that arrives after the pause RPC was selected by the PP can still be self-proposed ahead of the pause; previously the RPC branch called propose_pause_and_fence directly before returning to the select loop, so this extra buffering window regresses the pause/fencing ordering.

Useful? React with 👍 / 👎.

This seems like a real problem, will need to read the code around the fencing a bit to figure out a way around it.

I think this is not a problem as we want to fence invoker effects once the pause command has been written to Bifrost. That way we prevent that effects from a previous attempt mess up with a resumed invocation attempt. With the newly added channel we simply delay writing the pause command to Bifrost a little bit.

tillrohrmann

Thanks for creating this PR @MohamedBassem. The network rpc bounded channel of size 1 shouldn't have a big impact on the overall rpc throughput rate, right? The only rpc for which this could be relevant is the ingestion rpc as the others are expected to be of low frequency. I would assume that it will increase the tail latencies of the RPC calls a little bit as the select statement needs to select the leadership::run path on a following loop iteration.

What's the plan for when we have the leader events scheduler? Would the rpc channel keep a limit of 1? Would the other channels get a different capacity? If yes, would this affect the fairness of the scheduler?

We can also answer these questions once we get to the scheduler. So +1 for merging this PR :-)

tillrohrmann · 2026-07-21T21:48:46Z

+    ) -> Result<
+        tokio::sync::mpsc::OwnedPermit<NetworkServiceEvent>,
+        tokio::sync::mpsc::error::TrySendError<tokio::sync::mpsc::Sender<NetworkServiceEvent>>,
+    > {


Looking at the call sites, we could simplify this method by returning Option<OwnedPermit<NetworkServiceEvent>> since the caller will probably not do anything with the cloned network_events_tx anyway.

tillrohrmann · 2026-07-21T22:03:41Z

+            let network_processing_permit =
+                self.leadership_state.try_reserve_rpc_processing_permit();


There is no problem that we are only trying to reserve the permit at the beginning of every loop iteration because draining the network rpc tx and thereby creating new capacity/permits will always complete one of the select arms, right?

If I understand your question correctly, we only need at most one permit per loop because the permit is consumed only from a single select branch (the network branch), and that branch produces at most one message per branch.

I was thinking whether we need to select over the availability of the rpc processing permit to trigger a new loop iteration in case we didn't obtain a permit at the beginning. But I think this is not needed as we will finish the loop once we consume from the network rpc channel (and thereby creating new permits).

Ah I see what you mean. The permit actually can only be available once the LeaderState::run branch completes as it's the one that consumes the prev value. So we're guaranteed a loop iteration once the permit is available.

tillrohrmann · 2026-07-21T22:11:15Z

+        let decision = rpc::RpcHandler::handle(context, body).await;
+
+        match decision {
+            rpc::Decision::Propose(proposal) => permit.buffer_rpc_proposal(proposal, response_tx),


I think this is not a problem as we want to fence invoker effects once the pause command has been written to Bifrost. That way we prevent that effects from a previous attempt mess up with a resumed invocation attempt. With the newly added channel we simply delay writing the pause command to Bifrost a little bit.

tillrohrmann · 2026-07-21T22:29:14Z

@codex review

chatgpt-codex-connector

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: fbb5177d4d

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

Open a pull request for review
Mark a draft as ready
Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

chatgpt-codex-connector · 2026-07-21T22:33:44Z

 mod state_machine;
 pub mod types;

+use self::leadership::RpcProcessingPermit;


Move the local import into the required stanza

Move RpcProcessingPermit to the local-crate import stanza after the restate_* imports; placing this self import before the std stanza violates the repository's mandatory Rust import ordering.

AGENTS.md reference: AGENTS.md:L38-L42

Useful? React with 👍 / 👎.

tillrohrmann · 2026-07-22T07:32:01Z

            let config = self.node_ctx.config.live_load();
            let max_batching_size = config.worker.max_command_batch_size();
            let bytes_limit = config.worker.max_command_batch_bytes.as_usize();
+            let network_processing_permit =


Are we dropping the network_processing_permit on every loop iteration? It probably falls into the category of micro/pre-mature optimizations to cache the permit as the channel is not contended to avoid it to safe a few atomic operations.

Yes, we can cache it. The only caveat is that we'll need to remember to drop it when leader_state changes (e.g. leader to follower transitions), otherwise messages will get lost. Let me see if there's an easy way I can prevent this mismatch.

I think for simplicity, I'll leave it as is, I think it'll become a footgun as its lifecycle is not associated with the underlying leader state and can easily diverge.

tillrohrmann · 2026-07-22T07:36:30Z

+        // Close the network events channel and drain it to respond to all pending requests.
+        // The drain must not await: the partition processor's main loop reserves a permit
+        // before its select! and holds it across the arm that calls stop(), and recv()
+        // doesn't return `None` while a permit is outstanding — awaiting here would
+        // deadlock. Since that loop is the only sender and processes one arm at a time,
+        // the outstanding permit is guaranteed unused and every sent event is already
+        // buffered, so a non-blocking drain loses nothing.
+        self.network_events_stream.close();
+        while let Some(Some(event)) = self.network_events_stream.next().now_or_never() {
+            match event {
+                NetworkServiceEvent::RpcProposal { reciprocal, .. } => reciprocal.send(Err(
+                    PartitionProcessorRpcError::LostLeadership(self.partition_id),
+                )),
+                NetworkServiceEvent::IngestRecords { reciprocal, .. } => reciprocal.send(
+                    ResponseStatus::NotLeader {
+                        of: self.partition_id,
+                    }
+                    .into(),
+                ),
+            }
+        }


Given that we don't do message passing between different tasks, could we replace the network_events_stream channel with a simpler VecDeque which is being passed as part of the Processor trait? It probably falls under the category of micro optimizations and might not work if in a later PR we are relying on the fact that we have the network_events_stream channel.

Yes we can. I actually started with an Option<..> for a single slot message, then I thought we can do VecDeque to make the size configurable, but then I didn't like the API contract that the capacity check happens in the select guard and then inside it, we blindly write to the vecdeque hoping that we're within capacity. Instead of building a permit system on top, I decided to just reach for a channel for simplicity. This was all before writing the scheduler, though. The scheduler as you'll see in a later PR, is poll based, so the VecDeque wrapper will need to get a Notify as well to wake the scheduler if a write happens to this VecDeque (or at least persist the waker and call it on writes). So it slowly converges into a channel. I honestly don't have a sense of how much perf gains this can make, I can try to test it, and see if it makes a difference.

No need to do that if we later need the channel API in the scheduler.

This PR changes the handling of the PP RPCs and ingestion requests from happening inside the main PP loop and into the leader state run loop. This is done by introducing a new channel that gets consumed in the large stream_select of `LeaderState` run. With that all main callsites to `SelfProposer` are encapsulated inside `LeaderState` and in a lter PR will be dispatched by a new Self proposer scheduler (*). A couple of caveats to be aware of: 1. The new channel is of size one to maintain the semantics of keeping only one inflight RPC not enqueued in self proposer. I don't forsee any problems in increasing it though if needed. 2. The network service arm is now guarded by having a permit into this new channel. This means that although we could service stuff like `DedupQquery` (which doesn't propose commands), it'll currenly be blocked until `SelfProposer` is unblocked. This is not very different from now though. 3. `ActionEffects` is no longer effects, and I renamed it to `LeaderEvents` now that it carries also network service actions. Now the new `NetworkServiceEvent` is huge (~0.5KB) and it made the entire `LeaderEvent` wider. Luckily, we only keep 10 of those in memory at a time. In a later PR, I'm getting rid of this enum altogether. (*) Two other callsites to self proposer remain: the one callsite when candidate to send the `AnnounceLeader` command, and after becoming the leader to do any `VersionBarrier` command. Those two are rare and will be excluded from the SelfProposer backpressure.

This was referenced Jul 17, 2026

[worker][refactor] change rpc handlers to return their next steps instead of taking actions inline #5043

Merged

[MemoryPool] Add support for overdrafts in MemoryPool #5057

Merged

MohamedBassem force-pushed the pr5055 branch from b7b8cf3 to fbb5177 Compare July 20, 2026 14:59

MohamedBassem marked this pull request as ready for review July 20, 2026 15:01

chatgpt-codex-connector Bot reviewed Jul 20, 2026

View reviewed changes

MohamedBassem requested a review from tillrohrmann July 21, 2026 18:32

tillrohrmann approved these changes Jul 21, 2026

View reviewed changes

chatgpt-codex-connector Bot reviewed Jul 21, 2026

View reviewed changes

tillrohrmann reviewed Jul 22, 2026

View reviewed changes

This was referenced Jul 22, 2026

[worker] Introduce a new byte aware self proposer scheduler #5088

Open

[shuffle] Only report latest outbox truncation idx from the shuffle #5091

Merged

[timers] Fix missing wakeups in timer service #5093

Merged

MohamedBassem force-pushed the pr5055 branch from fbb5177 to e608fa0 Compare July 23, 2026 08:56

MohamedBassem force-pushed the pr5055 branch from e608fa0 to 05a3446 Compare July 24, 2026 08:45

MohamedBassem merged commit 63c65e8 into main Jul 24, 2026
57 checks passed

MohamedBassem deleted the pr5055 branch July 24, 2026 09:40

github-actions Bot locked and limited conversation to collaborators Jul 24, 2026

		let network_processing_permit =
		self.leadership_state.try_reserve_rpc_processing_permit();

Uh oh!

Conversation

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

Uh oh!

Uh oh!

github-actions Bot commented Jul 17, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Test Results

Uh oh!

chatgpt-codex-connector Bot left a comment

Choose a reason for hiding this comment

💡 Codex Review

Uh oh!

chatgpt-codex-connector Bot Jul 20, 2026

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!

tillrohrmann left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

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!

tillrohrmann commented Jul 21, 2026

Uh oh!

chatgpt-codex-connector Bot left a comment

Choose a reason for hiding this comment

💡 Codex Review

Uh oh!

chatgpt-codex-connector Bot Jul 21, 2026

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!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

MohamedBassem commented Jul 17, 2026 •

edited

Loading

github-actions Bot commented Jul 17, 2026 •

edited

Loading