Use ingestion-client in the `Shuffler` #4024

muhamadazmy · 2025-11-17T17:10:37Z

Use ingestion-client in the Shuffler

Avoid direct writes to bifrost in shuffler by using a
dedicated ingestion-client instance.

Stack created with Sapling. Best reviewed with ReviewStack.

tillrohrmann

Thanks a lot for replacing the direct Bifrost write with the IngressClient in the Shuffle @muhamadazmy. Maybe the name IngressClient does not fit 100% given that now also the Shuffle uses it. Maybe something like IngestionClient or so works better. Given that we don't use the send window of the IngressClient yet, I wouldn't expect a different runtime behavior of the shuffle. Once we have this, I would be interested in how the overall shuffle throughput increases by using the IngressClient.

I left a few minor comments for your consideration.

crates/worker/src/lib.rs

crates/worker/src/partition/leadership/mod.rs

crates/worker/src/partition/shuffle.rs

crates/worker/src/lib.rs

tillrohrmann · 2025-11-25T22:43:37Z

There seem to be a few test failures on GHA.

tillrohrmann · 2025-11-26T08:35:03Z

crates/worker/src/partition/shuffle.rs

+                    ingress
+                        .ingest(
+                            msg.partition_key(),
+                            IngestRecord::from_parts(msg.record_keys(), msg),
+                        )
+                        .await?;


What about support for rolling upgrades?

ingest does not fail unless the ingession client is closed. This means worst case is that it will block until leaders are responsive.

I think to avoid this situation it's possible we release support for the Ingest messages in PP first before actually using them in the following release.

tillrohrmann

Thanks a lot for updating the shuffle to use the new ingestion client @muhamadazmy. LGTM. +1 for merging.

Did you ran any benchmarks/tests for a workload that needs to shuffle a lot to check whether we see an improvement with these changes?

tillrohrmann · 2025-12-19T13:50:03Z

crates/worker/src/partition_processor_manager.rs


        let partition_store_manager = PartitionStoreManager::create().await?;

+        let ingress = IngestionClient::new(


ingestion_client?

tillrohrmann · 2025-12-19T13:54:03Z

crates/types/src/config/worker.rs

+    ///
+    /// Settings for the shared ingestion client used by all workers to
+    /// manage record ingestion across partitions (shuffle).
+    pub shuffle: IngestionOptions,


The IngestionOptions options refer to the Kafka ingestion. Should the description be updated to also include the shuffle or can it be generic enough?

I will try to update the docs of IngestionOptions to make it more generic.

tillrohrmann · 2025-12-19T14:02:29Z

crates/worker/src/partition/shuffle.rs

+        let mut stream =
+            state_machine::StateMachine::new(metadata, ingestion_client, outbox_reader, hint_rx);


Nit: Is stream still the best variable name?

tillrohrmann · 2025-12-19T14:08:48Z

crates/worker/src/partition/shuffle.rs

+            loop {
+                match &mut self.state {
+                    State::Idle => {
+                        loop {


Is this inner loop needed or could it be removed and let the outer loop handle things?

You are right, it's totally not needed.

- Use IngestionClient instead of bifrost to write to partitions logs - Remove deprecated `delete_invocation`

Avoid direct writes to bifrost in shuffler by using a dedicated ingestion-client instance.

muhamadazmy force-pushed the pr4024 branch from 819226a to d433166 Compare November 18, 2025 14:50

muhamadazmy changed the title ~~[wip] shuffler with ingress client~~ Use ingress-client in the Shuffler Nov 18, 2025

muhamadazmy marked this pull request as ready for review November 18, 2025 14:51

muhamadazmy requested a review from tillrohrmann November 18, 2025 14:51

muhamadazmy force-pushed the pr4024 branch 2 times, most recently from 8008c9c to 91f6046 Compare November 25, 2025 12:54

tillrohrmann approved these changes Nov 25, 2025

View reviewed changes

crates/worker/src/lib.rs Outdated Show resolved Hide resolved

crates/worker/src/partition/leadership/mod.rs Outdated Show resolved Hide resolved

crates/worker/src/partition/shuffle.rs Outdated Show resolved Hide resolved

crates/worker/src/lib.rs Outdated Show resolved Hide resolved

tillrohrmann reviewed Nov 26, 2025

View reviewed changes

muhamadazmy force-pushed the pr4024 branch 15 times, most recently from d6e9955 to 8c0797e Compare December 2, 2025 09:00

muhamadazmy force-pushed the pr4024 branch 5 times, most recently from fb4ba71 to df301ad Compare December 17, 2025 11:38

muhamadazmy mentioned this pull request Dec 17, 2025

[Ingest] Use bilrost encoding for ingestion message #4113

Merged

muhamadazmy force-pushed the pr4024 branch 9 times, most recently from 8921aa2 to ef64a0b Compare December 18, 2025 13:14

muhamadazmy requested a review from tillrohrmann December 18, 2025 13:40

muhamadazmy force-pushed the pr4024 branch from ef64a0b to 6df166b Compare December 19, 2025 10:22

tillrohrmann approved these changes Dec 19, 2025

View reviewed changes

muhamadazmy force-pushed the pr4024 branch 2 times, most recently from 58a3264 to 68b90af Compare December 29, 2025 13:02

muhamadazmy force-pushed the pr4024 branch 4 times, most recently from 3c9a9a4 to b5a503b Compare January 6, 2026 09:06

muhamadazmy added 2 commits January 6, 2026 10:41

[AdminAPI] Use IngestionClient for invocation and state mgmt

b648647

- Use IngestionClient instead of bifrost to write to partitions logs - Remove deprecated `delete_invocation`

Use ingestion-client in the Shuffler

8836d3c

Avoid direct writes to bifrost in shuffler by using a dedicated ingestion-client instance.

muhamadazmy force-pushed the pr4024 branch from b5a503b to 8836d3c Compare January 6, 2026 09:41

muhamadazmy merged commit d3d9e6a into restatedev:main Jan 6, 2026
58 checks passed

muhamadazmy deleted the pr4024 branch January 6, 2026 10:11

github-actions bot locked and limited conversation to collaborators Jan 6, 2026


		let partition_store_manager = PartitionStoreManager::create().await?;

		let ingress = IngestionClient::new(

		let mut stream =
		state_machine::StateMachine::new(metadata, ingestion_client, outbox_reader, hint_rx);

Use ingestion-client in the Shuffler #4024

Use ingestion-client in the Shuffler #4024

Uh oh!

Conversation

muhamadazmy commented Nov 17, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

tillrohrmann left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

tillrohrmann commented Nov 25, 2025

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 left a comment

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

Use ingestion-client in the `Shuffler` #4024

Use ingestion-client in the `Shuffler` #4024

muhamadazmy commented Nov 17, 2025 •

edited

Loading