From 7f2952924b640ffb5f4efcc01f587894a0dacade Mon Sep 17 00:00:00 2001 From: Pascal Thesing Date: Mon, 6 Jul 2026 00:39:56 +0200 Subject: [PATCH] Add weighted fair scheduling for vqueues via scope rules --- cli/src/commands/rules/list.rs | 7 +- cli/src/commands/rules/mod.rs | 12 +- cli/src/commands/rules/set.rs | 14 +- crates/limiter/src/rule_book.rs | 3 + crates/limiter/src/user_limits.rs | 59 +- .../storage-query-datafusion/src/rules/row.rs | 3 + .../src/rules/schema.rs | 4 + crates/types/src/config/common.rs | 8 + crates/types/src/partitions/features.rs | 15 + crates/vqueues/src/lib.rs | 2 +- crates/vqueues/src/scheduler.rs | 238 +++- crates/vqueues/src/scheduler/drr.rs | 1103 ++++++++++++++++- crates/vqueues/src/scheduler/eligible.rs | 428 +++++-- .../vqueues/src/scheduler/resource_manager.rs | 25 +- .../resource_manager/grouped_waiters.rs | 354 ++++++ .../src/scheduler/resource_manager/invoker.rs | 143 ++- crates/worker-api/src/processor/features.rs | 19 + crates/worker/src/partition/leadership/mod.rs | 23 + .../state_machine/entries/call_commands.rs | 12 +- .../lifecycle/version_barrier.rs | 4 + .../worker/src/partition/state_machine/mod.rs | 32 +- .../src/partition/state_machine/tests/mod.rs | 79 ++ 22 files changed, 2422 insertions(+), 165 deletions(-) create mode 100644 crates/vqueues/src/scheduler/resource_manager/grouped_waiters.rs diff --git a/cli/src/commands/rules/list.rs b/cli/src/commands/rules/list.rs index 7d24af9703..1715501352 100644 --- a/cli/src/commands/rules/list.rs +++ b/cli/src/commands/rules/list.rs @@ -41,7 +41,7 @@ async fn list(env: &CliEnv, opts: &List) -> Result<()> { let client = DataFusionHttpClient::new(env).await?; let rows: Vec = client .run_json_query( - "SELECT pattern, concurrency, description, disabled, version, last_modified \ + "SELECT pattern, concurrency, scheduling_weight, description, disabled, version, last_modified \ FROM sys_rules ORDER BY pattern" .to_string(), ) @@ -57,13 +57,14 @@ async fn list(env: &CliEnv, opts: &List) -> Result<()> { table.set_styled_header(vec![ "PATTERN", "CONCURRENCY", + "WEIGHT", "DISABLED", "DESCRIPTION", "VERSION", "LAST MODIFIED", ]); } else { - table.set_styled_header(vec!["PATTERN", "CONCURRENCY", "DISABLED"]); + table.set_styled_header(vec!["PATTERN", "CONCURRENCY", "WEIGHT", "DISABLED"]); } for row in rows { @@ -73,6 +74,7 @@ async fn list(env: &CliEnv, opts: &List) -> Result<()> { table.add_row(vec![ Cell::new(row.pattern), Cell::new(render_concurrency(row.concurrency)), + Cell::new(row.scheduling_weight.unwrap_or(1)), Cell::new(disabled), Cell::new(row.description.unwrap_or_default()), Cell::new(row.version), @@ -82,6 +84,7 @@ async fn list(env: &CliEnv, opts: &List) -> Result<()> { table.add_row(vec![ Cell::new(row.pattern), Cell::new(render_concurrency(row.concurrency)), + Cell::new(row.scheduling_weight.unwrap_or(1)), Cell::new(disabled), ]); } diff --git a/cli/src/commands/rules/mod.rs b/cli/src/commands/rules/mod.rs index d14a1f4759..1d1f8adcfb 100644 --- a/cli/src/commands/rules/mod.rs +++ b/cli/src/commands/rules/mod.rs @@ -53,6 +53,8 @@ pub(crate) struct RuleRow { #[serde(default)] pub concurrency: Option, #[serde(default)] + pub scheduling_weight: Option, + #[serde(default)] pub description: Option, pub disabled: bool, pub version: u32, @@ -65,6 +67,11 @@ impl RuleRow { fn concurrency(&self) -> Option { self.concurrency.and_then(NonZeroU32::new) } + + /// The scheduling weight as a `NonZeroU32` (the runtime shape). + fn scheduling_weight(&self) -> Option { + self.scheduling_weight.and_then(std::num::NonZeroU32::new) + } } /// Renders a concurrency limit for display (`unlimited` when unset). @@ -88,7 +95,7 @@ pub(crate) async fn fetch_rule( canonical_pattern: &str, ) -> Result> { let query = format!( - "SELECT pattern, concurrency, description, disabled, version, last_modified \ + "SELECT pattern, concurrency, scheduling_weight, description, disabled, version, last_modified \ FROM sys_rules WHERE pattern = '{}'", escape_sql(canonical_pattern) ); @@ -142,7 +149,8 @@ pub(crate) async fn toggle_disabled(env: &CliEnv, pattern: &str, disabled: bool) let client = AdminClient::new(env).await?; let request = UpsertRuleRequest { pattern, - limits: UserLimits::new(current.concurrency()), + limits: UserLimits::new(current.concurrency()) + .with_scheduling_weight(current.scheduling_weight()), description: current.description.clone(), disabled, precondition: Precondition::Matches(Version::from(current.version)), diff --git a/cli/src/commands/rules/set.rs b/cli/src/commands/rules/set.rs index 5e81621989..407a78bd7e 100644 --- a/cli/src/commands/rules/set.rs +++ b/cli/src/commands/rules/set.rs @@ -37,6 +37,14 @@ pub struct Set { #[clap(long)] unlimited: bool, + /// Scheduling weight (>= 1) for this rule's scope in the weighted + /// round-robin scheduler: weight N receives N dispatch slots per cycle + /// relative to weight-1 groups. On an existing rule, omitting this leaves + /// the current weight unchanged. Requires the vqueues scheduler; only + /// scope-level exact patterns are consulted. + #[clap(long)] + weight: Option, + /// Description for the rule #[clap(long)] description: Option, @@ -62,7 +70,8 @@ pub async fn run_set(State(env): State, opts: &Set) -> Result<()> { None } else { opts.concurrency - }), + }) + .with_scheduling_weight(opts.weight), description: opts.description.clone(), disabled: opts.disabled, precondition: Precondition::DoesNotExist, @@ -82,9 +91,10 @@ pub async fn run_set(State(env): State, opts: &Set) -> Result<()> { .description .clone() .or_else(|| rule.description.clone()); + let scheduling_weight = opts.weight.or_else(|| rule.scheduling_weight()); UpsertRuleRequest { pattern, - limits: UserLimits::new(concurrency), + limits: UserLimits::new(concurrency).with_scheduling_weight(scheduling_weight), description, disabled: rule.disabled, precondition: Precondition::Matches(Version::from(rule.version)), diff --git a/crates/limiter/src/rule_book.rs b/crates/limiter/src/rule_book.rs index a02c2b3663..414e9365cb 100644 --- a/crates/limiter/src/rule_book.rs +++ b/crates/limiter/src/rule_book.rs @@ -483,6 +483,7 @@ mod tests { RuleUpsert { limits: UserLimits { concurrency: NonZeroU32::new(concurrency), + scheduling_weight: None, }, description: None, disabled: false, @@ -516,6 +517,7 @@ mod tests { PersistedRule { limits: UserLimits { concurrency: NonZeroU32::new(1000), + scheduling_weight: None, }, description: Some("global default".to_owned()), disabled: false, @@ -528,6 +530,7 @@ mod tests { PersistedRule { limits: UserLimits { concurrency: NonZeroU32::new(10), + scheduling_weight: None, }, description: None, disabled: true, diff --git a/crates/limiter/src/user_limits.rs b/crates/limiter/src/user_limits.rs index 7ea561edf0..dc2bb565ec 100644 --- a/crates/limiter/src/user_limits.rs +++ b/crates/limiter/src/user_limits.rs @@ -45,11 +45,35 @@ pub struct UserLimits { )] #[cfg_attr(feature = "schema", schema(value_type = Option, minimum = 1))] pub concurrency: Option, + + /// Scheduling weight of this rule's scope in the scheduler's weighted + /// round-robin: a scope with weight N receives N dispatch slots per cycle + /// relative to weight-1 groups, regardless of how many queues it has. + /// `None` means the default weight of 1. Only scope-level exact patterns + /// are consulted. Requires the vqueues scheduler. + /// + /// Upserts replace the whole limits object: omitting this field resets + /// the weight (the CLI preserves it by re-reading the current rule). + #[cfg_attr(feature = "bilrost", bilrost(tag(2)))] + #[cfg_attr( + feature = "serde", + serde(default, skip_serializing_if = "Option::is_none") + )] + #[cfg_attr(feature = "schema", schema(value_type = Option, minimum = 1))] + pub scheduling_weight: Option, } impl UserLimits { pub fn new(concurrency: Option) -> Self { - Self { concurrency } + Self { + concurrency, + scheduling_weight: None, + } + } + + pub fn with_scheduling_weight(mut self, scheduling_weight: Option) -> Self { + self.scheduling_weight = scheduling_weight; + self } } @@ -69,3 +93,36 @@ pub enum RuleUpdate { /// Remove a rule by its pattern. Remove { pattern: RulePattern }, } + +#[cfg(test)] +mod tests { + use std::num::NonZeroU32; + + use bilrost::{Message, OwnedMessage}; + + use super::*; + + /// Wire compatibility: pre-weight encodings (concurrency only) decode with + /// `scheduling_weight = None`, and a default weight is omitted on the wire + /// so old readers see the exact pre-weight byte layout. + #[test] + fn scheduling_weight_wire_compat() { + // "old" shape: only tag(1) present + let old = UserLimits::new(NonZeroU32::new(100)); + let old_bytes = old.encode_to_bytes(); + let decoded = ::decode(old_bytes.clone()).unwrap(); + assert_eq!(decoded.scheduling_weight, None); + assert_eq!(decoded.concurrency, NonZeroU32::new(100)); + + // new shape round-trips + let new = UserLimits::new(NonZeroU32::new(100)).with_scheduling_weight(NonZeroU32::new(10)); + let new_bytes = new.encode_to_bytes(); + let decoded = ::decode(new_bytes).unwrap(); + assert_eq!(decoded.scheduling_weight, NonZeroU32::new(10)); + + // None weight encodes identically to the old shape (old readers + // decoding new writers see no unknown data) + let new_default = UserLimits::new(NonZeroU32::new(100)); + assert_eq!(new_default.encode_to_bytes(), old_bytes); + } +} diff --git a/crates/storage-query-datafusion/src/rules/row.rs b/crates/storage-query-datafusion/src/rules/row.rs index 421be17a9c..6e767ab6f8 100644 --- a/crates/storage-query-datafusion/src/rules/row.rs +++ b/crates/storage-query-datafusion/src/rules/row.rs @@ -23,6 +23,9 @@ pub(crate) fn append_rule_row( if let Some(concurrency) = rule.limits.concurrency { row.concurrency(concurrency.get()); } + if let Some(weight) = rule.limits.scheduling_weight { + row.scheduling_weight(weight.get()); + } if let Some(description) = rule.description.as_deref() { row.description(description); } diff --git a/crates/storage-query-datafusion/src/rules/schema.rs b/crates/storage-query-datafusion/src/rules/schema.rs index 97df0d4d8e..f1ef03fa9b 100644 --- a/crates/storage-query-datafusion/src/rules/schema.rs +++ b/crates/storage-query-datafusion/src/rules/schema.rs @@ -20,6 +20,10 @@ define_table!(sys_rules( /// rule does not constrain concurrency. concurrency: DataType::UInt32, + /// Scheduling weight of this rule's scope in the weighted round-robin + /// scheduler. Null means the default weight of 1. + scheduling_weight: DataType::UInt32, + /// Free-form description set by the operator. description: DataType::LargeUtf8, diff --git a/crates/types/src/config/common.rs b/crates/types/src/config/common.rs index 010e56ce87..deaf277892 100644 --- a/crates/types/src/config/common.rs +++ b/crates/types/src/config/common.rs @@ -662,6 +662,14 @@ experimental! { /// Since v1.7.0 scoped_virtual_objects, + /// # Enables scope inheritance along the call chain + /// + /// When enabled, a child invocation created via ctx.call/ctx.send inherits + /// its caller's scope when the child target carries no scope of its own, + /// so scope-level rules govern the whole call tree rooted at a scoped + /// ingress invocation. Requires `vqueues`. + scope_inheritance, + /// # Enables Kafka header support for scoped invocations /// /// When enabled, Kafka subscriptions read `x-restate-scope` and diff --git a/crates/types/src/partitions/features.rs b/crates/types/src/partitions/features.rs index aa4cef55f7..ffdadb1094 100644 --- a/crates/types/src/partitions/features.rs +++ b/crates/types/src/partitions/features.rs @@ -49,6 +49,11 @@ pub enum PartitionFeatureChange { /// /// *Since v1.7.0* EnableUniqueRandomSeeds = 3, + /// Child invocations created via ctx.call inherit their caller's scope + /// when the child target has none. + /// + /// *Since v1.7.0* + EnableScopeInheritance = 4, } impl PartitionFeatureChange { @@ -65,6 +70,7 @@ impl PartitionFeatureChange { Self::EnableJournalV2 => &RESTATE_VERSION_1_6_0, Self::EnableVqueues => &RESTATE_VERSION_1_7_0, Self::EnableUniqueRandomSeeds => &RESTATE_VERSION_1_7_0, + Self::EnableScopeInheritance => &RESTATE_VERSION_1_7_0, } } @@ -77,6 +83,9 @@ impl PartitionFeatureChange { Self::EnableUniqueRandomSeeds => { !std::mem::replace(&mut features.unique_random_seeds, true) } + Self::EnableScopeInheritance => { + !std::mem::replace(&mut features.scope_inheritance, true) + } } } } @@ -121,6 +130,11 @@ pub struct PersistedFeatures { /// *Since v1.7.0* #[bilrost(tag(3))] pub unique_random_seeds: bool, + /// Child invocations inherit their caller's scope when they have none. + /// + /// *Since v1.7.0* + #[bilrost(tag(4))] + pub scope_inheritance: bool, } impl PersistedFeatures { @@ -132,6 +146,7 @@ impl PersistedFeatures { self.journal_v2.then_some("journal_v2"), self.vqueues.then_some("vqueues"), self.unique_random_seeds.then_some("unique_random_seeds"), + self.scope_inheritance.then_some("scope_inheritance"), ] .into_iter() .flatten() diff --git a/crates/vqueues/src/lib.rs b/crates/vqueues/src/lib.rs index e3abaf6001..d17811a6ab 100644 --- a/crates/vqueues/src/lib.rs +++ b/crates/vqueues/src/lib.rs @@ -21,7 +21,7 @@ use std::time::Duration; pub use cache::{VQueueHandle, VQueuesMeta, VQueuesMetaCache}; pub use metric_definitions::describe_metrics; pub use restate_worker_api::{ResourceKind, SchedulingStatus, VQueueSchedulerStatus}; -pub use scheduler::{ResourceManager, SchedulerService}; +pub use scheduler::{ResourceManager, SchedulerService, ScopeWeights, scope_weight_resolver}; pub use util::*; use smallvec::SmallVec; diff --git a/crates/vqueues/src/scheduler.rs b/crates/vqueues/src/scheduler.rs index 966a50aec7..fb58e2d4ee 100644 --- a/crates/vqueues/src/scheduler.rs +++ b/crates/vqueues/src/scheduler.rs @@ -9,13 +9,13 @@ // by the Apache License, Version 2.0. use std::collections::BTreeMap; -use std::num::NonZeroU16; +use std::num::NonZeroU32; use std::pin::Pin; use std::future::poll_fn; use std::task::Poll; -use restate_limiter::RuleUpdate; +use restate_limiter::{Pattern, RulePattern, RuleUpdate}; use restate_storage_api::StorageError; use restate_storage_api::vqueue_table::scheduler::{RunAction, SchedulerAction, YieldAction}; use restate_storage_api::vqueue_table::{EntryKey, EntryMetadata, ScanVQueueTable, VQueueStore}; @@ -41,8 +41,58 @@ mod resource_manager; mod vqueue_state; // Re-exports +pub use eligible::{SchedulingGroup, WeightResolver}; pub use resource_manager::ResourceManager; +/// Live map of scope name → scheduling weight, fed from scope-level exact rules +/// (`restate rules set --weight N`). Shared between the scheduler service +/// (which maintains it from rule updates) and the [`WeightResolver`] closures +/// used by the eligibility ring and the grouped waiter lists. +pub type ScopeWeights = std::sync::Arc>>; + +/// Builds the weight resolver over a shared scope-weights map: scope groups get +/// their rule weight (default 1); service groups always run at weight 1 — +/// configurable weights are scope-keyed only. +pub fn scope_weight_resolver(weights: ScopeWeights) -> WeightResolver { + std::sync::Arc::new(move |group: &SchedulingGroup| match group { + SchedulingGroup::Scope(scope) => weights + .read() + .expect("scope weights lock poisoned") + .get(scope.as_str()) + .copied() + .unwrap_or(NonZeroU32::MIN), + SchedulingGroup::Service(_) => NonZeroU32::MIN, + }) +} + +/// Folds rule updates into the scope-weights map. Only scope-level EXACT +/// patterns carry scheduling weights (wildcards and L1/L2 limit-key patterns +/// are concurrency-only); an upsert without a weight resets to the default. +pub fn apply_rule_updates_to_scope_weights(weights: &ScopeWeights, updates: &[RuleUpdate]) { + let mut map = weights.write().expect("scope weights lock poisoned"); + for update in updates { + match update { + RuleUpdate::Upsert { pattern, limit } => { + if let RulePattern::Scope(Pattern::Exact(scope)) = pattern { + match limit.scheduling_weight { + Some(weight) => { + map.insert(scope.as_str().to_owned(), weight); + } + None => { + map.remove(scope.as_str()); + } + } + } + } + RuleUpdate::Remove { pattern } => { + if let RulePattern::Scope(Pattern::Exact(scope)) = pattern { + map.remove(scope.as_str()); + } + } + } + } +} + type UnconfirmedAssignments = hashbrown::HashMap; fn status_from_detailed_eligibility(value: DetailedEligibility) -> SchedulingStatus { @@ -119,6 +169,9 @@ enum State { pub struct SchedulerService { state: State, + /// Scope → weight map kept in sync from rule updates; read by the + /// weight-resolver closures inside the DRR scheduler. + scope_weights: ScopeWeights, } impl SchedulerService { @@ -128,6 +181,7 @@ impl SchedulerService { { Self { state: State::::Disabled, + scope_weights: ScopeWeights::default(), } } @@ -135,6 +189,8 @@ impl SchedulerService { resource_manager: ResourceManager, storage: S, vqueues_cache: &VQueuesMetaCache, + weight_resolver: WeightResolver, + scope_weights: ScopeWeights, ) -> Result where S: ScanVQueueTable, @@ -148,14 +204,18 @@ impl SchedulerService { // // Note that propose_many will error out if the number of commands is greater than the // channel's max capacity. - NonZeroU16::new(25).unwrap(), + NonZeroU32::new(25).unwrap(), // currently constant but we can make it configurable if needed - NonZeroU16::new(1000).unwrap(), + NonZeroU32::new(1000).unwrap(), resource_manager, storage, vqueues_cache.view(), + weight_resolver, ))); - Ok(Self { state }) + Ok(Self { + state, + scope_weights, + }) } pub fn on_inbox_event(&mut self, metas: VQueuesMeta<'_>, event: VQueueEvent) { @@ -165,8 +225,10 @@ impl SchedulerService { } /// Forward a batch of rule-book updates to the embedded resource - /// manager. No-op when the scheduler is disabled (followers). + /// manager, keeping the scope-weight map in sync first. No-op when the + /// scheduler is disabled (followers). pub fn on_rules_updated(&self, updates: Box<[RuleUpdate]>) { + apply_rule_updates_to_scope_weights(&self.scope_weights, &updates); if let State::Active(ref drr_scheduler) = self.state { drr_scheduler.on_rules_updated(updates); } @@ -254,3 +316,167 @@ impl SchedulerService { } } } + +#[cfg(test)] +mod scope_weight_tests { + use std::num::NonZeroU32; + + use restate_limiter::UserLimits; + use restate_types::Scope; + use restate_util_string::ReString; + + use super::*; + + fn scope_pattern(name: &'static str) -> RulePattern { + name.parse().expect("valid pattern") + } + + fn weight(n: u32) -> Option { + NonZeroU32::new(n) + } + + /// Scope-level EXACT rules feed the weight map; wildcard and limit-key + /// (L1/L2) patterns are concurrency-only and must be ignored; an upsert + /// without a weight resets to the default; removes clear. + #[test] + fn rule_updates_feed_scope_weights() { + let weights = ScopeWeights::default(); + let resolver = scope_weight_resolver(weights.clone()); + let payment = SchedulingGroup::Scope(Scope::try_non_interned("payment").unwrap()); + let other = SchedulingGroup::Scope(Scope::try_non_interned("other").unwrap()); + + // default weight before any rule + assert_eq!(resolver(&payment), NonZeroU32::MIN); + + apply_rule_updates_to_scope_weights( + &weights, + &[ + RuleUpdate::Upsert { + pattern: scope_pattern("payment"), + limit: UserLimits::new(NonZeroU32::new(100)).with_scheduling_weight(weight(10)), + }, + // wildcard scope: ignored for weights + RuleUpdate::Upsert { + pattern: scope_pattern("*"), + limit: UserLimits::new(None).with_scheduling_weight(weight(7)), + }, + // limit-key level: ignored for weights + RuleUpdate::Upsert { + pattern: scope_pattern("payment/sepa"), + limit: UserLimits::new(NonZeroU32::new(50)).with_scheduling_weight(weight(9)), + }, + ], + ); + assert_eq!(resolver(&payment), weight(10).unwrap()); + assert_eq!(resolver(&other), NonZeroU32::MIN, "wildcard must not leak"); + + // upsert WITHOUT a weight resets to the default + apply_rule_updates_to_scope_weights( + &weights, + &[RuleUpdate::Upsert { + pattern: scope_pattern("payment"), + limit: UserLimits::new(NonZeroU32::new(100)), + }], + ); + assert_eq!(resolver(&payment), NonZeroU32::MIN); + + // re-set then remove clears + apply_rule_updates_to_scope_weights( + &weights, + &[RuleUpdate::Upsert { + pattern: scope_pattern("payment"), + limit: UserLimits::new(None).with_scheduling_weight(weight(3)), + }], + ); + assert_eq!(resolver(&payment), weight(3).unwrap()); + apply_rule_updates_to_scope_weights( + &weights, + &[RuleUpdate::Remove { + pattern: scope_pattern("payment"), + }], + ); + assert_eq!(resolver(&payment), NonZeroU32::MIN); + } + + /// Service groups never resolve to configurable weights. + #[test] + fn service_groups_are_always_weight_one() { + let weights = ScopeWeights::default(); + weights + .write() + .unwrap() + .insert("SomeService".to_owned(), weight(10).unwrap()); + let resolver = scope_weight_resolver(weights); + let group = SchedulingGroup::Service(restate_types::ServiceName::new("SomeService")); + assert_eq!(resolver(&group), NonZeroU32::MIN); + } +} + +#[cfg(test)] +mod scope_weight_lifecycle_tests { + use std::num::NonZeroU32; + + use slotmap::SlotMap; + + use restate_limiter::UserLimits; + use restate_types::Scope; + use restate_util_string::ReString; + + use super::*; + use crate::scheduler::resource_manager::test_grouped_waiters; + + /// The WeightResolver docstring claims weights are "consulted whenever a + /// group is (re-)created, so rule changes take effect as soon as a group + /// cycles" — pin BOTH halves: an EXISTING group keeps its creation-time + /// weight after a rule update, and a FRESH group picks the new weight up + /// immediately. + #[test] + fn rule_update_applies_on_group_recreation_only() { + let weights = ScopeWeights::default(); + let resolver = scope_weight_resolver(weights.clone()); + let mut waiters = test_grouped_waiters(resolver); + + let mut map = SlotMap::::with_key(); + let handles: Vec<_> = (0..6).map(|_| map.insert(())).collect(); + let payment = SchedulingGroup::Scope(Scope::try_non_interned("payment").unwrap()); + let other = SchedulingGroup::Service(restate_types::ServiceName::new("other")); + + // group created at default weight 1 + waiters.push_back(handles[0], &payment); + waiters.push_back(handles[1], &payment); + waiters.push_back(handles[2], &other); + + // rule update arrives while the group EXISTS: weight 10 + let pattern: restate_limiter::RulePattern = "payment".parse().unwrap(); + apply_rule_updates_to_scope_weights( + &weights, + &[RuleUpdate::Upsert { + pattern, + limit: UserLimits::new(NonZeroU32::new(100)) + .with_scheduling_weight(NonZeroU32::new(10)), + }], + ); + + // existing group still behaves as weight 1: yields after ONE grant + assert_eq!(waiters.pop_front(), Some(handles[0])); + assert_eq!( + waiters.pop_front(), + Some(handles[2]), + "existing payment group must keep its creation-time weight (1) until recreated" + ); + assert_eq!(waiters.pop_front(), Some(handles[1])); + assert!(waiters.is_empty()); + + // fresh group (created after the update) picks up weight 10 immediately + waiters.push_back(handles[3], &payment); + waiters.push_back(handles[4], &payment); + waiters.push_back(handles[5], &other); + assert_eq!(waiters.pop_front(), Some(handles[3])); + assert_eq!( + waiters.pop_front(), + Some(handles[4]), + "fresh payment group runs at the new weight (10): both grants before yielding" + ); + assert_eq!(waiters.pop_front(), Some(handles[5])); + } +} diff --git a/crates/vqueues/src/scheduler/drr.rs b/crates/vqueues/src/scheduler/drr.rs index 473e3e1d2a..fcb875dad3 100644 --- a/crates/vqueues/src/scheduler/drr.rs +++ b/crates/vqueues/src/scheduler/drr.rs @@ -8,7 +8,7 @@ // the Business Source License, use of this software will be governed // by the Apache License, Version 2.0. -use std::num::NonZeroU16; +use std::num::NonZeroU32; use std::pin::Pin; use std::task::Poll; use std::task::Waker; @@ -36,7 +36,7 @@ use crate::cache; use crate::cache::VQueueHandle; use crate::metric_definitions::VQUEUE_ENQUEUE; use crate::metric_definitions::VQUEUE_RUN_CONFIRMED; -use crate::scheduler::eligible::EligibilityTracker; +use crate::scheduler::eligible::{EligibilityTracker, SchedulingGroup, WeightResolver}; use crate::scheduler::vqueue_state::Pop; use super::Decisions; @@ -57,9 +57,9 @@ pub struct DRRScheduler { last_report: Instant, /// Limits the number of queues picked per scheduler's poll - limit_qid_per_poll: NonZeroU16, + limit_qid_per_poll: NonZeroU32, /// Limits the number of items included in a single decision across all queues - max_items_per_decision: NonZeroU16, + max_items_per_decision: NonZeroU32, // SAFETY NOTE: **must** Keep this at the end since it needs to outlive all readers. storage: S, @@ -67,25 +67,26 @@ pub struct DRRScheduler { impl DRRScheduler { pub fn new( - limit_qid_per_poll: NonZeroU16, - max_items_per_decision: NonZeroU16, + limit_qid_per_poll: NonZeroU32, + max_items_per_decision: NonZeroU32, resource_manager: ResourceManager, storage: S, vqueues: VQueuesMeta<'_>, + weight_resolver: WeightResolver, ) -> Self { let mut total_running = 0; let mut total_waiting = 0; let mut q = SecondaryMap::with_capacity(vqueues.capacity()); let mut eligible: EligibilityTracker = - EligibilityTracker::with_capacity(vqueues.capacity()); + EligibilityTracker::with_capacity(vqueues.capacity(), weight_resolver); for (handle, qid, meta) in vqueues.iter_active_vqueues() { total_running += meta.num_running(); total_waiting += meta.total_waiting(); q.insert(handle, VQueueState::new(qid, &storage, meta.num_running())); // We init all active vqueues as eligible first - eligible.insert_eligible(handle); + eligible.insert_eligible(handle, SchedulingGroup::of(meta)); } debug!( @@ -409,6 +410,12 @@ impl DRRScheduler { } } + /// Verifies the eligibility tracker's internal invariants. Test-only. + #[cfg(test)] + pub fn assert_scheduler_invariants(&self) { + self.eligible.assert_invariants(); + } + /// Snapshot of every user-limit counter currently tracked by this scheduler's /// resource manager. Stamped with the owning partition's key. pub fn scan_user_limit_counters( @@ -422,7 +429,7 @@ impl DRRScheduler { #[cfg(test)] mod tests { - use std::num::{NonZeroU16, NonZeroU32, NonZeroUsize}; + use std::num::{NonZeroU32, NonZeroUsize}; use std::task::Poll; use restate_clock::RoughTimestamp; @@ -455,6 +462,11 @@ mod tests { const TEST_VQUEUES_CAPACITY: usize = 1024; + /// All services get the default weight of 1 unless a test overrides the resolver. + fn test_weight_resolver() -> WeightResolver { + std::sync::Arc::new(|_group: &SchedulingGroup| NonZeroU32::MIN) + } + use super::*; const BASE_RUN_AT_MS: u64 = 1_744_000_000_000; @@ -493,6 +505,47 @@ mod tests { enqueue_entry_with_run_at(txn, cache, qid, id, run_at, action_collector).await } + /// Like [`enqueue_entry`] but linking the vqueue to a caller-chosen service, so + /// tests can exercise the per-service group ring. + async fn enqueue_entry_for_service( + txn: &mut PartitionStoreTransaction<'_>, + cache: &mut VQueuesMetaCache, + qid: &VQueueId, + service: &ServiceName, + id: u8, + action_collector: Option<&mut Vec>, + ) -> EntryKey { + let run_at = MillisSinceEpoch::new(BASE_RUN_AT_MS); + let created_at = UniqueTimestamp::try_from(1000u64 + id as u64).unwrap(); + let seq = id as u64; + let entry_id = EntryId::new(EntryKind::Invocation, [id; EntryId::REMAINDER_LEN]); + let run_at_rough = RoughTimestamp::from_unix_millis_clamped(run_at); + + let mut vqueue = VQueue::get_or_create_vqueue( + created_at, + qid, + txn, + cache, + action_collector, + service, + &None, + &LimitKey::None, + &None, + ) + .await + .expect("vqueue should be created"); + + vqueue.enqueue_new( + created_at, + seq, + Some(run_at), + entry_id, + EntryMetadata::default(), + ); + + EntryKey::new(false, run_at_rough, seq, entry_id) + } + async fn enqueue_entry_with_run_at( txn: &mut PartitionStoreTransaction<'_>, cache: &mut VQueuesMetaCache, @@ -644,6 +697,16 @@ mod tests { db: &PartitionDb, concurrency: Concurrency, global_throttling: Option, + ) -> ResourceManager { + create_resource_manager_full(db, concurrency, global_throttling, test_weight_resolver()) + .await + } + + async fn create_resource_manager_full( + db: &PartitionDb, + concurrency: Concurrency, + global_throttling: Option, + weight_resolver: WeightResolver, ) -> ResourceManager { ResourceManager::create( db.clone(), @@ -651,6 +714,7 @@ mod tests { global_throttling, MemoryPool::unlimited(), NonZeroByteCount::new(NonZeroUsize::MIN), + weight_resolver, ) .await .expect("resource manager creation should succeed") @@ -668,11 +732,12 @@ mod tests { cache: &VQueuesMetaCache, ) -> DRRScheduler { DRRScheduler::new( - NonZeroU16::new(100).unwrap(), - NonZeroU16::new(100).unwrap(), + NonZeroU32::new(100).unwrap(), + NonZeroU32::new(100).unwrap(), create_resource_manager(db, Concurrency::new_unlimited()).await, db.clone(), cache.view(), + test_weight_resolver(), ) } @@ -682,8 +747,8 @@ mod tests { concurrency_limit: usize, ) -> DRRScheduler { DRRScheduler::new( - NonZeroU16::new(100).unwrap(), - NonZeroU16::new(100).unwrap(), + NonZeroU32::new(100).unwrap(), + NonZeroU32::new(100).unwrap(), create_resource_manager( db, Concurrency::new(Some(NonZeroUsize::new(concurrency_limit).unwrap())), @@ -691,6 +756,7 @@ mod tests { .await, db.clone(), cache.view(), + test_weight_resolver(), ) } @@ -913,6 +979,7 @@ mod tests { assert!(events.is_empty()); txn.commit().await.expect("commit should succeed"); + drop(txn); } #[restate_core::test] @@ -940,11 +1007,12 @@ mod tests { let db = rocksdb.partition_db(); let mut scheduler = DRRScheduler::new( - NonZeroU16::new(100).unwrap(), - NonZeroU16::new(3).unwrap(), + NonZeroU32::new(100).unwrap(), + NonZeroU32::new(3).unwrap(), create_resource_manager(db, Concurrency::new_unlimited()).await, db.clone(), cache.view(), + test_weight_resolver(), ); let Poll::Ready(Ok(decision)) = poll_scheduler(Pin::new(&mut scheduler), cache.view()) @@ -1075,8 +1143,8 @@ mod tests { let db = rocksdb.partition_db(); let mut scheduler = DRRScheduler::new( - NonZeroU16::new(100).unwrap(), - NonZeroU16::new(100).unwrap(), + NonZeroU32::new(100).unwrap(), + NonZeroU32::new(100).unwrap(), create_resource_manager_with_throttling( db, Concurrency::new_unlimited(), @@ -1085,6 +1153,7 @@ mod tests { .await, db.clone(), cache.view(), + test_weight_resolver(), ); let Poll::Ready(Ok(decision)) = poll_scheduler(Pin::new(&mut scheduler), cache.view()) @@ -1156,8 +1225,8 @@ mod tests { let db = rocksdb.partition_db(); let mut scheduler = DRRScheduler::new( - NonZeroU16::new(100).unwrap(), - NonZeroU16::new(100).unwrap(), + NonZeroU32::new(100).unwrap(), + NonZeroU32::new(100).unwrap(), create_resource_manager_with_throttling( db, Concurrency::new_unlimited(), @@ -1166,6 +1235,7 @@ mod tests { .await, db.clone(), cache.view(), + test_weight_resolver(), ); let Poll::Ready(Ok(_decision)) = poll_scheduler(Pin::new(&mut scheduler), cache.view()) @@ -1209,11 +1279,12 @@ mod tests { let db = rocksdb.partition_db(); let mut scheduler = DRRScheduler::new( - NonZeroU16::new(100).unwrap(), - NonZeroU16::new(2).unwrap(), + NonZeroU32::new(100).unwrap(), + NonZeroU32::new(2).unwrap(), create_resource_manager(db, Concurrency::new_unlimited()).await, db.clone(), cache.view(), + test_weight_resolver(), ); if let Poll::Ready(Ok(decision)) = poll_scheduler(Pin::new(&mut scheduler), cache.view()) { @@ -1238,11 +1309,12 @@ mod tests { let db = rocksdb.partition_db(); let mut scheduler = DRRScheduler::new( - NonZeroU16::new(100).unwrap(), - NonZeroU16::new(2).unwrap(), + NonZeroU32::new(100).unwrap(), + NonZeroU32::new(2).unwrap(), create_resource_manager(db, Concurrency::new_unlimited()).await, db.clone(), cache.view(), + test_weight_resolver(), ); let Poll::Ready(Ok(decision)) = poll_scheduler(Pin::new(&mut scheduler), cache.view()) @@ -1313,12 +1385,13 @@ mod tests { let db = rocksdb.partition_db(); let mut scheduler = DRRScheduler::new( - NonZeroU16::new(100).unwrap(), - NonZeroU16::new(10).unwrap(), + NonZeroU32::new(100).unwrap(), + NonZeroU32::new(10).unwrap(), create_resource_manager(db, Concurrency::new(Some(NonZeroUsize::new(1).unwrap()))) .await, db.clone(), cache.view(), + test_weight_resolver(), ); let h_qid1 = cache.view().handle_for(&qid1).unwrap(); @@ -1392,4 +1465,984 @@ mod tests { Poll::Pending )); } + + // --- Two-level weighted round-robin tests ------------------------------------- + // + // The tests below cover the service-group WRR layered on top of the per-queue DRR: + // backwards compatibility at weight=1, weighted dispatch ratios, work conservation, + // group lifecycle, and internal data-structure invariants. + + /// Weight resolver used by WRR tests: derives the weight from the service name + /// suffix, e.g. "svc-w10" gets weight 10. Anything else gets weight 1. + fn suffix_weight_resolver() -> WeightResolver { + std::sync::Arc::new(|group: &SchedulingGroup| { + let SchedulingGroup::Service(service_name) = group else { + return NonZeroU32::MIN; + }; + let name: &str = service_name.as_ref(); + name.rsplit_once("-w") + .and_then(|(_, w)| w.parse::().ok()) + .and_then(NonZeroU32::new) + .unwrap_or(NonZeroU32::MIN) + }) + } + + /// Builds `queues_per_service` single-entry vqueues for each given service. Queue + /// partition keys are `base + service_index * 1000 + queue_index` so a dispatched + /// qid can be mapped back to its service. + async fn setup_services( + rocksdb: &mut PartitionStore, + cache: &mut VQueuesMetaCache, + base: u64, + services: &[&str], + queues_per_service: u64, + ) { + let mut txn = rocksdb.transaction(); + for (si, service) in services.iter().enumerate() { + let service_name = ServiceName::new(service); + for qi in 0..queues_per_service { + let pk = base + si as u64 * 1000 + qi; + let qid = test_qid(pk); + enqueue_entry_for_service( + &mut txn, + cache, + &qid, + &service_name, + // entry ids only need to be unique within a queue + ((si * 40 + qi as usize) % 255) as u8 + 1, + None, + ) + .await; + } + } + txn.commit().await.expect("commit should succeed"); + drop(txn); + } + + /// Polls the scheduler with max_items_per_decision=1 until it goes Pending, and + /// counts dispatched items per service index (derived from the partition key). + fn drain_and_count( + scheduler: &mut DRRScheduler, + cache: &VQueuesMetaCache, + base: u64, + num_services: usize, + max_polls: usize, + ) -> Vec { + let mut counts = vec![0usize; num_services]; + for _ in 0..max_polls { + match poll_scheduler(Pin::new(scheduler), cache.view()) { + Poll::Ready(Ok(decision)) => { + for qid in decision.qids.keys() { + let si = ((qid.partition_key() - base) / 1000) as usize; + counts[si] += 1; + } + scheduler.assert_scheduler_invariants(); + } + Poll::Ready(Err(err)) => panic!("scheduler error: {err}"), + Poll::Pending => break, + } + } + counts + } + + fn scheduler_with_resolver( + db: &PartitionDb, + cache: &VQueuesMetaCache, + resource_manager: ResourceManager, + resolver: WeightResolver, + ) -> DRRScheduler { + DRRScheduler::new( + NonZeroU32::new(100).unwrap(), + NonZeroU32::new(1).unwrap(), // one item per decision to observe dispatch order + resource_manager, + db.clone(), + cache.view(), + resolver, + ) + } + + /// Backwards compatibility: with every service at weight 1 the groups share slots + /// equally, and every enqueued item is dispatched (nothing lost, nothing starved). + #[restate_core::test] + async fn wrr_all_weight_one_equal_service_share() { + let mut rocksdb = storage_test_environment().await; + let mut cache = VQueuesMetaCache::new_empty(TEST_VQUEUES_CAPACITY); + const BASE: u64 = 30_000; + + setup_services( + &mut rocksdb, + &mut cache, + BASE, + &["svc-a", "svc-b", "svc-c"], + 5, + ) + .await; + + let db = rocksdb.partition_db(); + let mut scheduler = scheduler_with_resolver( + db, + &cache, + create_resource_manager(db, Concurrency::new_unlimited()).await, + test_weight_resolver(), + ); + + let counts = drain_and_count(&mut scheduler, &cache, BASE, 3, 100); + assert_eq!(counts, vec![5, 5, 5], "each weight-1 service drains fully"); + } + + /// A single service must drain exactly like the previous flat ring: all items + /// dispatched, scheduler Pending afterwards, group dropped when empty. + #[restate_core::test] + async fn wrr_single_service_drains_completely() { + let mut rocksdb = storage_test_environment().await; + let mut cache = VQueuesMetaCache::new_empty(TEST_VQUEUES_CAPACITY); + const BASE: u64 = 32_000; + + setup_services(&mut rocksdb, &mut cache, BASE, &["svc-solo"], 10).await; + + let db = rocksdb.partition_db(); + let mut scheduler = scheduler_with_resolver( + db, + &cache, + create_resource_manager(db, Concurrency::new_unlimited()).await, + test_weight_resolver(), + ); + + let counts = drain_and_count(&mut scheduler, &cache, BASE, 1, 100); + assert_eq!(counts, vec![10]); + assert!(matches!( + poll_scheduler(Pin::new(&mut scheduler), cache.view()), + Poll::Pending + )); + scheduler.assert_scheduler_invariants(); + } + + /// Weighted ratio: weight-10 service gets ~10 slots per weight-1 slot while both + /// have work queued. Observed over the first dispatches before either drains. + #[restate_core::test] + async fn wrr_two_groups_dispatch_ratio_one_to_ten() { + let mut rocksdb = storage_test_environment().await; + let mut cache = VQueuesMetaCache::new_empty(TEST_VQUEUES_CAPACITY); + const BASE: u64 = 34_000; + + // 40 queues each so neither group drains within the observation window + setup_services( + &mut rocksdb, + &mut cache, + BASE, + &["payment-w1", "indexer-w10"], + 40, + ) + .await; + + let db = rocksdb.partition_db(); + let mut scheduler = scheduler_with_resolver( + db, + &cache, + create_resource_manager(db, Concurrency::new_unlimited()).await, + suffix_weight_resolver(), + ); + + // observe 22 dispatch slots: exactly two full WRR cycles (1 + 10 per cycle) + let counts = drain_and_count(&mut scheduler, &cache, BASE, 2, 22); + assert_eq!( + counts[0] + counts[1], + 22, + "both groups still had work; every poll must dispatch" + ); + assert_eq!(counts[0], 2, "weight-1 group gets 1 slot per cycle"); + assert_eq!(counts[1], 20, "weight-10 group gets 10 slots per cycle"); + } + + /// Work conservation: when the high-weight group drains, the low-weight group gets + /// every remaining slot and drains completely (no starvation, no lost items). + #[restate_core::test] + async fn wrr_work_conservation_after_group_drains() { + let mut rocksdb = storage_test_environment().await; + let mut cache = VQueuesMetaCache::new_empty(TEST_VQUEUES_CAPACITY); + const BASE: u64 = 36_000; + + let mut txn = rocksdb.transaction(); + let heavy = ServiceName::new("heavy-w100"); + let light = ServiceName::new("light-w1"); + for qi in 0..2u64 { + let qid = test_qid(BASE + qi); + enqueue_entry_for_service(&mut txn, &mut cache, &qid, &heavy, qi as u8 + 1, None).await; + } + for qi in 0..20u64 { + let qid = test_qid(BASE + 1000 + qi); + enqueue_entry_for_service(&mut txn, &mut cache, &qid, &light, qi as u8 + 10, None) + .await; + } + txn.commit().await.expect("commit should succeed"); + drop(txn); + + let db = rocksdb.partition_db(); + let mut scheduler = scheduler_with_resolver( + db, + &cache, + create_resource_manager(db, Concurrency::new_unlimited()).await, + suffix_weight_resolver(), + ); + + let counts = drain_and_count(&mut scheduler, &cache, BASE, 2, 100); + assert_eq!( + counts, + vec![2, 20], + "heavy group drains its 2 items, light group still drains all 20" + ); + assert!(matches!( + poll_scheduler(Pin::new(&mut scheduler), cache.view()), + Poll::Pending + )); + } + + /// Group lifecycle: a drained group is removed from the ring and correctly + /// re-created when new work for its service arrives. + #[restate_core::test] + async fn wrr_group_removed_and_recreated_on_new_enqueue() { + let mut rocksdb = storage_test_environment().await; + let mut cache = VQueuesMetaCache::new_empty(TEST_VQUEUES_CAPACITY); + const BASE: u64 = 38_000; + let service = ServiceName::new("cycler"); + let qid = test_qid(BASE); + + let mut txn = rocksdb.transaction(); + enqueue_entry_for_service(&mut txn, &mut cache, &qid, &service, 1, None).await; + txn.commit().await.expect("commit should succeed"); + drop(txn); + + let db = rocksdb.partition_db(); + let mut scheduler = scheduler_with_resolver( + db, + &cache, + create_resource_manager(db, Concurrency::new_unlimited()).await, + test_weight_resolver(), + ); + + let counts = drain_and_count(&mut scheduler, &cache, BASE, 1, 10); + assert_eq!(counts, vec![1]); + scheduler.assert_scheduler_invariants(); + + // new work for the same service re-creates the group + let mut txn = rocksdb.transaction(); + let mut events = Vec::new(); + enqueue_entry_for_service(&mut txn, &mut cache, &qid, &service, 2, Some(&mut events)).await; + txn.commit().await.expect("commit should succeed"); + drop(txn); + for event in events.drain(..) { + scheduler.on_inbox_event(cache.view(), event); + } + + let counts = drain_and_count(&mut scheduler, &cache, BASE, 1, 10); + assert_eq!(counts, vec![1], "re-created group dispatches new work"); + scheduler.assert_scheduler_invariants(); + } + + /// Blocked queues leave the ring without stalling their group: with a concurrency + /// limit of 1, other queues of the same service keep dispatching one at a time. + #[restate_core::test] + async fn wrr_blocked_queue_does_not_stall_group() { + let mut rocksdb = storage_test_environment().await; + let mut cache = VQueuesMetaCache::new_empty(TEST_VQUEUES_CAPACITY); + const BASE: u64 = 40_000; + + setup_services(&mut rocksdb, &mut cache, BASE, &["limited"], 3).await; + + let db = rocksdb.partition_db(); + let mut scheduler = scheduler_with_resolver( + db, + &cache, + create_resource_manager(db, Concurrency::new(Some(NonZeroUsize::new(1).unwrap()))) + .await, + test_weight_resolver(), + ); + + // Only one Run fits within the concurrency limit; the other two queues get + // blocked on InvokerConcurrency and must leave the ring without panicking. + let Poll::Ready(Ok(decision)) = poll_scheduler(Pin::new(&mut scheduler), cache.view()) + else { + panic!("expected a decision"); + }; + assert_eq!(decision.num_run(), 1); + scheduler.assert_scheduler_invariants(); + + assert!(matches!( + poll_scheduler(Pin::new(&mut scheduler), cache.view()), + Poll::Pending + )); + scheduler.assert_scheduler_invariants(); + + // Confirming the running item frees the concurrency slot; a blocked queue is + // woken up and dispatches next. + let first_qid = decision.qids.keys().next().unwrap().clone(); + let first_key = run_keys(&decision)[0]; + let metas = cache.view(); + let h_first = metas.handle_for(&first_qid).unwrap(); + let first_slot = metas.get(h_first).unwrap(); + let mut scheduler = Pin::new(&mut scheduler); + let resources = scheduler + .as_mut() + .confirm_run_attempt(h_first, first_slot, &first_key); + assert!(resources.is_some()); + drop(resources); + + let Poll::Ready(Ok(decision2)) = poll_scheduler(scheduler.as_mut(), cache.view()) else { + panic!("expected a second decision"); + }; + assert_eq!(decision2.num_run(), 1); + scheduler.assert_scheduler_invariants(); + } + + /// Weight changes take effect when a group cycles: the resolver is re-consulted on + /// group re-creation, mirroring an admin PATCH between benchmark runs. + #[restate_core::test] + async fn wrr_weight_change_applies_on_group_recreation() { + use std::sync::Arc; + use std::sync::atomic::{AtomicU32, Ordering}; + + let mut rocksdb = storage_test_environment().await; + let mut cache = VQueuesMetaCache::new_empty(TEST_VQUEUES_CAPACITY); + const BASE: u64 = 42_000; + let service = ServiceName::new("mutable"); + let qid = test_qid(BASE); + + let mut txn = rocksdb.transaction(); + enqueue_entry_for_service(&mut txn, &mut cache, &qid, &service, 1, None).await; + txn.commit().await.expect("commit should succeed"); + drop(txn); + + let weight = Arc::new(AtomicU32::new(1)); + let resolver_weight = Arc::clone(&weight); + let resolver: WeightResolver = Arc::new(move |_group: &SchedulingGroup| { + NonZeroU32::new(resolver_weight.load(Ordering::Relaxed)).unwrap() + }); + + let db = rocksdb.partition_db(); + let mut scheduler = scheduler_with_resolver( + db, + &cache, + create_resource_manager(db, Concurrency::new_unlimited()).await, + resolver, + ); + + // drain the group so it gets dropped, then bump the weight + let counts = drain_and_count(&mut scheduler, &cache, BASE, 1, 10); + assert_eq!(counts, vec![1]); + weight.store(10, Ordering::Relaxed); + + // new work re-creates the group; the new weight must be picked up (verified + // indirectly: the dispatch succeeds and invariants hold — the ratio itself is + // covered by wrr_two_groups_dispatch_ratio_one_to_ten) + let mut txn = rocksdb.transaction(); + let mut events = Vec::new(); + enqueue_entry_for_service(&mut txn, &mut cache, &qid, &service, 2, Some(&mut events)).await; + txn.commit().await.expect("commit should succeed"); + drop(txn); + for event in events.drain(..) { + scheduler.on_inbox_event(cache.view(), event); + } + let counts = drain_and_count(&mut scheduler, &cache, BASE, 1, 10); + assert_eq!(counts, vec![1]); + scheduler.assert_scheduler_invariants(); + } + + /// Bulk stress on the invariants: many services, many queues, drain everything. + #[restate_core::test] + async fn wrr_invariants_under_bulk_load() { + let mut rocksdb = storage_test_environment().await; + let mut cache = VQueuesMetaCache::new_empty(TEST_VQUEUES_CAPACITY); + const BASE: u64 = 44_000; + + setup_services( + &mut rocksdb, + &mut cache, + BASE, + &["s0-w1", "s1-w2", "s2-w4", "s3-w1", "s4-w8"], + 8, + ) + .await; + + let db = rocksdb.partition_db(); + let mut scheduler = scheduler_with_resolver( + db, + &cache, + create_resource_manager(db, Concurrency::new_unlimited()).await, + suffix_weight_resolver(), + ); + + // drain_and_count asserts invariants after every poll + let counts = drain_and_count(&mut scheduler, &cache, BASE, 5, 200); + assert_eq!(counts, vec![8, 8, 8, 8, 8], "all queues fully drained"); + assert!(matches!( + poll_scheduler(Pin::new(&mut scheduler), cache.view()), + Poll::Pending + )); + } + + /// Three groups with weights 1/2/4: within one full WRR cycle (7 slots) the groups + /// receive exactly 1, 2 and 4 dispatches. + #[restate_core::test] + async fn wrr_three_groups_weighted_ratio() { + let mut rocksdb = storage_test_environment().await; + let mut cache = VQueuesMetaCache::new_empty(TEST_VQUEUES_CAPACITY); + const BASE: u64 = 46_000; + + setup_services( + &mut rocksdb, + &mut cache, + BASE, + &["a-w1", "b-w2", "c-w4"], + 20, + ) + .await; + + let db = rocksdb.partition_db(); + let mut scheduler = scheduler_with_resolver( + db, + &cache, + create_resource_manager(db, Concurrency::new_unlimited()).await, + suffix_weight_resolver(), + ); + + // two full cycles: 2 * (1 + 2 + 4) = 14 slots, no group drains (20 queues each) + let counts = drain_and_count(&mut scheduler, &cache, BASE, 3, 14); + assert_eq!(counts.iter().sum::(), 14); + assert_eq!( + counts, + vec![2, 4, 8], + "1:2:4 weighted split over two cycles" + ); + } + + /// The thesis starvation scenario in miniature: a huge weight-1 group must not + /// starve a tiny weight-10 group. The small group's share depends only on the + /// weights, not on the queue-count imbalance. + #[restate_core::test] + async fn wrr_small_high_weight_group_not_starved_by_large_group() { + let mut rocksdb = storage_test_environment().await; + let mut cache = VQueuesMetaCache::new_empty(TEST_VQUEUES_CAPACITY); + const BASE: u64 = 48_000; + + let mut txn = rocksdb.transaction(); + let payment = ServiceName::new("payment-w1"); + let batcher = ServiceName::new("batcher-w10"); + // 100 payment queues vs 2 batcher queues with plenty of work: enqueue two + // entries per batcher queue so the batcher group survives multiple cycles + for qi in 0..100u64 { + let qid = test_qid(BASE + qi); + enqueue_entry_for_service(&mut txn, &mut cache, &qid, &payment, qi as u8, None).await; + } + for qi in 0..2u64 { + let qid = test_qid(BASE + 1000 + qi); + enqueue_entry_for_service(&mut txn, &mut cache, &qid, &batcher, 200 + qi as u8, None) + .await; + } + txn.commit().await.expect("commit should succeed"); + drop(txn); + + let db = rocksdb.partition_db(); + let mut scheduler = scheduler_with_resolver( + db, + &cache, + create_resource_manager(db, Concurrency::new_unlimited()).await, + suffix_weight_resolver(), + ); + + // Observe the first 4 slots: the flat ring would have given the batcher a + // ~2/102 chance per slot; the WRR must dispatch both batcher queues within the + // first cycle (1 payment slot + up to 10 batcher slots). + let counts = drain_and_count(&mut scheduler, &cache, BASE, 2, 4); + assert_eq!( + counts[1], 2, + "both batcher queues dispatched within the first cycle despite 100 payment queues" + ); + } + + /// max_items_per_decision caps the whole decision across groups, not per group. + #[restate_core::test] + async fn wrr_max_items_per_decision_across_groups() { + let mut rocksdb = storage_test_environment().await; + let mut cache = VQueuesMetaCache::new_empty(TEST_VQUEUES_CAPACITY); + const BASE: u64 = 50_000; + + setup_services( + &mut rocksdb, + &mut cache, + BASE, + &["x-w1", "y-w1", "z-w1"], + 10, + ) + .await; + + let db = rocksdb.partition_db(); + let mut scheduler = DRRScheduler::new( + NonZeroU32::new(100).unwrap(), + NonZeroU32::new(5).unwrap(), + create_resource_manager(db, Concurrency::new_unlimited()).await, + db.clone(), + cache.view(), + suffix_weight_resolver(), + ); + + let Poll::Ready(Ok(decision)) = poll_scheduler(Pin::new(&mut scheduler), cache.view()) + else { + panic!("expected decision"); + }; + assert!( + decision.total_items() <= 5, + "decision item cap applies across all groups, got {}", + decision.total_items() + ); + scheduler.assert_scheduler_invariants(); + } + + /// Delayed items (run_at in the future) leave the ring into the delay queue and the + /// overdue item preempts within the same queue — original behavior, now per group. + #[restate_core::test] + async fn wrr_run_at_scheduling_unchanged() { + let mut rocksdb = storage_test_environment().await; + let mut cache = VQueuesMetaCache::new_empty(TEST_VQUEUES_CAPACITY); + let qid = test_qid(52_000); + let now = MillisSinceEpoch::now().as_u64(); + + let mut txn = rocksdb.transaction(); + let future_key = enqueue_entry_with_run_at( + &mut txn, + &mut cache, + &qid, + 1, + MillisSinceEpoch::new(now.saturating_add(60_000)), + None, + ) + .await; + let overdue_key = enqueue_entry_with_run_at( + &mut txn, + &mut cache, + &qid, + 2, + MillisSinceEpoch::new(now.saturating_sub(1_000)), + None, + ) + .await; + txn.commit().await.expect("commit should succeed"); + drop(txn); + + let db = rocksdb.partition_db(); + let mut scheduler = create_scheduler(db, &cache).await; + + let Poll::Ready(Ok(decision)) = poll_scheduler(Pin::new(&mut scheduler), cache.view()) + else { + panic!("expected decision"); + }; + let keys = run_keys(&decision); + assert_eq!(keys.len(), 1); + assert_eq!(keys[0], overdue_key); + assert_ne!(keys[0], future_key); + scheduler.assert_scheduler_invariants(); + } + + /// External removal (RemovedFromInbox) of a queue's only item makes the queue + /// dormant; its group must not retain a stale handle and the sibling queue in the + /// same group keeps dispatching. + #[restate_core::test] + async fn wrr_external_removal_purges_handle_from_group() { + let mut rocksdb = storage_test_environment().await; + let mut cache = VQueuesMetaCache::new_empty(TEST_VQUEUES_CAPACITY); + const BASE: u64 = 54_000; + let service = ServiceName::new("removal"); + let qid_kept = test_qid(BASE); + let qid_removed = test_qid(BASE + 1); + + let mut txn = rocksdb.transaction(); + enqueue_entry_for_service(&mut txn, &mut cache, &qid_kept, &service, 1, None).await; + let removed_key = + enqueue_entry_for_service(&mut txn, &mut cache, &qid_removed, &service, 2, None).await; + txn.commit().await.expect("commit should succeed"); + drop(txn); + + let db = rocksdb.partition_db(); + let mut scheduler = scheduler_with_resolver( + db, + &cache, + create_resource_manager(db, Concurrency::new_unlimited()).await, + test_weight_resolver(), + ); + + // Externally remove the only item of the second queue. + let mut txn = rocksdb.transaction(); + let mut events = Vec::new(); + let at = UniqueTimestamp::try_from(2_000u64).unwrap(); + let header = txn + .get_vqueue_entry_status(qid_removed.partition_key(), removed_key.entry_id()) + .await + .expect("lookup succeeds") + .expect("entry exists"); + let mut vqueue = VQueue::get_or_create_vqueue( + at, + &qid_removed, + &mut txn, + &mut cache, + Some(&mut events), + &service, + &None, + &LimitKey::None, + &None, + ) + .await + .expect("vqueue exists"); + vqueue.end( + at, + &header, + restate_storage_api::vqueue_table::Status::Killed, + Duration::ZERO, + ); + txn.commit().await.expect("commit should succeed"); + drop(txn); + for event in events.drain(..) { + scheduler.on_inbox_event(cache.view(), event); + } + scheduler.assert_scheduler_invariants(); + + // Only the kept queue's item is dispatched; the removed one never appears. + let counts = drain_and_count(&mut scheduler, &cache, BASE, 1, 10); + assert_eq!(counts, vec![1], "only the kept queue dispatches"); + assert!(matches!( + poll_scheduler(Pin::new(&mut scheduler), cache.view()), + Poll::Pending + )); + scheduler.assert_scheduler_invariants(); + } + + /// Vqueues without a service link land in the shared unlinked group and are + /// scheduled normally alongside service-linked groups. + #[restate_core::test] + async fn wrr_unlinked_and_linked_queues_coexist() { + let mut rocksdb = storage_test_environment().await; + let mut cache = VQueuesMetaCache::new_empty(TEST_VQUEUES_CAPACITY); + const BASE: u64 = 56_000; + + let mut txn = rocksdb.transaction(); + // enqueue_entry uses the fixed service "test" — plus one differently-named + // service, exercising two coexisting groups with default weight + enqueue_entry(&mut txn, &mut cache, &test_qid(BASE), 1, 0, None).await; + let other = ServiceName::new("other"); + enqueue_entry_for_service( + &mut txn, + &mut cache, + &test_qid(BASE + 1000), + &other, + 2, + None, + ) + .await; + txn.commit().await.expect("commit should succeed"); + drop(txn); + + let db = rocksdb.partition_db(); + let mut scheduler = scheduler_with_resolver( + db, + &cache, + create_resource_manager(db, Concurrency::new_unlimited()).await, + test_weight_resolver(), + ); + + let counts = drain_and_count(&mut scheduler, &cache, BASE, 2, 10); + assert_eq!(counts, vec![1, 1], "both groups fully served"); + scheduler.assert_scheduler_invariants(); + } + + /// Higher-priority (earlier run_at) items still preempt within a queue between + /// polls — regression guard for the original `higher_priority_preempts` behavior + /// under the group ring. + #[restate_core::test] + async fn wrr_priority_preemption_within_group_unchanged() { + let mut rocksdb = storage_test_environment().await; + let mut cache = VQueuesMetaCache::new_empty(TEST_VQUEUES_CAPACITY); + let qid = test_qid(58_000); + let service = ServiceName::new("preempt"); + let mut events = Vec::new(); + + let mut txn = rocksdb.transaction(); + for i in 1..=3 { + enqueue_entry_for_service(&mut txn, &mut cache, &qid, &service, i, None).await; + } + txn.commit().await.expect("commit should succeed"); + drop(txn); + + let db = rocksdb.partition_db(); + let mut scheduler = DRRScheduler::new( + NonZeroU32::new(100).unwrap(), + NonZeroU32::new(2).unwrap(), + create_resource_manager(db, Concurrency::new_unlimited()).await, + db.clone(), + cache.view(), + test_weight_resolver(), + ); + + let Poll::Ready(Ok(decision)) = poll_scheduler(Pin::new(&mut scheduler), cache.view()) + else { + panic!("expected decision"); + }; + assert_eq!(decision.num_run(), 2); + scheduler.assert_scheduler_invariants(); + + // enqueue a fourth item while the first two are in flight + let mut txn = rocksdb.transaction(); + events.clear(); + enqueue_entry_for_service(&mut txn, &mut cache, &qid, &service, 4, Some(&mut events)).await; + txn.commit().await.expect("commit should succeed"); + drop(txn); + for event in events.drain(..) { + scheduler.on_inbox_event(cache.view(), event); + } + + let Poll::Ready(Ok(decision)) = poll_scheduler(Pin::new(&mut scheduler), cache.view()) + else { + panic!("expected decision"); + }; + assert_eq!(decision.num_run(), 2, "remaining items keep flowing"); + scheduler.assert_scheduler_invariants(); + } + + /// End-to-end: under invoker-concurrency pressure (limit 1), the WRR waiter + /// list — not arrival order — decides who acquires freed permits. Payment + /// queues arrive first and outnumber the batcher 6:2; a FIFO waiter list + /// would grant all payment permits before any batcher permit. The batcher + /// must instead be dispatched within the first WRR cycle after a release. + #[restate_core::test] + async fn wrr_waiter_list_decides_dispatch_under_concurrency_pressure() { + let mut rocksdb = storage_test_environment().await; + let mut cache = VQueuesMetaCache::new_empty(TEST_VQUEUES_CAPACITY); + const BASE: u64 = 62_000; + let payment = ServiceName::new("payment-w1"); + let batcher = ServiceName::new("batcher-w10"); + + let mut txn = rocksdb.transaction(); + // payments enqueued FIRST — they flood the waiter list on the first poll + for qi in 0..6u64 { + let qid = test_qid(BASE + qi); + enqueue_entry_for_service(&mut txn, &mut cache, &qid, &payment, qi as u8 + 1, None) + .await; + } + for qi in 0..2u64 { + let qid = test_qid(BASE + 1000 + qi); + enqueue_entry_for_service(&mut txn, &mut cache, &qid, &batcher, 50 + qi as u8, None) + .await; + } + txn.commit().await.expect("commit should succeed"); + drop(txn); + + let db = rocksdb.partition_db(); + let mut scheduler = DRRScheduler::new( + NonZeroU32::new(100).unwrap(), + NonZeroU32::new(1).unwrap(), + create_resource_manager_full( + db, + Concurrency::new(Some(NonZeroUsize::new(1).unwrap())), + None, + suffix_weight_resolver(), + ) + .await, + db.clone(), + cache.view(), + suffix_weight_resolver(), + ); + + // Dispatch one item at a time; after each Run, confirm it to release + // the single permit, then record which service the next dispatch hits. + let mut order: Vec<&str> = Vec::new(); + for _ in 0..8 { + let Poll::Ready(Ok(decision)) = poll_scheduler(Pin::new(&mut scheduler), cache.view()) + else { + panic!("expected a decision, got Pending after {order:?}"); + }; + assert_eq!(decision.num_run(), 1); + let qid = decision.qids.keys().next().unwrap().clone(); + let service = if qid.partition_key() >= BASE + 1000 { + "batcher" + } else { + "payment" + }; + order.push(service); + scheduler.assert_scheduler_invariants(); + + // release the permit so the waiter list decides the next grant + let key = run_keys(&decision)[0]; + let metas = cache.view(); + let handle = metas.handle_for(&qid).unwrap(); + let slot = metas.get(handle).unwrap(); + let resources = Pin::new(&mut scheduler).confirm_run_attempt(handle, slot, &key); + assert!(resources.is_some()); + drop(resources); + } + + assert_eq!(order.len(), 8, "all queues eventually dispatched"); + assert_eq!( + order.iter().filter(|s| **s == "batcher").count(), + 2, + "both batcher queues dispatched" + ); + let first_batcher = order.iter().position(|s| *s == "batcher").unwrap(); + assert!( + first_batcher <= 2, + "batcher must be granted within the first WRR cycle despite arriving \ + last behind 6 payment queues; order: {order:?}" + ); + } + + /// Scope groups: scoped queues of DIFFERENT services schedule as ONE group + /// whose weight comes from the (rule-fed) scope-weights map, competing + /// against unscoped per-service groups at weight 1 — the end-to-end shape + /// of `restate rules set payment --weight 10` + scope inheritance. + #[restate_core::test] + async fn wrr_scope_group_spans_services_with_rule_weight() { + use crate::scheduler::{ScopeWeights, scope_weight_resolver}; + use restate_types::Scope; + + let mut rocksdb = storage_test_environment().await; + let mut cache = VQueuesMetaCache::new_empty(TEST_VQUEUES_CAPACITY); + const BASE: u64 = 64_000; + let payment_scope = Scope::try_non_interned("payment").unwrap(); + + let mut txn = rocksdb.transaction(); + // scoped: two different services under ONE scope, 20 queues total + for qi in 0..10u64 { + for (si, svc) in ["PaymentWorkflow", "SepaService"].iter().enumerate() { + let qid = test_qid(BASE + si as u64 * 100 + qi); + let service = ServiceName::new(svc); + let at = UniqueTimestamp::try_from(1000u64 + qi + si as u64 * 50).unwrap(); + let entry_id = EntryId::new( + EntryKind::Invocation, + [(si as u8) * 100 + qi as u8 + 1; EntryId::REMAINDER_LEN], + ); + let run_at = MillisSinceEpoch::new(BASE_RUN_AT_MS); + let mut vqueue = VQueue::get_or_create_vqueue( + at, + &qid, + &mut txn, + &mut cache, + None::<&mut Vec>, + &service, + &Some(payment_scope.clone()), + &LimitKey::None, + &None, + ) + .await + .expect("vqueue should be created"); + vqueue.enqueue_new( + at, + qi + si as u64 * 100, + Some(run_at), + entry_id, + EntryMetadata::default(), + ); + } + } + // unscoped competitor: one service, 20 queues + for qi in 0..20u64 { + let qid = test_qid(BASE + 1000 + qi); + enqueue_entry_for_service( + &mut txn, + &mut cache, + &qid, + &ServiceName::new("indexer"), + 200 + qi as u8, + None, + ) + .await; + } + txn.commit().await.expect("commit should succeed"); + drop(txn); + + // rule-fed weights: scope "payment" -> 10 + let weights = ScopeWeights::default(); + weights + .write() + .unwrap() + .insert("payment".to_owned(), NonZeroU32::new(10).unwrap()); + let resolver = scope_weight_resolver(weights); + + let db = rocksdb.partition_db(); + let mut scheduler = DRRScheduler::new( + NonZeroU32::new(100).unwrap(), + NonZeroU32::new(1).unwrap(), + create_resource_manager(db, Concurrency::new_unlimited()).await, + db.clone(), + cache.view(), + resolver, + ); + + // one full WRR cycle = 10 (scope) + 1 (indexer service group) slots + let mut scope_count = 0usize; + let mut service_count = 0usize; + for _ in 0..22 { + let Poll::Ready(Ok(decision)) = poll_scheduler(Pin::new(&mut scheduler), cache.view()) + else { + panic!("expected a decision"); + }; + for qid in decision.qids.keys() { + if qid.partition_key() >= BASE + 1000 { + service_count += 1; + } else { + scope_count += 1; + } + } + scheduler.assert_scheduler_invariants(); + } + assert_eq!( + (scope_count, service_count), + (20, 2), + "scope group (weight 10, spanning two services) gets 10 slots per cycle \ + vs 1 for the unscoped service group" + ); + } + + /// Rapid enqueue/drain cycles across many services: groups are created and dropped + /// repeatedly without desyncing ring and map. + #[restate_core::test] + async fn wrr_repeated_group_churn_invariants() { + let mut rocksdb = storage_test_environment().await; + let mut cache = VQueuesMetaCache::new_empty(TEST_VQUEUES_CAPACITY); + const BASE: u64 = 60_000; + let db = rocksdb.partition_db().clone(); + + let mut scheduler = scheduler_with_resolver( + &db, + &cache, + create_resource_manager(&db, Concurrency::new_unlimited()).await, + suffix_weight_resolver(), + ); + + for round in 0..5u64 { + let mut txn = rocksdb.transaction(); + let mut events = Vec::new(); + for si in 0..3u64 { + let service = ServiceName::new(&format!("churn{si}-w{}", si + 1)); + let qid = test_qid(BASE + si * 1000 + round); + enqueue_entry_for_service( + &mut txn, + &mut cache, + &qid, + &service, + (round * 10 + si) as u8 + 1, + Some(&mut events), + ) + .await; + } + txn.commit().await.expect("commit should succeed"); + drop(txn); + for event in events.drain(..) { + scheduler.on_inbox_event(cache.view(), event); + } + + let counts = drain_and_count(&mut scheduler, &cache, BASE, 3, 20); + assert_eq!(counts, vec![1, 1, 1], "round {round}: all services served"); + } + assert!(matches!( + poll_scheduler(Pin::new(&mut scheduler), cache.view()), + Poll::Pending + )); + scheduler.assert_scheduler_invariants(); + } } diff --git a/crates/vqueues/src/scheduler/eligible.rs b/crates/vqueues/src/scheduler/eligible.rs index 4c8f1194d6..c5e5eed00c 100644 --- a/crates/vqueues/src/scheduler/eligible.rs +++ b/crates/vqueues/src/scheduler/eligible.rs @@ -9,6 +9,7 @@ // by the Apache License, Version 2.0. use std::collections::VecDeque; +use std::num::NonZeroU32; use std::task::Poll; use std::time::Duration; @@ -17,10 +18,12 @@ use slotmap::secondary::Entry; use tokio_util::time::{DelayQueue, delay_queue}; use restate_limiter::RuleHandle; +use restate_platform::hash::HashMap; use restate_storage_api::StorageError; use restate_storage_api::vqueue_table::VQueueStore; use restate_storage_api::vqueue_table::metadata::VQueueMeta; use restate_types::time::MillisSinceEpoch; +use restate_types::{Scope, ServiceName}; use restate_util_string::ReString; use restate_worker_api::{ResourceKind, SchedulingStatus}; @@ -40,6 +43,65 @@ use crate::scheduler::vqueue_state::Eligibility; /// I applaud your server's uptime. const MAX_PARK: Duration = Duration::from_secs(365 * 24 * 60 * 60); +/// The scheduler's fairness key: a vqueue belongs to its invocation SCOPE when +/// one is attached (Restate's flow-control scopes — the operator-facing fairness +/// key, weighted via `restate rules set --weight N`), and falls back to +/// its service otherwise. Service groups always run at the default weight 1; +/// only scope groups have configurable weights. +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub enum SchedulingGroup { + Scope(Scope), + Service(ServiceName), +} + +impl SchedulingGroup { + /// Derives the scheduling group of a queue from its metadata. + pub fn of(meta: &VQueueMeta) -> Self { + match meta.scope() { + Some(scope) => Self::Scope(scope.clone()), + None => Self::Service(meta.service_name().cloned().unwrap_or_else(unlinked_group)), + } + } +} + +/// Resolves a scheduling group's weight. Consulted whenever a group is +/// (re-)created, so rule changes take effect as soon as a group cycles. +/// Shared (Arc) between the eligibility tracker and the resource manager's +/// grouped waiter lists. +pub type WeightResolver = std::sync::Arc NonZeroU32 + Send + Sync>; + +/// Group name used for vqueues that are not linked to any service. +pub(in crate::scheduler) fn unlinked_group() -> ServiceName { + ServiceName::new("") +} + +/// Weighted-round-robin quota shared by the eligibility ring's [`ServiceGroup`] +/// and the resource manager's `GroupedWaiters`: a group receives `weight` slots +/// per cycle before yielding its turn. +#[derive(Debug)] +pub(in crate::scheduler) struct WrrQuota { + weight: NonZeroU32, + used: u32, +} + +impl WrrQuota { + pub(in crate::scheduler) fn new(weight: NonZeroU32) -> Self { + Self { weight, used: 0 } + } + + /// Consumes one slot; returns true when the quota is exhausted and the + /// group ring should rotate. + pub(in crate::scheduler) fn consume_slot(&mut self) -> bool { + self.used += 1; + if self.used >= self.weight.get() { + self.used = 0; + true + } else { + false + } + } +} + #[derive(Debug, Copy, Clone)] pub(super) struct WakeUp { ts: MillisSinceEpoch, @@ -61,27 +123,116 @@ pub(super) enum State { BlockedOn(ResourceKind), } +/// Per-service group in the two-level weighted round-robin. +/// +/// The scheduler cycles through groups (outer ring); within a group, vqueues cycle in +/// the inner ring exactly like the flat DRR. A group with weight N receives N +/// dispatch slots before rotating to the next group. +#[derive(Debug)] +struct ServiceGroup { + quota: WrrQuota, + ring: VecDeque, +} + +impl ServiceGroup { + fn new(weight: NonZeroU32, handle: VQueueHandle) -> Self { + let mut ring = VecDeque::with_capacity(4); + ring.push_back(handle); + Self { + quota: WrrQuota::new(weight), + ring, + } + } +} + #[derive(derive_more::Debug)] pub(crate) struct EligibilityTracker { #[debug(skip)] delayed_eligibility: DelayQueue, - ready_ring: VecDeque, + /// Outer WRR ring: rotation order over service groups. Invariant: contains exactly + /// the keys of `groups`, each once. + /// + /// Note: group keys hash their interned string CONTENT (not a pointer), + /// so `groups` lookups cost O(name length) — acceptable because lookups are + /// per dispatch step, not per queued key. + group_ring: VecDeque, + #[debug(skip)] + groups: HashMap, + /// Which group a handle belongs to. Set when the scheduler first learns about a + /// queue and kept for the lifetime of the cache slot (a queue's scope/service + /// never changes), so wake-up paths that only carry the handle can find the group. + #[debug(skip)] + handle_group: SecondaryMap, #[debug(skip)] states: SecondaryMap, + #[debug(skip)] + weight_resolver: WeightResolver, } impl EligibilityTracker { - pub fn with_capacity(capacity: usize) -> Self { + pub fn with_capacity(capacity: usize, weight_resolver: WeightResolver) -> Self { Self { delayed_eligibility: DelayQueue::with_capacity(capacity), - ready_ring: VecDeque::with_capacity(capacity), + group_ring: VecDeque::with_capacity(16), + groups: HashMap::default(), + handle_group: SecondaryMap::with_capacity(capacity), states: SecondaryMap::with_capacity(capacity), + weight_resolver, + } + } + + /// Adds the handle to its service group's ring, creating the group (and resolving + /// its weight) if needed. O(1); the weight resolver only runs on group creation. + fn push_to_group(&mut self, handle: VQueueHandle) { + let group_name = self + .handle_group + .get(handle) + .cloned() + .unwrap_or_else(|| SchedulingGroup::Service(unlinked_group())); + if let Some(group) = self.groups.get_mut(&group_name) { + group.ring.push_back(handle); + } else { + let weight = (self.weight_resolver)(&group_name); + self.groups + .insert(group_name.clone(), ServiceGroup::new(weight, handle)); + self.group_ring.push_back(group_name); } } - pub fn insert_eligible(&mut self, handle: VQueueHandle) { + /// Removes the front group if its ring is empty. The group being removed is always + /// at the front of the group ring, so this is O(1). + fn drop_front_group_if_empty(&mut self) { + let Some(front) = self.group_ring.front() else { + return; + }; + if self + .groups + .get(front) + .is_some_and(|group| group.ring.is_empty()) + { + let name = self.group_ring.pop_front().expect("front exists"); + self.groups.remove(&name); + } + } + + /// The handle currently at the dispatch position: front of the front group's ring. + fn front_handle(&self) -> Option { + let front_group = self.group_ring.front()?; + self.groups.get(front_group)?.ring.front().copied() + } + + pub fn insert_eligible(&mut self, handle: VQueueHandle, group: SchedulingGroup) { self.states.insert(handle, State::NeedsPoll); - self.ready_ring.push_back(handle); + self.handle_group.insert(handle, group); + self.push_to_group(handle); + } + + /// Records (or refreshes) the group a handle belongs to. Used by paths that have + /// access to the vqueue metadata. + fn record_group(&mut self, handle: VQueueHandle, meta: &VQueueMeta) { + if !self.handle_group.contains_key(handle) { + self.handle_group.insert(handle, SchedulingGroup::of(meta)); + } } pub fn get_status( @@ -118,8 +269,13 @@ impl EligibilityTracker { match state { State::Scheduled { wake_up } if wake_up.timer_key == timer_key => { *state = State::NeedsPoll; - debug_assert!(!self.ready_ring.contains(&handle)); - self.ready_ring.push_back(handle); + debug_assert!( + self.groups + .values() + .all(|group| !group.ring.contains(&handle)), + "scheduled handle must not already be in a ring" + ); + self.push_to_group(handle); } _ => {} } @@ -134,83 +290,129 @@ impl EligibilityTracker { storage: &S, vqueues: &mut SecondaryMap>, ) -> Result, StorageError> { - let n = self.ready_ring.len(); - // avoid rescanning the ready ring multiple rounds - for _ in 0..n { - // what is my current status - let Some(handle) = self.ready_ring.front().copied() else { + // avoid rescanning the group ring multiple rounds: every iteration either fully + // scans + rotates a group (at most num_groups times) or drops an emptied group + // (at most num_groups times), so 2 * num_groups bounds the scan. With a single + // group no rotation or drop-and-continue can occur, so one scan suffices — + // matching the flat ring's single pass. + let num_groups = self.group_ring.len(); + let scan_bound = if num_groups <= 1 { + num_groups + } else { + num_groups.saturating_mul(2) + }; + 'groups: for _ in 0..scan_bound { + let Some(group_name) = self.group_ring.front().cloned() else { return Ok(None); }; - - let Some(qstate) = vqueues.get_mut(handle) else { - // the vqueue has been removed probably it's empty therefore, it's - // safe to assume that it's not eligible. - self.remove(handle); - self.ready_ring.pop_front(); - continue; - }; - - let Some(current_state) = self.states.get_mut(handle) else { - // The vqueue is not eligible anymore. This can happen if the vqueue became dormant - // due to items being removed externally (killed, etc.) - self.ready_ring.pop_front(); - continue; - }; - - let slot = metas.get(handle).unwrap(); - - match current_state { - State::NeedsPoll => { - // update the state based on eligibility. - match qstate.poll_eligibility(cx, slot, storage) { - Poll::Ready(Ok(Eligibility::Eligible)) => { - *current_state = State::Ready; - return Ok(Some(handle)); - } - Poll::Ready(Ok(Eligibility::EligibleAt(ts))) => { - let ts = ts.as_unix_millis(); - let timer_key = self.delayed_eligibility.insert( - handle, - // Cap the parked delay. See`MAX_PARK`. - ts.duration_since(SchedulerClock.now_millis()).min(MAX_PARK), - ); - *current_state = State::Scheduled { - wake_up: WakeUp { ts, timer_key }, - }; - self.ready_ring.pop_front(); - continue; - } - Poll::Ready(Ok(Eligibility::NotEligible)) => { - self.ready_ring.pop_front(); - self.remove(handle); - continue; - } - Poll::Ready(Err(err)) => { - return Err(err); - } - Poll::Pending => { - self.ready_ring.rotate_left(1); - continue; + let group = self.groups.get(&group_name).expect("ring/groups in sync"); + + // avoid rescanning this group's ring multiple rounds + let n = group.ring.len(); + for _ in 0..n { + if self.group_ring.front() != Some(&group_name) { + // the group drained and was dropped during the scan; the new front + // group gets its own full scan without an outer rotation + continue 'groups; + } + // what is my current status + let Some(handle) = self.front_handle() else { + // group drained during the scan + break; + }; + + let Some(qstate) = vqueues.get_mut(handle) else { + // the vqueue has been removed probably it's empty therefore, it's + // safe to assume that it's not eligible. + self.remove(handle); + self.pop_front_handle(); + continue; + }; + + let Some(current_state) = self.states.get_mut(handle) else { + // The vqueue is not eligible anymore. This can happen if the vqueue became dormant + // due to items being removed externally (killed, etc.) + self.pop_front_handle(); + continue; + }; + + let slot = metas.get(handle).unwrap(); + + match current_state { + State::NeedsPoll => { + // update the state based on eligibility. + match qstate.poll_eligibility(cx, slot, storage) { + Poll::Ready(Ok(Eligibility::Eligible)) => { + *current_state = State::Ready; + return Ok(Some(handle)); + } + Poll::Ready(Ok(Eligibility::EligibleAt(ts))) => { + let ts = ts.as_unix_millis(); + // Cap the parked delay. See `MAX_PARK`. + let duration = + ts.duration_since(SchedulerClock.now_millis()).min(MAX_PARK); + let timer_key = self.delayed_eligibility.insert(handle, duration); + *self.states.get_mut(handle).expect("state exists") = + State::Scheduled { + wake_up: WakeUp { ts, timer_key }, + }; + self.pop_front_handle(); + continue; + } + Poll::Ready(Ok(Eligibility::NotEligible)) => { + self.pop_front_handle(); + self.remove(handle); + continue; + } + Poll::Ready(Err(err)) => { + return Err(err); + } + Poll::Pending => { + self.rotate_one(); + continue; + } } } - } - State::BlockedOn(_) | State::Scheduled { .. } => { - self.ready_ring.pop_front(); - } - State::Ready => { - return Ok(Some(handle)); + State::BlockedOn(_) | State::Scheduled { .. } => { + self.pop_front_handle(); + } + State::Ready => { + return Ok(Some(handle)); + } } } + + // No eligible handle in this group; move on to the next group. + self.drop_front_group_if_empty(); + if self.group_ring.len() <= 1 { + continue 'groups; + } + self.group_ring.rotate_left(1); } Ok(None) } + /// Pops the handle at the dispatch position (front of the front group's ring) and + /// removes the group if it became empty. O(1). + fn pop_front_handle(&mut self) { + if let Some(front_group) = self.group_ring.front() + && let Some(group) = self.groups.get_mut(front_group) + { + group.ring.pop_front(); + self.drop_front_group_if_empty(); + } + } + pub fn remove(&mut self, handle: VQueueHandle) { // cancel scheduled wake ups if let Some(State::Scheduled { wake_up }) = self.states.remove(handle) { self.delayed_eligibility.remove(&wake_up.timer_key); } + // Note: the handle is lazily purged from its group's ring on the next + // `next_eligible` scan (missing state => pop), exactly like the previous + // flat-ring implementation. `handle_group` is intentionally kept: the queue's + // service never changes for the lifetime of the cache slot. } // Places the vqueue back on the ready ring only if it was not already there. @@ -221,7 +423,7 @@ impl EligibilityTracker { Entry::Occupied(mut occupied_entry) => match occupied_entry.get() { State::BlockedOn(_resource) => { occupied_entry.insert(State::NeedsPoll); - self.ready_ring.push_back(vqueue); + self.push_to_group(vqueue); true } State::NeedsPoll => false, @@ -234,13 +436,13 @@ impl EligibilityTracker { State::Scheduled { wake_up } => { self.delayed_eligibility.remove(&wake_up.timer_key); occupied_entry.insert(State::NeedsPoll); - self.ready_ring.push_back(vqueue); + self.push_to_group(vqueue); true } }, Entry::Vacant(vacant_entry) => { vacant_entry.insert(State::NeedsPoll); - self.ready_ring.push_back(vqueue); + self.push_to_group(vqueue); true } } @@ -269,7 +471,7 @@ impl EligibilityTracker { Entry::Occupied(mut occupied_entry) => match occupied_entry.get() { State::BlockedOn(_resource) => { occupied_entry.insert(State::NeedsPoll); - self.ready_ring.push_back(vqueue); + self.push_to_group(vqueue); } _ => { // do nothing. @@ -277,7 +479,7 @@ impl EligibilityTracker { }, Entry::Vacant(vacant_entry) => { vacant_entry.insert(State::NeedsPoll); - self.ready_ring.push_back(vqueue); + self.push_to_group(vqueue); } } } @@ -291,6 +493,8 @@ impl EligibilityTracker { meta: &VQueueMeta, vqueue: &VQueueState, ) { + self.record_group(handle, meta); + let Some(current_state) = self.states.entry(handle) else { // the vqueue handle was removed from the original slot map. return; @@ -299,7 +503,7 @@ impl EligibilityTracker { match current_state { Entry::Vacant(vacant_entry) => { vacant_entry.insert(State::NeedsPoll); - self.ready_ring.push_back(handle); + self.push_to_group(handle); } Entry::Occupied(mut occupied_entry) => { let eligibility = vqueue.check_eligibility(meta).as_compact(); @@ -315,13 +519,13 @@ impl EligibilityTracker { // happens, we err on the safe side and switch to polling. self.delayed_eligibility.remove(&wake_up.timer_key); occupied_entry.insert(State::NeedsPoll); - self.ready_ring.push_back(handle); + self.push_to_group(handle); } (State::Scheduled { wake_up }, Eligibility::Eligible) => { // wake up now self.delayed_eligibility.remove(&wake_up.timer_key); self.states.insert(handle, State::NeedsPoll); - self.ready_ring.push_back(handle); + self.push_to_group(handle); } (State::Scheduled { wake_up }, Eligibility::EligibleAt(eligible_at_ts)) => { let eligible_at_ts = eligible_at_ts.as_unix_millis(); @@ -367,7 +571,7 @@ impl EligibilityTracker { if matches!(current_state, State::BlockedOn(_)) { let blocked_on = std::mem::replace(current_state, State::NeedsPoll); - self.ready_ring.push_back(handle); + self.push_to_group(handle); let State::BlockedOn(resource) = blocked_on else { unreachable!(); }; @@ -378,31 +582,85 @@ impl EligibilityTracker { } } + /// Rotates the current group's inner ring (DRR "needs credit" step). Group quota is + /// not consumed since nothing was dispatched. pub fn rotate_one(&mut self) { - self.ready_ring.rotate_left(1); + if let Some(front_group) = self.group_ring.front() + && let Some(group) = self.groups.get_mut(front_group) + { + group.ring.rotate_left(1); + } } pub fn len(&self) -> usize { - self.ready_ring.len() + // Only used for periodic reporting and tests; group count is small (bounded by + // the number of services), so summing is cheap. + self.groups.values().map(|group| group.ring.len()).sum() } + /// Called after a successful dispatch (Run/Yield): flips the dispatched queue back + /// to NeedsPoll and consumes one unit of the group's WRR quota, rotating the group + /// ring when the quota is exhausted. pub fn front_needs_poll(&mut self) { - if let Some(handle) = self.ready_ring.front() - && let Some(state) = self.states.get_mut(*handle) - { + let Some(handle) = self.front_handle() else { + return; + }; + if let Some(state) = self.states.get_mut(handle) { *state = State::NeedsPoll; } + if let Some(front_group) = self.group_ring.front() + && let Some(group) = self.groups.get_mut(front_group) + && group.quota.consume_slot() + && self.group_ring.len() > 1 + { + self.group_ring.rotate_left(1); + } } pub fn front_blocked(&mut self, resource: ResourceKind) { - if let Some(handle) = self.ready_ring.front() - && let Some(state) = self.states.get_mut(*handle) - { + let Some(handle) = self.front_handle() else { + return; + }; + if let Some(state) = self.states.get_mut(handle) { *state = State::BlockedOn(resource); } + // A blocked queue leaves the ring until a resource release wakes it up. + self.pop_front_handle(); } pub fn compact_memory(&mut self) { self.delayed_eligibility.compact(); } + + /// Verifies internal data-structure invariants. Test-only. + #[cfg(test)] + pub(crate) fn assert_invariants(&self) { + use std::collections::HashSet; + // group_ring contains exactly the keys of groups, each once + let ring_set: HashSet<&SchedulingGroup> = self.group_ring.iter().collect(); + assert_eq!( + ring_set.len(), + self.group_ring.len(), + "group_ring has duplicates" + ); + assert_eq!( + ring_set.len(), + self.groups.len(), + "group_ring and groups out of sync" + ); + for name in self.groups.keys() { + assert!(ring_set.contains(name), "group {name:?} missing from ring"); + } + // no handle appears in two different rings; every ringed handle has a group mapping + let mut seen: HashSet = HashSet::new(); + for group in self.groups.values() { + for &handle in &group.ring { + assert!(seen.insert(handle), "handle {handle:?} in multiple rings"); + assert!( + self.handle_group.contains_key(handle), + "ringed handle {handle:?} missing handle_group mapping" + ); + } + } + } } diff --git a/crates/vqueues/src/scheduler/resource_manager.rs b/crates/vqueues/src/scheduler/resource_manager.rs index 83081790f7..c1848dd500 100644 --- a/crates/vqueues/src/scheduler/resource_manager.rs +++ b/crates/vqueues/src/scheduler/resource_manager.rs @@ -8,7 +8,17 @@ // the Business Source License, use of this software will be governed // by the Apache License, Version 2.0. +mod grouped_waiters; mod invoker; + +/// Test-only handle to the grouped waiter list (weight-lifecycle tests live in +/// `scheduler.rs` next to the scope-weights plumbing they exercise). +#[cfg(test)] +pub(super) fn test_grouped_waiters( + weight_resolver: super::eligible::WeightResolver, +) -> grouped_waiters::GroupedWaiters { + grouped_waiters::GroupedWaiters::new(weight_resolver) +} mod invoker_memory; mod invoker_throttle; mod locks; @@ -45,7 +55,7 @@ use self::locks::Locks; use self::permit::ProvisionalPermit; use self::user_limiter::UserLimiter; use super::VQueueHandle; -use super::eligible::EligibilityTracker; +use super::eligible::{EligibilityTracker, SchedulingGroup, WeightResolver}; use crate::GlobalTokenBucket; // A set of queues waiting on a resource @@ -77,13 +87,17 @@ impl ResourceManager { global_throttling: Option, memory_pool: MemoryPool, initial_invocation_memory: NonZeroByteCount, + weight_resolver: WeightResolver, ) -> Result { let locks = Locks::create(storage).await?; let (_tx, rx) = mpsc::unbounded_channel(); Ok(Self { - invoker_concurrency: InvokerConcurrencyLimiter::new(concurrency_limiter), + invoker_concurrency: InvokerConcurrencyLimiter::new( + concurrency_limiter, + weight_resolver, + ), invoker_throttling: InvokerThrottlingLimiter::new(global_throttling), invoker_memory: InvokerMemoryLimiter::new(memory_pool, initial_invocation_memory), user_limiter: UserLimiter::create(), @@ -258,8 +272,11 @@ impl ResourceManager { // Do we have one? if !current_permit.has_invoker_permit() { // poll for one or die trying - let Some(invoker_permit) = self.invoker_concurrency.poll_acquire(cx, vqueue) - else { + let Some(invoker_permit) = self.invoker_concurrency.poll_acquire( + cx, + vqueue, + &SchedulingGroup::of(meta), + ) else { return AcquireOutcome::BlockedOn(ResourceKind::InvokerConcurrency); }; current_permit.set_invoker_permit(invoker_permit); diff --git a/crates/vqueues/src/scheduler/resource_manager/grouped_waiters.rs b/crates/vqueues/src/scheduler/resource_manager/grouped_waiters.rs new file mode 100644 index 0000000000..1789862f3d --- /dev/null +++ b/crates/vqueues/src/scheduler/resource_manager/grouped_waiters.rs @@ -0,0 +1,354 @@ +// Copyright (c) 2023 - 2026 Restate Software, Inc., Restate GmbH. +// All rights reserved. +// +// Use of this software is governed by the Business Source License +// included in the LICENSE file. +// +// As of the Change Date specified in that file, in accordance with +// the Business Source License, use of this software will be governed +// by the Apache License, Version 2.0. + +use std::collections::VecDeque; + +use slotmap::SecondaryMap; + +use crate::scheduler::VQueueHandle; +use crate::scheduler::eligible::{SchedulingGroup, WeightResolver, WrrQuota}; + +/// Two-level weighted round-robin waiter list. +/// +/// A plain FIFO waiter list reintroduces the starvation the two-level WRR +/// eligibility ring eliminates: when a resource (e.g. invoker concurrency) is +/// exhausted, *arrival order* decides who acquires next, so a service with +/// thousands of queued keys pushes every other service to the back regardless +/// of scheduling weights. This structure applies the same service-group WRR to +/// the waiter list itself: the head is the front of the front group's queue, +/// and a group's quota (its scheduling weight) controls how many acquisitions +/// it gets before the group ring rotates. +/// +/// Invariants: `ring` contains exactly the keys of `groups`; a handle appears +/// at most once (tracked in `membership`); `total` == sum of group queue lens. +pub(in crate::scheduler) struct GroupedWaiters { + ring: VecDeque, + groups: hashbrown::HashMap, + /// Which group each waiting handle sits in. SecondaryMap (direct slotmap + /// index) rather than a hash map: `VQueueHandle` is a slotmap key, matching + /// the `EligibilityTracker::handle_group` idiom and avoiding per-op hashing. + membership: SecondaryMap, + total: usize, + weight_resolver: WeightResolver, +} + +struct GroupQueue { + quota: WrrQuota, + queue: VecDeque, +} + +impl GroupedWaiters { + pub fn new(weight_resolver: WeightResolver) -> Self { + Self { + ring: VecDeque::new(), + groups: hashbrown::HashMap::new(), + membership: SecondaryMap::new(), + total: 0, + weight_resolver, + } + } + + pub fn is_empty(&self) -> bool { + self.total == 0 + } + + pub fn len(&self) -> usize { + self.total + } + + /// The handle at the current WRR head (front of the front group). Only used + /// by tests; permit claims are not head-gated. + #[cfg(test)] + pub fn front(&self) -> Option { + let group = self.ring.front()?; + self.groups.get(group)?.queue.front().copied() + } + + /// Adds a waiter to its scheduling group's queue (no-op if already waiting). + pub fn push_back(&mut self, handle: VQueueHandle, group: &SchedulingGroup) { + // single membership probe: entry() combines the duplicate check + insert + let Some(entry) = self.membership.entry(handle) else { + return; + }; + let slotmap::secondary::Entry::Vacant(vacant) = entry else { + return; + }; + let group = group.clone(); + if let Some(gq) = self.groups.get_mut(&group) { + gq.queue.push_back(handle); + vacant.insert(group); + } else { + let weight = (self.weight_resolver)(&group); + let mut queue = VecDeque::with_capacity(4); + queue.push_back(handle); + self.groups.insert( + group.clone(), + GroupQueue { + quota: WrrQuota::new(weight), + queue, + }, + ); + self.ring.push_back(group.clone()); + vacant.insert(group); + } + self.total += 1; + } + + /// Pops the current head (front of the front group), consuming one unit of + /// the group's quota. Rotates the group ring when the quota is exhausted + /// and drops the group when its queue empties. + pub fn pop_front(&mut self) -> Option { + let group_name = self.ring.front()?.clone(); + let gq = self.groups.get_mut(&group_name)?; + let handle = gq.queue.pop_front()?; + self.membership.remove(handle); + self.total -= 1; + + let quota_exhausted = gq.quota.consume_slot(); + + if gq.queue.is_empty() { + self.groups.remove(&group_name); + self.ring.pop_front(); + } else if quota_exhausted && self.ring.len() > 1 { + self.ring.rotate_left(1); + } + + Some(handle) + } + + /// Removes a waiter wherever it sits (external cancellation). O(group len). + pub fn remove(&mut self, handle: VQueueHandle) { + let Some(group_name) = self.membership.remove(handle) else { + return; + }; + if let Some(gq) = self.groups.get_mut(&group_name) { + gq.queue.retain(|h| *h != handle); + self.total -= 1; + if gq.queue.is_empty() { + self.groups.remove(&group_name); + self.ring.retain(|g| g != &group_name); + } + } + } +} + +#[cfg(test)] +mod tests { + use std::num::NonZeroU32; + use std::sync::Arc; + + use slotmap::SlotMap; + + use restate_types::ServiceName; + + use super::*; + + /// Test groups: derive the weight from a `-w` suffix on the group's + /// service name so weighted scenarios don't need a live rule map. + fn resolver() -> WeightResolver { + Arc::new(|group: &SchedulingGroup| { + let SchedulingGroup::Service(service_name) = group else { + return NonZeroU32::MIN; + }; + let name: &str = service_name.as_ref(); + name.rsplit_once("-w") + .and_then(|(_, w)| w.parse::().ok()) + .and_then(NonZeroU32::new) + .unwrap_or(NonZeroU32::MIN) + }) + } + + fn group(name: &str) -> SchedulingGroup { + SchedulingGroup::Service(ServiceName::new(name)) + } + + fn handles(n: usize) -> Vec { + let mut map = SlotMap::::with_key(); + (0..n).map(|_| map.insert(())).collect() + } + + #[test] + fn fifo_within_single_group() { + let hs = handles(3); + let mut w = GroupedWaiters::new(resolver()); + let g = group("svc"); + for h in &hs { + w.push_back(*h, &g); + } + assert_eq!(w.len(), 3); + assert_eq!(w.pop_front(), Some(hs[0])); + assert_eq!(w.pop_front(), Some(hs[1])); + assert_eq!(w.pop_front(), Some(hs[2])); + assert!(w.is_empty()); + } + + #[test] + fn weighted_interleave_between_groups() { + // payment-w1 has 20 waiters queued FIRST; batcher-w10 has 4 queued after. + // A FIFO would grant all 20 payment permits before any batcher permit; + // the grouped list must serve batchers within the first WRR cycle. + let payment_handles = handles(24); + let (pay, batch) = payment_handles.split_at(20); + let mut w = GroupedWaiters::new(resolver()); + let payment = group("payment-w1"); + let batcher = group("batcher-w10"); + for h in pay { + w.push_back(*h, &payment); + } + for h in batch { + w.push_back(*h, &batcher); + } + + let mut first_batcher_position = None; + for i in 0.. { + let Some(h) = w.pop_front() else { break }; + if batch.contains(&h) && first_batcher_position.is_none() { + first_batcher_position = Some(i); + } + } + // payment (weight 1) yields after 1 grant; batcher group is next + assert_eq!(first_batcher_position, Some(1)); + } + + #[test] + fn duplicate_push_is_noop_and_remove_clears() { + let hs = handles(2); + let mut w = GroupedWaiters::new(resolver()); + let g = group("svc"); + w.push_back(hs[0], &g); + w.push_back(hs[0], &g); + assert_eq!(w.len(), 1); + w.push_back(hs[1], &g); + w.remove(hs[0]); + assert_eq!(w.len(), 1); + assert_eq!(w.front(), Some(hs[1])); + w.remove(hs[1]); + assert!(w.is_empty()); + assert!(w.ring.is_empty()); + assert!(w.groups.is_empty()); + } + + /// Removing the only waiter of a NON-front group must drop that group from + /// both `groups` and `ring` (exercises the retain-based removal branch). + #[test] + fn remove_last_member_of_non_front_group() { + let hs = handles(3); + let mut w = GroupedWaiters::new(resolver()); + let a = group("a"); + let b = group("b"); + w.push_back(hs[0], &a); // front group + w.push_back(hs[1], &a); + w.push_back(hs[2], &b); // non-front group, single member + w.remove(hs[2]); + assert_eq!(w.len(), 2); + assert!(!w.ring.contains(&b), "ring must drop the emptied group"); + assert!( + !w.groups.contains_key(&b), + "groups must drop the emptied group" + ); + // remaining group still drains normally + assert_eq!(w.pop_front(), Some(hs[0])); + assert_eq!(w.pop_front(), Some(hs[1])); + assert!(w.is_empty()); + } + + /// Weight changes apply on group re-creation (after full drain), not + /// retroactively — same semantics as the eligibility ring. + #[test] + fn weight_change_applies_after_group_recreation() { + use std::sync::atomic::{AtomicU32, Ordering}; + let hs = handles(6); + let weight = Arc::new(AtomicU32::new(1)); + let w_ref = Arc::clone(&weight); + let dyn_resolver: WeightResolver = Arc::new(move |_: &SchedulingGroup| { + NonZeroU32::new(w_ref.load(Ordering::Relaxed)).unwrap() + }); + let mut w = GroupedWaiters::new(dyn_resolver); + let a = group("dyn"); + let b = group("other"); + + // group created at weight 1: after 1 grant it must yield to `other` + w.push_back(hs[0], &a); + w.push_back(hs[1], &a); + w.push_back(hs[2], &b); + assert_eq!(w.pop_front(), Some(hs[0])); + assert_eq!( + w.pop_front(), + Some(hs[2]), + "weight-1 group yields after one grant" + ); + assert_eq!(w.pop_front(), Some(hs[1])); + assert!(w.is_empty()); + + // bump the weight; the re-created group must get 2 grants before yielding + weight.store(2, Ordering::Relaxed); + w.push_back(hs[3], &a); + w.push_back(hs[4], &a); + w.push_back(hs[5], &b); + assert_eq!(w.pop_front(), Some(hs[3])); + assert_eq!( + w.pop_front(), + Some(hs[4]), + "weight-2 group gets both grants first" + ); + assert_eq!(w.pop_front(), Some(hs[5])); + } + + /// All-weight-1 with many groups: pop order must be a clean round-robin + /// cycle across services, not FIFO within a cycle. + #[test] + fn equal_weights_round_robin_across_many_groups() { + const GROUPS: usize = 5; + const PER_GROUP: usize = 10; + let hs = handles(GROUPS * PER_GROUP); + let names: Vec = (0..GROUPS).map(|g| group(&format!("svc{g}"))).collect(); + let mut w = GroupedWaiters::new(resolver()); + // enqueue all of group 0 first, then group 1, etc. (worst case for FIFO) + for (g, name) in names.iter().enumerate() { + for q in 0..PER_GROUP { + w.push_back(hs[g * PER_GROUP + q], name); + } + } + // every consecutive window of GROUPS pops must contain one handle from + // each group + let mut popped = Vec::new(); + while let Some(h) = w.pop_front() { + popped.push(h); + } + assert_eq!(popped.len(), GROUPS * PER_GROUP); + for (cycle, window) in popped.chunks(GROUPS).enumerate() { + let mut seen = [false; GROUPS]; + for h in window { + let idx = hs.iter().position(|x| x == h).unwrap() / PER_GROUP; + assert!( + !seen[idx], + "cycle {cycle}: group {idx} served twice within one round" + ); + seen[idx] = true; + } + } + } + + #[test] + fn work_conserving_when_group_drains() { + let hs = handles(6); + let mut w = GroupedWaiters::new(resolver()); + let a = group("a-w10"); + let b = group("b-w1"); + w.push_back(hs[0], &a); + for h in &hs[1..] { + w.push_back(*h, &b); + } + // a's single waiter drains; every remaining grant goes to b + let order: Vec<_> = std::iter::from_fn(|| w.pop_front()).collect(); + assert_eq!(order.len(), 6); + assert_eq!(order[0], hs[0]); + } +} diff --git a/crates/vqueues/src/scheduler/resource_manager/invoker.rs b/crates/vqueues/src/scheduler/resource_manager/invoker.rs index 684c822af9..9c49fffb1d 100644 --- a/crates/vqueues/src/scheduler/resource_manager/invoker.rs +++ b/crates/vqueues/src/scheduler/resource_manager/invoker.rs @@ -12,65 +12,69 @@ use std::task::Poll; use restate_futures_util::concurrency::{Concurrency, Permit}; -use super::Waiters; +use super::grouped_waiters::GroupedWaiters; use crate::scheduler::VQueueHandle; +use crate::scheduler::eligible::{SchedulingGroup, WeightResolver}; pub struct InvokerConcurrencyLimiter { limiter: Concurrency, - waiters: Waiters, + // Weighted round-robin over service groups instead of a flat FIFO: with a + // FIFO, whichever service floods the waiter list first monopolizes freed + // permits and starves every other service regardless of scheduling weights. + waiters: GroupedWaiters, cached_permit: Permit, } impl InvokerConcurrencyLimiter { - pub fn new(limiter: Concurrency) -> Self { + pub fn new(limiter: Concurrency, weight_resolver: WeightResolver) -> Self { Self { limiter, - waiters: Default::default(), + waiters: GroupedWaiters::new(weight_resolver), cached_permit: Permit::new_empty(), } } pub fn remove_from_waiters(&mut self, vqueue: VQueueHandle) { - self.waiters.retain(|h| *h != vqueue); - } - - fn pop_and_advance(&mut self, cx: &mut std::task::Context<'_>) { - // pop ourselves. - self.waiters.pop_front(); - - if !self.waiters.is_empty() { - cx.waker().wake_by_ref(); - } + self.waiters.remove(vqueue); } + /// Attempts to claim a permit for a queue the scheduler decided to dispatch. + /// + /// Unlike the previous FIFO design, this does NOT gate the claim on the + /// caller being the waiter-list head: the eligibility ring already applies + /// weighted round-robin when choosing which queue polls, and with a + /// rotating (WRR) waiter head the head-only gate livelocks — the woken + /// queue arrives at dispatch after the head has rotated past it, fails, + /// re-queues, and the freed permit is never claimed. The waiter list's job + /// is reduced to picking the wake-up order (see `poll_head`), which is + /// where the group weights apply. pub(super) fn poll_acquire( &mut self, cx: &mut std::task::Context<'_>, vqueue: VQueueHandle, + group: &SchedulingGroup, ) -> Option { - if self.waiters.is_empty() { - self.waiters.push_back(vqueue); + // cached permit exists (set aside by poll_head when it woke a waiter) + if let Some(permit) = self.cached_permit.split(1) { + self.claimed(cx, vqueue); + return Some(permit); } - if self.waiters.front().is_some_and(|waiter| waiter == &vqueue) { - tracing::trace!("Checking the head vqueue {vqueue:?} for invoker concurrency permits"); - // cached permit exists, return it. - if let Some(permit) = self.cached_permit.split(1) { - self.pop_and_advance(cx); - return Some(permit); - } + if let Poll::Ready(invoker_permit) = self.limiter.poll_acquire(cx) { + self.claimed(cx, vqueue); + return Some(invoker_permit); + } - if let Poll::Ready(invoker_permit) = self.limiter.poll_acquire(cx) { - self.pop_and_advance(cx); - Some(invoker_permit) - } else { - // waker registered - None - } - } else { - // We need to stand behind - self.waiters.push_back(vqueue); - None + // No permit available: park this queue in its service group's wait list + // (push_back de-duplicates). + self.waiters.push_back(vqueue, group); + None + } + + fn claimed(&mut self, cx: &mut std::task::Context<'_>, vqueue: VQueueHandle) { + self.waiters.remove(vqueue); + if !self.waiters.is_empty() { + cx.waker().wake_by_ref(); } } @@ -106,3 +110,74 @@ impl InvokerConcurrencyLimiter { Poll::Pending } } + +#[cfg(test)] +mod tests { + use std::num::{NonZeroU32, NonZeroUsize}; + use std::sync::Arc; + + use slotmap::SlotMap; + + use restate_types::ServiceName; + + use super::*; + use crate::scheduler::eligible::WeightResolver; + + fn resolver() -> WeightResolver { + Arc::new(|_: &SchedulingGroup| NonZeroU32::MIN) + } + + fn group(name: &str) -> SchedulingGroup { + SchedulingGroup::Service(ServiceName::new(name)) + } + + /// The wake-then-steal race is intentional and livelock-free: `poll_head` + /// sets a permit aside and wakes waiter A, but whichever queue the + /// scheduler dispatches first may claim it. The loser simply re-parks and + /// the WRR waiter list routes a later permit to it — no permit is ever + /// stranded (the livelock the previous head-gated design had). + #[test] + fn woken_permit_can_be_claimed_by_another_queue_without_stranding() { + let mut handles = SlotMap::::with_key(); + let holder = handles.insert(()); + let vq_a = handles.insert(()); + let vq_b = handles.insert(()); + let group_a = group("a"); + let group_b = group("b"); + + let waker = std::task::Waker::noop(); + let mut cx = std::task::Context::from_waker(waker); + let mut limiter = InvokerConcurrencyLimiter::new( + Concurrency::new(Some(NonZeroUsize::new(1).unwrap())), + resolver(), + ); + + // holder takes the only permit; A and B park in the waiter list + let permit = limiter + .poll_acquire(&mut cx, holder, &group("holder")) + .expect("permit"); + assert!(limiter.poll_acquire(&mut cx, vq_a, &group_a).is_none()); + assert!(limiter.poll_acquire(&mut cx, vq_b, &group_b).is_none()); + + // release; poll_head caches the freed permit and wakes A (WRR head) + drop(permit); + let woken = limiter.poll_head(&mut cx); + assert!(matches!(woken, Poll::Ready(Some(h)) if h == vq_a)); + + // B "wins the race" to dispatch first and steals the cached permit + let stolen = limiter + .poll_acquire(&mut cx, vq_b, &group_b) + .expect("B claims the cached permit"); + + // A loses, re-parks — and the next released permit reaches A: nothing + // is stranded and the loser is not starved + assert!(limiter.poll_acquire(&mut cx, vq_a, &group_a).is_none()); + drop(stolen); + let woken = limiter.poll_head(&mut cx); + assert!(matches!(woken, Poll::Ready(Some(h)) if h == vq_a)); + assert!( + limiter.poll_acquire(&mut cx, vq_a, &group_a).is_some(), + "A claims the permit on its wake" + ); + } +} diff --git a/crates/worker-api/src/processor/features.rs b/crates/worker-api/src/processor/features.rs index 5a33863939..358aea7e92 100644 --- a/crates/worker-api/src/processor/features.rs +++ b/crates/worker-api/src/processor/features.rs @@ -33,6 +33,12 @@ pub trait PartitionFeatures { /// /// *Since v1.7.0* fn is_unique_random_seeds_enabled(&self) -> bool; + + /// Whether child invocations created via ctx.call inherit their caller's + /// scope when the child target carries none. + /// + /// *Since v1.7.0* + fn is_scope_inheritance_enabled(&self) -> bool; } impl PartitionFeatures for PersistedFeatures { @@ -50,6 +56,11 @@ impl PartitionFeatures for PersistedFeatures { fn is_unique_random_seeds_enabled(&self) -> bool { self.unique_random_seeds } + + #[inline] + fn is_scope_inheritance_enabled(&self) -> bool { + self.scope_inheritance + } } // -- Boilerplate -- @@ -66,6 +77,10 @@ impl PartitionFeatures for &T { fn is_unique_random_seeds_enabled(&self) -> bool { (**self).is_unique_random_seeds_enabled() } + + fn is_scope_inheritance_enabled(&self) -> bool { + (**self).is_scope_inheritance_enabled() + } } impl PartitionFeatures for &mut T { @@ -80,4 +95,8 @@ impl PartitionFeatures for &mut T { fn is_unique_random_seeds_enabled(&self) -> bool { (**self).is_unique_random_seeds_enabled() } + + fn is_scope_inheritance_enabled(&self) -> bool { + (**self).is_scope_inheritance_enabled() + } } diff --git a/crates/worker/src/partition/leadership/mod.rs b/crates/worker/src/partition/leadership/mod.rs index 44adb0db8c..8117dfc4a7 100644 --- a/crates/worker/src/partition/leadership/mod.rs +++ b/crates/worker/src/partition/leadership/mod.rs @@ -476,6 +476,14 @@ where })? .into_guard(); + // Scheduling weights are scope-keyed and arrive via the rules system + // (`restate rules set --weight N`): the scheduler service keeps + // the shared scope-weight map in sync from rule updates, and the + // resolver closures read it whenever a scheduling group is (re-)created. + // Service-name groups always run at the default weight 1. + let scope_weights = restate_vqueues::ScopeWeights::default(); + let weight_resolver = restate_vqueues::scope_weight_resolver(scope_weights.clone()); + let scheduler_service = SchedulerService::create( ResourceManager::create( partition_store.partition_db().clone(), @@ -483,10 +491,13 @@ where self.invoker_capacity.invocation_token_bucket.clone(), self.invoker_capacity.memory_pool.clone(), self.invoker_capacity.initial_invocation_memory, + weight_resolver.clone(), ) .await?, partition_store.partition_db().clone(), vqueues_cache, + weight_resolver, + scope_weights, ) .await?; @@ -586,6 +597,14 @@ where feature_changes.push(PartitionFeatureChange::EnableUniqueRandomSeeds); } + // Children of scoped invocations join their caller's scope, so + // scope-level rules govern whole call trees. + if config.common.experimental.is_scope_inheritance_enabled() + && !state_machine_features.is_scope_inheritance_enabled() + { + feature_changes.push(PartitionFeatureChange::EnableScopeInheritance); + } + if !feature_changes.is_empty() { // Smallest version that supports every listed feature, but never below // the partition's current min_restate_version. @@ -993,6 +1012,10 @@ mod tests { fn is_unique_random_seeds_enabled(&self) -> bool { false } + + fn is_scope_inheritance_enabled(&self) -> bool { + false + } } #[test(restate_core::test)] diff --git a/crates/worker/src/partition/state_machine/entries/call_commands.rs b/crates/worker/src/partition/state_machine/entries/call_commands.rs index e98677ca53..37b3fc64af 100644 --- a/crates/worker/src/partition/state_machine/entries/call_commands.rs +++ b/crates/worker/src/partition/state_machine/entries/call_commands.rs @@ -20,6 +20,7 @@ use restate_types::journal_v2::command::{CallCommand, CallRequest, OneWayCallCom use restate_types::journal_v2::raw::RawEntry; use restate_types::journal_v2::{CallInvocationIdCompletion, CompletionId, Entry}; use restate_types::time::MillisSinceEpoch; +use restate_worker_api::processor::PartitionFeatures; use std::collections::VecDeque; pub(super) type ApplyCallCommand<'e> = ApplyJournalCommandEffect<'e, CallCommand>; @@ -94,7 +95,7 @@ where let CallRequest { invocation_id, - invocation_target, + mut invocation_target, span_context, parameter, headers, @@ -104,6 +105,15 @@ where limit_key, } = self.request; + // A child without its own scope joins its caller's scope. Deterministic: + // gated on the partition-persisted feature set, replicated state only. + if ctx.is_scope_inheritance_enabled() + && invocation_target.scope().is_none() + && let Some(scope) = caller_invocation_metadata.invocation_target.scope() + { + invocation_target = invocation_target.with_scope(Some(scope.clone())); + } + // Prepare the service invocation to propose let service_invocation = ServiceInvocation { argument: parameter, diff --git a/crates/worker/src/partition/state_machine/lifecycle/version_barrier.rs b/crates/worker/src/partition/state_machine/lifecycle/version_barrier.rs index 79a08fa0e9..f688043137 100644 --- a/crates/worker/src/partition/state_machine/lifecycle/version_barrier.rs +++ b/crates/worker/src/partition/state_machine/lifecycle/version_barrier.rs @@ -60,6 +60,9 @@ where // point. Pre-existing invocations without a stored random seed keep working via the // `to_random_seed()` fallback in `invoker_storage_reader.rs`. PartitionFeatureChange::EnableUniqueRandomSeeds => Ok(false), + // Flipping scope inheritance on only affects child invocations created + // after the apply point; pre-existing invocations are untouched. + PartitionFeatureChange::EnableScopeInheritance => Ok(false), } } @@ -468,6 +471,7 @@ mod tests { journal_v2: false, vqueues: true, unique_random_seeds: false, + scope_inheritance: false, }, Default::default(), std::sync::Arc::new(RuleBook::default()), diff --git a/crates/worker/src/partition/state_machine/mod.rs b/crates/worker/src/partition/state_machine/mod.rs index a2857d5549..1520c11255 100644 --- a/crates/worker/src/partition/state_machine/mod.rs +++ b/crates/worker/src/partition/state_machine/mod.rs @@ -148,6 +148,10 @@ impl PartitionFeatures for StateMachine { fn is_unique_random_seeds_enabled(&self) -> bool { self.enabled_features.unique_random_seeds } + + fn is_scope_inheritance_enabled(&self) -> bool { + self.enabled_features.scope_inheritance + } } impl PartitionFeatures for StateMachineApplyContext<'_, S> { @@ -162,6 +166,10 @@ impl PartitionFeatures for StateMachineApplyContext<'_, S> { fn is_unique_random_seeds_enabled(&self) -> bool { self.enabled_features.unique_random_seeds } + + fn is_scope_inheritance_enabled(&self) -> bool { + self.enabled_features.scope_inheritance + } } pub struct StateMachine { // initialized from persistent storage @@ -3887,9 +3895,19 @@ impl StateMachineApplyContext<'_, S> { journal_entry.deserialize_entry_ref::()? ); + // Scope inheritance (journal-v2 equivalent: call_commands.rs). + let mut callee_invocation_target = callee_invocation_target.clone(); + if self.is_scope_inheritance_enabled() + && callee_invocation_target.scope().is_none() + && let Some(scope) = invocation_metadata.invocation_target.scope() + { + callee_invocation_target = + callee_invocation_target.with_scope(Some(scope.clone())); + } + let service_invocation = Box::new(ServiceInvocation { invocation_id: *callee_invocation_id, - invocation_target: callee_invocation_target.clone(), + invocation_target: callee_invocation_target, argument: request.parameter, source: Source::Service( invocation_id, @@ -3942,9 +3960,19 @@ impl StateMachineApplyContext<'_, S> { Some(MillisSinceEpoch::new(invoke_time)) }; + // Scope inheritance: see the Call arm above. + let mut callee_invocation_target = callee_invocation_target.clone(); + if self.is_scope_inheritance_enabled() + && callee_invocation_target.scope().is_none() + && let Some(scope) = invocation_metadata.invocation_target.scope() + { + callee_invocation_target = + callee_invocation_target.with_scope(Some(scope.clone())); + } + let service_invocation = Box::new(ServiceInvocation { invocation_id: *callee_invocation_id, - invocation_target: callee_invocation_target.clone(), + invocation_target: callee_invocation_target, argument: request.parameter, source: Source::Service( invocation_id, diff --git a/crates/worker/src/partition/state_machine/tests/mod.rs b/crates/worker/src/partition/state_machine/tests/mod.rs index a018be3378..03de124cf2 100644 --- a/crates/worker/src/partition/state_machine/tests/mod.rs +++ b/crates/worker/src/partition/state_machine/tests/mod.rs @@ -1323,3 +1323,82 @@ async fn yield_effect_resumes_invocation() { test_env.shutdown().await; } + +/// With the scope_inheritance partition feature enabled, a Call entry's child +/// invocation inherits the caller's scope when the callee target carries none. +#[test(restate_core::test)] +async fn call_child_inherits_caller_scope() -> TestResult { + run_call_child_scope_check(true).await +} + +/// With the feature disabled the child stays unscoped (regression guard). +#[test(restate_core::test)] +async fn call_child_unscoped_without_feature() -> TestResult { + run_call_child_scope_check(false).await +} + +async fn run_call_child_scope_check(feature_on: bool) -> TestResult { + use restate_types::Scope; + use restate_types::partitions::PartitionFeatureChange; + + let features = if feature_on { + PersistedFeatures::from_iter([PartitionFeatureChange::EnableScopeInheritance]) + } else { + PersistedFeatures::default() + }; + let mut test_env = TestEnv::create_with_features(features).await; + + let scope = Scope::try_non_interned("payment").unwrap(); + let caller_target = + InvocationTarget::service("ScopedCaller", "run").with_scope(Some(scope.clone())); + let invocation_id = + fixtures::mock_start_invocation_with_invocation_target(&mut test_env, caller_target).await; + + let callee_service_id = ServiceId::mock_random(); + let actions = test_env + .apply(commands::InvokerEffectCommand::test_envelope(Effect { + invocation_id, + kind: InvokerEffectKind::JournalEntry { + entry_index: 1, + entry: ProtobufRawEntryCodec::serialize_enriched(Entry::invoke( + InvokeRequest { + service_name: callee_service_id.service_name, + handler_name: "MyMethod".into(), + parameter: Bytes::default(), + headers: vec![], + key: callee_service_id.key, + idempotency_key: None, + }, + None, + )), + }, + })) + .await; + + let outbox_child = actions + .iter() + .find_map(|action| match action { + Action::NewOutboxMessage { + message: restate_storage_api::outbox_table::OutboxMessage::ServiceInvocation(si), + .. + } => Some(si.clone()), + _ => None, + }) + .expect("call entry must produce a child invocation"); + + if feature_on { + assert_eq!( + outbox_child.invocation_target.scope(), + Some(&scope), + "feature on: child inherits the caller's scope" + ); + } else { + assert_eq!( + outbox_child.invocation_target.scope(), + None, + "feature off: child stays unscoped" + ); + } + test_env.shutdown().await; + Ok(()) +}