From 95164d7be0d3fe9eb54611a1a6be025c484f3330 Mon Sep 17 00:00:00 2001 From: Muhamad Awad Date: Tue, 12 May 2026 17:52:05 +0200 Subject: [PATCH] wal-protocol: store UpsertRuleBookCommand as opaque bilrost bytes in v1 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary - Convert `UpsertRuleBookCommand` to a `bilrost::Message` carrying `Arc` directly, instead of a serde struct wrapping an already-bilrost-encoded `Bytes`. Eliminates the previous double encoding (serde-around-bilrost) on the proposal path. - In v1, `Command::UpsertRuleBook` now wraps an opaque `Bytes` payload (the bilrost-encoded `UpsertRuleBookCommand`) rather than the typed struct. Because `UpsertRuleBook` is `*Since v1.7.0` and still unreleased, this payload change is backward compatible. - v1→v2 compatibility no longer decodes/re-encodes the rule book; it forwards the bilrost payload via `Envelope::from_bytes_unchecked`, matching the existing pattern used for `VQueues*` / `VQSchedulerDecisions`. - `record_keys()` for the v1 variant falls back to `Keys::Single(self.partition_key())` since the `partition_key_range` is no longer reachable without decoding. - Leader proposer encodes the command into the shared arena buffer before proposing, avoiding an extra allocation. - State machine now ignores `UpsertRuleBook` messages whose `partition_key_range` does not overlap the local partition's range, mirroring how other range-targeted commands are filtered. --- crates/wal-protocol/src/control.rs | 33 ++++++++++++++++--- crates/wal-protocol/src/v1.rs | 17 ++++++++-- crates/wal-protocol/src/v2/commands.rs | 31 +++-------------- crates/wal-protocol/src/v2/compatibility.rs | 27 ++++----------- .../src/partition/leadership/leader_state.rs | 20 +++++++---- 5 files changed, 68 insertions(+), 60 deletions(-) diff --git a/crates/wal-protocol/src/control.rs b/crates/wal-protocol/src/control.rs index 0fcc794268..f81ffd962b 100644 --- a/crates/wal-protocol/src/control.rs +++ b/crates/wal-protocol/src/control.rs @@ -8,7 +8,12 @@ // the Business Source License, use of this software will be governed // by the Apache License, Version 2.0. -use bytes::Bytes; +use std::sync::Arc; + +use bytes::{Buf, BufMut, Bytes}; + +use restate_encoding::Arced; +use restate_limiter::RuleBook; use restate_storage_api::fsm_table::{CurrentReplicaSetState, NextReplicaSetState}; use restate_types::identifiers::{LeaderEpoch, PartitionId}; use restate_types::logs::{HasRecordKeys, Keys, Lsn, SequenceNumber}; @@ -241,8 +246,28 @@ impl HasRecordKeys for UpsertSchemaCommand { /// The state machine decodes once on apply. /// /// Since v1.7.0. -#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] +#[derive(Debug, Clone, bilrost::Message)] pub struct UpsertRuleBookCommand { - pub partition_key_range: KeyRange, - pub rule_book: Bytes, + #[bilrost(tag(1), encoding(Arced))] + pub rule_book: Arc, +} + +bilrost_storage_encode_decode!(UpsertRuleBookCommand); + +impl UpsertRuleBookCommand { + pub fn bilrost_encode(&self, b: &mut B) -> Result<(), bilrost::EncodeError> { + bilrost::Message::encode(self, b) + } + + pub fn encoded_len(&self) -> usize { + bilrost::Message::encoded_len(self) + } + + pub fn bilrost_encode_to_bytes(&self) -> Bytes { + bilrost::Message::encode_to_bytes(self) + } + + pub fn bilrost_decode(buf: B) -> Result { + bilrost::OwnedMessage::decode(buf) + } } diff --git a/crates/wal-protocol/src/v1.rs b/crates/wal-protocol/src/v1.rs index 3c60419c46..bf9750fcbe 100644 --- a/crates/wal-protocol/src/v1.rs +++ b/crates/wal-protocol/src/v1.rs @@ -19,11 +19,12 @@ use restate_types::invocation::{ }; use restate_types::logs::{self, HasRecordKeys, Keys, MatchKeyQuery}; use restate_types::message::MessageIndex; +use restate_types::sharding::KeyRange; use restate_types::state_mut::ExternalStateMutation; use crate::control::{ - AnnounceLeaderCommand, UpdatePartitionDurabilityCommand, UpsertRuleBookCommand, - UpsertSchemaCommand, VersionBarrierCommand, + AnnounceLeaderCommand, UpdatePartitionDurabilityCommand, UpsertSchemaCommand, + VersionBarrierCommand, }; use crate::timer::TimerKeyValue; @@ -193,8 +194,9 @@ pub enum Command { /// Upsert the cluster-global rule book for consistent rules across /// replicas; the apply path persists it to the partition store and /// notifies the leader's `UserLimiter` of the diff. + /// /// *Since v1.7.0 - UpsertRuleBook(UpsertRuleBookCommand), + UpsertRuleBook(UpsertRuleBookCommandWrapper), // # Commands for VQueues management // ---------------------------------- /// A command to attempt a run an entry in the vqueue (invocation, or otherwise) @@ -275,3 +277,12 @@ impl MatchKeyQuery for Envelope { self.record_keys().matches_key_query(query) } } + +/// A temporary wrapper for [`UpsertRuleBookCommand`] +/// only used in v1 to only supply the partition_key_range +#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] +pub struct UpsertRuleBookCommandWrapper { + pub partition_key_range: KeyRange, + /// Bytes are bilrost encoded [`UpsertRuleBookCommand`] + pub command: Bytes, +} diff --git a/crates/wal-protocol/src/v2/commands.rs b/crates/wal-protocol/src/v2/commands.rs index 15e56f01cf..cbecfd0e18 100644 --- a/crates/wal-protocol/src/v2/commands.rs +++ b/crates/wal-protocol/src/v2/commands.rs @@ -8,29 +8,24 @@ // the Business Source License, use of this software will be governed // by the Apache License, Version 2.0. -// Re-epxort vqueues commands -pub use crate::vqueues::{VQueuesPauseCommand, VQueuesResumeCommand}; -pub use restate_storage_api::vqueue_table::scheduler::SchedulerDecisionsCommand; - -use std::sync::Arc; - use serde::{Deserialize, Serialize}; -use restate_encoding::Arced; -use restate_limiter::RuleBook; +pub use restate_storage_api::vqueue_table::scheduler::SchedulerDecisionsCommand; use restate_types::{ bilrost_storage_encode_decode, flexbuffers_storage_encode_decode, identifiers::{WithInvocationId, WithPartitionKey}, invocation, logs::{HasRecordKeys, Keys}, message::MessageIndex, - sharding::KeyRange, state_mut, }; use super::sealed::Sealed; use super::{Command, CommandKind}; -use crate::timer::{self}; +pub use crate::control::UpsertRuleBookCommand; +use crate::timer; +// Re-epxort vqueues commands +pub use crate::vqueues::{VQueuesPauseCommand, VQueuesResumeCommand}; pub use crate::control::{ AnnounceLeaderCommand, UpdatePartitionDurabilityCommand, UpsertSchemaCommand, @@ -340,22 +335,6 @@ impl HasRecordKeys for ProxyThroughCommand { } } -#[derive(Debug, Clone, bilrost::Message)] -pub struct UpsertRuleBookCommand { - #[bilrost(tag(1))] - pub partition_key_range: KeyRange, - #[bilrost(tag(2), encoding(Arced))] - pub rule_book: Arc, -} - -bilrost_storage_encode_decode!(UpsertRuleBookCommand); - -impl HasRecordKeys for UpsertRuleBookCommand { - fn record_keys(&self) -> Keys { - Keys::RangeInclusive(self.partition_key_range.into()) - } -} - // end types // define record types diff --git a/crates/wal-protocol/src/v2/compatibility.rs b/crates/wal-protocol/src/v2/compatibility.rs index b70d4969c5..9f4e0d6715 100644 --- a/crates/wal-protocol/src/v2/compatibility.rs +++ b/crates/wal-protocol/src/v2/compatibility.rs @@ -8,12 +8,7 @@ // the Business Source License, use of this software will be governed // by the Apache License, Version 2.0. -use std::sync::Arc; - -use bilrost::OwnedMessage; - use restate_encoding::U128; -use restate_limiter::RuleBook; use restate_storage_api::deduplication_table::{ DedupInformation, DedupSequenceNumber, EpochSequenceNumber, ProducerId, }; @@ -23,10 +18,7 @@ use restate_util_string::ReString; use super::{Raw, commands}; use crate::{ v1, - v2::{ - self, Envelope, - commands::{TruncateOutboxCommand, UpsertRuleBookCommand}, - }, + v2::{self, Envelope, commands::TruncateOutboxCommand}, }; impl TryFrom for v2::Envelope { @@ -139,17 +131,12 @@ impl TryFrom for v2::Envelope { } v1::Command::UpsertSchema(payload) => Envelope::new(dedup, payload).into_raw(), v1::Command::VersionBarrier(payload) => Envelope::new(dedup, payload).into_raw(), - v1::Command::UpsertRuleBook(payload) => { - let rule_book = RuleBook::decode(payload.rule_book)?; - Envelope::new( - dedup, - UpsertRuleBookCommand { - partition_key_range: payload.partition_key_range, - rule_book: Arc::new(rule_book), - }, - ) - .into_raw() - } + v1::Command::UpsertRuleBook(wrapper) => Envelope::from_bytes_unchecked( + v2::CommandKind::UpsertRuleBook, + StorageCodecKind::Bilrost, + dedup, + wrapper.command, + ), v1::Command::VQSchedulerDecisions(payload) => Envelope::from_bytes_unchecked( v2::CommandKind::VQSchedulerDecisions, StorageCodecKind::Bilrost, diff --git a/crates/worker/src/partition/leadership/leader_state.rs b/crates/worker/src/partition/leadership/leader_state.rs index 00bdb24750..69d47a4724 100644 --- a/crates/worker/src/partition/leadership/leader_state.rs +++ b/crates/worker/src/partition/leadership/leader_state.rs @@ -50,6 +50,7 @@ use restate_vqueues::scheduler::Decisions; use restate_vqueues::{SchedulerService, VQueuesMeta}; use restate_wal_protocol::Command; use restate_wal_protocol::control::UpsertSchemaCommand; +use restate_wal_protocol::v1::UpsertRuleBookCommandWrapper; use restate_worker_api::invoker::InvokerHandle; use restate_worker_api::resources::ReservedResources; use restate_worker_api::{SchedulerStatusEntry, UserLimitCounterEntry}; @@ -454,7 +455,7 @@ impl LeaderState { .await?; } } - ActionEffect::UpsertRuleBook(book) => { + ActionEffect::UpsertRuleBook(rule_book) => { // todo(tillrohrmann) also enable the feature once the partition has been migrated // to use vqueues and then rolling back to v1.7 if Configuration::pinned() @@ -462,15 +463,20 @@ impl LeaderState { .experimental .is_vqueues_enabled() { + let cmd = + restate_wal_protocol::control::UpsertRuleBookCommand { rule_book }; + + arena.reserve(cmd.encoded_len()); + // safe to unwrap because we reserved enough space + cmd.bilrost_encode(&mut arena).unwrap(); + self.self_proposer .self_propose( self.partition_key_range.start(), - Command::UpsertRuleBook( - restate_wal_protocol::control::UpsertRuleBookCommand { - partition_key_range: self.partition_key_range, - rule_book: book.bilrost_encode_to_bytes(), - }, - ), + Command::UpsertRuleBook(UpsertRuleBookCommandWrapper { + partition_key_range: self.partition_key_range, + command: arena.split().freeze(), + }), ) .await?; }