[Limiter] Persist rule book per partition via UpsertRuleBook log entry

Wires the cluster-global rule book into the partition processor state
machine so leader-driven distribution (Step 5) has somewhere to land:

  * `Command::UpsertRuleBook(UpsertRuleBook { partition_key_range,
    rule_book: Bytes })` — new wal-protocol command. The payload is
    bilrost-encoded `RuleBook` carried as opaque `Bytes` (precedent:
    `Command::VQSchedulerDecisions`) so flexbuffers-based `Envelope`
    serde does not need to drag full serde derive through every
    limiter type.

  * `ReadFsmTable::get_rule_book` / `WriteFsmTable::put_rule_book`
    (both `*Since v1.7.0*`) backed by a new FSM slot `RULE_BOOK = 9`
    in the partition store. Each partition writes the same logical
    rule book; readback on PP boot gives leader transitions the right
    state without an extra metadata-store round trip.

  * `RuleBook` gains `StorageEncode`/`StorageDecode` (bilrost) plus
    `bilrost_encode_to_bytes` / `bilrost_decode` helpers.

  * `StateMachine` carries `rule_book: RuleBook` in-memory; loaded
    from FSM table at PP boot. The `Command::UpsertRuleBook` apply
    path bilrost-decodes the bytes, idempotency-checks the version
    (skips when not strictly newer), diffs against the previous
    in-memory book, persists via `put_rule_book` within the same
    transaction, updates in-memory state, and emits
    `Action::RulesUpdated(Vec<RuleUpdate>)` when non-empty.

  * `Action::RulesUpdated` is dispatched in `leader_state` via a new
    `SchedulerService::on_rules_updated` API that forwards the batch
    through the existing resource-manager mpsc to `UserLimiter`.
    Followers don't dispatch actions, matching the "only the leader's
    UserLimiter is live" design.

  * `restate-wal-protocol` does NOT depend on `restate-limiter` — the
    opaque-bytes payload keeps it isolated.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Author: tillrohrmann

Cargo.lock

@@ -388,6 +388,20 @@ impl RuleBook {
         self.diff(&Self::empty())
     }
 
+    /// Encode the rule book as bilrost-encoded bytes. Used by callers that
+    /// want to embed the book in another wire format opaquely (e.g. the
+    /// `Command::UpsertRuleBook` envelope) without taking a direct
+    /// dependency on the `bilrost` crate.
+    pub fn bilrost_encode_to_bytes(&self) -> bytes::Bytes {
+        bilrost::Message::encode_to_bytes(self)
+    }
+
+    /// Decode a rule book from bilrost-encoded bytes. Inverse of
+    /// [`Self::bilrost_encode_to_bytes`].
+    pub fn bilrost_decode<B: bytes::Buf>(buf: B) -> Result<Self, bilrost::DecodeError> {
+        <Self as bilrost::OwnedMessage>::decode(buf)
+    }
+
     /// Test/internal helper.
     #[cfg(test)]
     pub fn from_parts(version: Version, rules: HashMap<RuleId, PersistedRule>) -> Self {
@@ -497,6 +511,37 @@ impl Versioned for RuleBook {
     }
 }
 
+mod storage {
+    use bytes::BytesMut;
+
+    use restate_types::storage::{
+        StorageCodecKind, StorageDecode, StorageDecodeError, StorageEncode, StorageEncodeError,
+        decode, encode,
+    };
+
+    use super::RuleBook;
+
+    impl StorageEncode for RuleBook {
+        fn encode(&self, buf: &mut BytesMut) -> Result<(), StorageEncodeError> {
+            encode::encode_bilrost(self, buf)
+        }
+
+        fn default_codec(&self) -> StorageCodecKind {
+            StorageCodecKind::Bilrost
+        }
+    }
+
+    impl StorageDecode for RuleBook {
+        fn decode<B: bytes::Buf>(
+            buf: &mut B,
+            kind: StorageCodecKind,
+        ) -> Result<Self, StorageDecodeError> {
+            assert_eq!(kind, StorageCodecKind::Bilrost);
+            decode::decode_bilrost(buf)
+        }
+    }
+}
+
 // `RulePattern<ReString>` rides through bilrost via its `Display`/`FromStr`
 // round-trip — same trick as `InvocationId`. This avoids spreading bilrost
 // glue across `Pattern`/`RestrictedValue` and keeps the wire form

crates/limiter/src/rule_book.rs

@@ -17,6 +17,7 @@ restate-workspace-hack = { workspace = true }
 restate-clock = { workspace = true }
 restate-core = { workspace = true }
 restate-errors = { workspace = true }
+restate-limiter = { workspace = true, features = ["rule-book"] }
 restate-memory = { workspace = true }
 restate-object-store-util = { workspace = true }
 restate-rocksdb = { workspace = true }

crates/partition-store/Cargo.toml

@@ -8,6 +8,7 @@
 // the Business Source License, use of this software will be governed
 // by the Apache License, Version 2.0.
 
+use restate_limiter::RuleBook;
 use restate_storage_api::Result;
 use restate_storage_api::fsm_table::{
     CachedEpochMetadata, PartitionDurability, ReadFsmTable, SequenceNumber, WriteFsmTable,
@@ -66,6 +67,13 @@ pub(crate) mod fsm_variable {
     /// deployments.
     /// *Since v1.6.3*
     pub(crate) const JC_ORPHAN_CLEANUP_DONE: u64 = 8;
+
+    /// Cluster-global rule book persisted per-partition. Each partition writes
+    /// the same logical rule book (via `Command::UpsertRuleBook` log entries),
+    /// and reads it back on PP startup so leader transitions inherit the same
+    /// rule set without an extra metadata-store round trip.
+    /// *Since v1.7.0*
+    pub(crate) const RULE_BOOK: u64 = 9;
 }
 
 fn get<T: PartitionStoreProtobufValue, S: StorageAccess>(
@@ -196,6 +204,11 @@ impl ReadFsmTable for PartitionStore {
         let key = create_key(self.partition_id(), fsm_variable::PARTITION_CONFIG_STATE);
         self.get_value_storage_codec(key)
     }
+
+    async fn get_rule_book(&mut self) -> Result<Option<RuleBook>> {
+        let key = create_key(self.partition_id(), fsm_variable::RULE_BOOK);
+        self.get_value_storage_codec(key)
+    }
 }
 
 impl WriteFsmTable for PartitionStoreTransaction<'_> {
@@ -253,4 +266,9 @@ impl WriteFsmTable for PartitionStoreTransaction<'_> {
         let key = create_key(self.partition_id(), fsm_variable::PARTITION_CONFIG_STATE);
         self.put_kv_storage_codec(key, state)
     }
+
+    fn put_rule_book(&mut self, rule_book: &RuleBook) -> Result<()> {
+        let key = create_key(self.partition_id(), fsm_variable::RULE_BOOK);
+        self.put_kv_storage_codec(key, rule_book)
+    }
 }

crates/partition-store/src/fsm_table/mod.rs

@@ -14,7 +14,7 @@ test-util = []
 restate-workspace-hack = { workspace = true }
 
 restate-clock = { workspace = true }
-restate-limiter = { workspace = true, features = ["bilrost"] }
+restate-limiter = { workspace = true, features = ["bilrost", "rule-book"] }
 restate-memory = { workspace = true }
 restate-sharding = { workspace = true }
 restate-types = { workspace = true }

crates/storage-api/Cargo.toml

@@ -12,6 +12,7 @@ use std::future::Future;
 
 use bytes::BytesMut;
 
+use restate_limiter::RuleBook;
 use restate_types::identifiers::LeaderEpoch;
 use restate_types::logs::Lsn;
 use restate_types::message::MessageIndex;
@@ -48,6 +49,12 @@ pub trait ReadFsmTable {
     fn get_partition_config_state(
         &mut self,
     ) -> impl Future<Output = Result<Option<CachedEpochMetadata>>> + Send + '_;
+
+    /// The rule book persisted for this partition. Returns `None` when the
+    /// partition has not yet observed any rule-book update, in which case
+    /// callers should treat it as the empty default.
+    /// *Since v1.7.0*
+    fn get_rule_book(&mut self) -> impl Future<Output = Result<Option<RuleBook>>> + Send + '_;
 }
 
 pub trait WriteFsmTable {
@@ -64,6 +71,10 @@ pub trait WriteFsmTable {
     fn put_schema(&mut self, schema: &Schema) -> Result<()>;
 
     fn put_partition_config_state(&mut self, state: &CachedEpochMetadata) -> Result<()>;
+
+    /// Persist the rule book for this partition.
+    /// *Since v1.7.0*
+    fn put_rule_book(&mut self, rule_book: &RuleBook) -> Result<()>;
 }
 
 #[derive(Debug, Clone, Copy, derive_more::From, derive_more::Into)]

crates/storage-api/src/fsm_table/mod.rs

@@ -15,6 +15,7 @@ use std::pin::Pin;
 use std::future::poll_fn;
 use std::task::Poll;
 
+use restate_limiter::RuleUpdate;
 use restate_storage_api::StorageError;
 use restate_storage_api::vqueue_table::scheduler::{RunAction, SchedulerAction, YieldAction};
 use restate_storage_api::vqueue_table::{EntryKey, ScanVQueueTable, VQueueStore};
@@ -166,6 +167,14 @@ impl<S: VQueueStore> SchedulerService<S> {
         }
     }
 
+    /// Forward a batch of rule-book updates to the embedded resource
+    /// manager. No-op when the scheduler is disabled (followers).
+    pub fn on_rules_updated(&self, updates: Vec<RuleUpdate>) {
+        if let State::Active(ref drr_scheduler) = self.state {
+            drr_scheduler.on_rules_updated(updates);
+        }
+    }
+
     /// Return reserved resources (concurrency permit + memory lease) for a given
     /// item hash if it was assigned by the scheduler.
     ///

crates/vqueues/src/scheduler.rs

@@ -243,6 +243,11 @@ impl<S: VQueueStore> DRRScheduler<S> {
         Some(permit.build(&self.resource_manager))
     }
 
+    /// Forward a batch of rule-book updates to the embedded resource manager.
+    pub fn on_rules_updated(&self, updates: Vec<restate_limiter::RuleUpdate>) {
+        self.resource_manager.on_rules_updated(updates);
+    }
+
     #[tracing::instrument(skip_all)]
     #[track_caller]
     pub fn on_inbox_event(&mut self, event: VQueueEvent) {

crates/vqueues/src/scheduler/drr.rs

@@ -25,6 +25,7 @@ use tracing::trace;
 
 use restate_futures_util::concurrency::Concurrency;
 use restate_limiter::RuleHandle;
+use restate_limiter::RuleUpdate;
 use restate_memory::{MemoryPool, NonZeroByteCount};
 use restate_storage_api::StorageError;
 use restate_storage_api::lock_table::LoadLocks;
@@ -92,6 +93,15 @@ impl ResourceManager {
         })
     }
 
+    /// Forward a batch of rule-book updates to the resource manager's
+    /// internal channel. Picked up on the next `poll_resources` tick and
+    /// applied to the `UserLimiter`.
+    pub fn on_rules_updated(&self, updates: Vec<RuleUpdate>) {
+        // Sending via `tx` (which is held alongside `rx` for keep-alive)
+        // never fails: the receiver is owned by `self`.
+        let _ = self.tx.send(ResourceManagerUpdate::RulesUpdated(updates));
+    }
+
     /// Removes the vqueue from the resource it's blocked on
     pub(super) fn remove_vqueue(&mut self, handle: VQueueHandle, blocked_resource: &ResourceKind) {
         match blocked_resource {

crates/vqueues/src/scheduler/resource_manager.rs

@@ -8,6 +8,8 @@
 // the Business Source License, use of this software will be governed
 // by the Apache License, Version 2.0.
 
+use bytes::Bytes;
+
 use restate_storage_api::fsm_table::{CurrentReplicaSetState, NextReplicaSetState};
 use restate_types::identifiers::{LeaderEpoch, PartitionId};
 use restate_types::logs::{Keys, Lsn, SequenceNumber};
@@ -170,3 +172,24 @@ pub struct UpsertSchema {
     pub partition_key_range: Keys,
     pub schema: Schema,
 }
+
+/// Consistently distribute the cluster-global rule book across partition
+/// replicas. Each partition's leader observes a node-level cache of the
+/// rule book stored in the metadata store and proposes this command when
+/// it sees a higher version than the partition's in-memory state.
+/// Followers and the leader (replaying) apply it idempotently — no-op if
+/// the carried rule book's version is not greater than the current
+/// in-memory version.
+///
+/// `rule_book` is the bilrost-encoded [`restate_limiter::RuleBook`]. It
+/// is carried as opaque bytes (same precedent as
+/// [`crate::Command::VQSchedulerDecisions`]) so the wal-protocol crate
+/// doesn't need to drag full serde derive through every limiter type.
+/// The state machine decodes once on apply.
+///
+/// Since v1.7.0.
+#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
+pub struct UpsertRuleBook {
+    pub partition_key_range: Keys,
+    pub rule_book: Bytes,
+}