diff --git a/Cargo.lock b/Cargo.lock index b8fbf8549f..8b948e1f9a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5329,7 +5329,9 @@ dependencies = [ "serde", "serde_json", "serde_path_to_error", + "si-data-nats", "telemetry-utils", + "test-log", "thiserror 2.0.12", "time", "tokio", @@ -5337,6 +5339,7 @@ dependencies = [ "tower 0.4.13", "tracing", "tracing-subscriber", + "uuid", ] [[package]] diff --git a/lib/naxum/BUCK b/lib/naxum/BUCK index 60fbbbfc10..79a96b4b51 100644 --- a/lib/naxum/BUCK +++ b/lib/naxum/BUCK @@ -1,4 +1,4 @@ -load("@prelude-si//:macros.bzl", "rust_library") +load("@prelude-si//:macros.bzl", "rust_library", "rust_test") rust_library( name = "naxum", @@ -12,6 +12,7 @@ rust_library( "//third-party/rust:serde", "//third-party/rust:serde_json", "//third-party/rust:serde_path_to_error", + "//third-party/rust:thiserror", "//third-party/rust:time", "//third-party/rust:tokio", "//third-party/rust:tokio-util", @@ -19,4 +20,28 @@ rust_library( "//third-party/rust:tracing", ], srcs = glob(["src/**/*.rs"]), + extra_test_targets = [":test-integration"], +) + +rust_test( + name = "test-integration", + deps = [ + "//lib/si-data-nats:si-data-nats", + "//third-party/rust:futures", + "//third-party/rust:test-log", + "//third-party/rust:tokio", + "//third-party/rust:tokio-util", + "//third-party/rust:tracing-subscriber", + "//third-party/rust:uuid", + ":naxum", + ], + srcs = glob([ + "tests/**/*.rs", + ]), + crate_root = "tests/integration/mod.rs", + env = { + "CARGO_PKG_NAME": "integration", + "RUSTC_BOOTSTRAP": "1", + "CI": "buildkite", + }, ) diff --git a/lib/naxum/Cargo.toml b/lib/naxum/Cargo.toml index f89dd1f867..31a0b5f8fe 100644 --- a/lib/naxum/Cargo.toml +++ b/lib/naxum/Cargo.toml @@ -18,6 +18,7 @@ serde = { workspace = true } # NOTE: if extracted serde_json = { workspace = true } # NOTE: if extracted, this is used for a json feature serde_path_to_error = { workspace = true } # NOTE: if extracted, this is used for a json feature telemetry-utils = { path = "../../lib/telemetry-utils-rs" } +thiserror = { workspace = true } time = { workspace = true } tokio = { workspace = true } tokio-util = { workspace = true } @@ -25,5 +26,7 @@ tower = { workspace = true } tracing = { workspace = true } # NOTE: left with vanilla tracing for potential future extraction [dev-dependencies] -thiserror = { workspace = true } +si-data-nats = { path = "../si-data-nats" } +test-log = { workspace = true } tracing-subscriber = { workspace = true } +uuid = { workspace = true } diff --git a/lib/naxum/src/fair/config.rs b/lib/naxum/src/fair/config.rs new file mode 100644 index 0000000000..b62f368708 --- /dev/null +++ b/lib/naxum/src/fair/config.rs @@ -0,0 +1,235 @@ +use std::{ + fmt, + marker::PhantomData, + time::Duration, +}; + +use async_nats::jetstream::stream::Stream as JsStream; + +pub type KeyExtractorFn = fn(&str, Option<&str>, &str) -> Option; +pub type ConsumerNameFn = fn(&K, &str) -> String; +pub type ConsumerFilterSubjectFn = fn(Option<&str>, &K, &str) -> String; + +pub struct FairSchedulingConfig { + pub(crate) service_name: String, + pub(crate) tasks_stream: JsStream, + pub(crate) requests_stream: JsStream, + pub(crate) subject_prefix: Option, + pub(crate) key_extractor: KeyExtractorFn, + pub(crate) consumer_name_fn: ConsumerNameFn, + pub(crate) consumer_filter_subject_fn: ConsumerFilterSubjectFn, + pub(crate) inactive_threshold: Duration, + pub(crate) task_listener_name: String, + pub(crate) tasks_filter_subject: String, + _key_marker: PhantomData, +} + +impl fmt::Debug for FairSchedulingConfig { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("FairSchedulingConfig") + .field("service_name", &self.service_name) + .field("subject_prefix", &self.subject_prefix) + .field("inactive_threshold", &self.inactive_threshold) + .field("task_listener_name", &self.task_listener_name) + .field("tasks_filter_subject", &self.tasks_filter_subject) + .finish_non_exhaustive() + } +} + +impl FairSchedulingConfig { + /// Creates a config for workspace-partitioned fair scheduling. + /// + /// Uses standard subject conventions where workspace_id is extracted from: + /// - Tasks: `{prefix?}.{service}.tasks.{workspace_id}[.{rest}]` + /// - Requests: `{prefix?}.{service}.requests.{workspace_id}.>` + /// + /// # Example + /// + /// ```ignore + /// let config = FairSchedulingConfig::for_workspace_partitioning( + /// "veritech", + /// tasks_stream, + /// requests_stream, + /// prefix, + /// ); + /// ``` + pub fn for_workspace_partitioning( + service_name: impl Into, + tasks_stream: JsStream, + requests_stream: JsStream, + subject_prefix: Option, + ) -> Self { + let service_name = service_name.into(); + + let tasks_filter_subject = + with_optional_prefix(subject_prefix.as_deref(), format!("{service_name}.tasks.*")); + + Self { + service_name: service_name.clone(), + tasks_stream, + requests_stream, + subject_prefix, + key_extractor: default_workspace_key_extractor, + consumer_name_fn: default_workspace_consumer_name, + consumer_filter_subject_fn: default_workspace_filter_subject, + inactive_threshold: Duration::from_secs(300), + task_listener_name: format!("{service_name}-task-listener"), + tasks_filter_subject, + _key_marker: PhantomData, + } + } +} + +fn with_optional_prefix(prefix: Option<&str>, base: impl AsRef) -> String { + match prefix { + Some(p) => format!("{}.{}", p, base.as_ref()), + None => base.as_ref().to_string(), + } +} + +/// Extracts workspace_id from position 2 (no prefix) or 3 (with prefix). +/// +/// Works for subjects like: `{prefix?}.{service}.tasks.{workspace_id}[.{rest}]` +fn default_workspace_key_extractor( + subject: &str, + prefix: Option<&str>, + _service_name: &str, +) -> Option { + let parts: Vec<&str> = subject.split('.').collect(); + let idx = if prefix.is_some() { 3 } else { 2 }; + parts.get(idx).map(|s| s.to_string()) +} + +/// Generates consumer name: `{service}-ws-{workspace_id}` +fn default_workspace_consumer_name(workspace_id: &String, service_name: &str) -> String { + format!("{service_name}-ws-{workspace_id}") +} + +/// Generates filter subject: `{prefix?}.{service}.requests.{workspace_id}.>` +fn default_workspace_filter_subject( + prefix: Option<&str>, + workspace_id: &String, + service_name: &str, +) -> String { + with_optional_prefix(prefix, format!("{service_name}.requests.{workspace_id}.>")) +} + +#[cfg(test)] +mod tests { + use super::{ + default_workspace_consumer_name, + default_workspace_filter_subject, + default_workspace_key_extractor, + }; + + #[test] + fn test_workspace_key_extractor_no_prefix() { + // Subject format: {service}.tasks.{workspace_id} + let subject = "veritech.tasks.workspace-123"; + let result = default_workspace_key_extractor(subject, None, "veritech"); + + assert_eq!(result, Some("workspace-123".to_string())); + } + + #[test] + fn test_workspace_key_extractor_with_prefix() { + // Subject format: {prefix}.{service}.tasks.{workspace_id} + let subject = "si.veritech.tasks.workspace-456"; + let result = default_workspace_key_extractor(subject, Some("si"), "veritech"); + + assert_eq!(result, Some("workspace-456".to_string())); + } + + #[test] + fn test_workspace_key_extractor_with_trailing_parts() { + // Subject format can have additional parts after workspace_id + let subject = "veritech.tasks.workspace-789.extra.parts"; + let result = default_workspace_key_extractor(subject, None, "veritech"); + + assert_eq!(result, Some("workspace-789".to_string())); + } + + #[test] + fn test_workspace_key_extractor_invalid_subject() { + // Too few parts + let subject = "veritech.tasks"; + let result = default_workspace_key_extractor(subject, None, "veritech"); + + assert_eq!(result, None); + } + + #[test] + fn test_workspace_key_extractor_empty_workspace_id() { + // Empty workspace ID should still be extracted (validation is elsewhere) + let subject = "veritech.tasks..something"; + let result = default_workspace_key_extractor(subject, None, "veritech"); + + assert_eq!(result, Some("".to_string())); + } + + #[test] + fn test_consumer_name_generation() { + let workspace_id = "workspace-123"; + let service_name = "veritech"; + + let result = default_workspace_consumer_name(&workspace_id.to_string(), service_name); + + assert_eq!(result, "veritech-ws-workspace-123"); + } + + #[test] + fn test_consumer_name_with_special_characters() { + let workspace_id = "workspace-test_123"; + let service_name = "my-service"; + + let result = default_workspace_consumer_name(&workspace_id.to_string(), service_name); + + assert_eq!(result, "my-service-ws-workspace-test_123"); + } + + #[test] + fn test_filter_subject_no_prefix() { + let workspace_id = "workspace-123"; + let service_name = "veritech"; + + let result = + default_workspace_filter_subject(None, &workspace_id.to_string(), service_name); + + assert_eq!(result, "veritech.requests.workspace-123.>"); + } + + #[test] + fn test_filter_subject_with_prefix() { + let workspace_id = "workspace-456"; + let service_name = "veritech"; + + let result = + default_workspace_filter_subject(Some("si"), &workspace_id.to_string(), service_name); + + assert_eq!(result, "si.veritech.requests.workspace-456.>"); + } + + #[test] + fn test_filter_subject_wildcard_suffix() { + // Verify the wildcard suffix is correctly added + let result = default_workspace_filter_subject(None, &"ws1".to_string(), "svc"); + + assert!( + result.ends_with(".>"), + "Filter subject must end with wildcard" + ); + } + + #[test] + fn test_key_extractor_consistency_with_subject_helpers() { + // Verify that the extractor can parse subjects generated by the helpers + let workspace_id = "test-workspace"; + + // Simulate what task_subject_for_workspace generates + let subject = format!("veritech.tasks.{workspace_id}"); + + // Should be extractable + let extracted = default_workspace_key_extractor(&subject, None, "veritech"); + assert_eq!(extracted, Some(workspace_id.to_string())); + } +} diff --git a/lib/naxum/src/fair/listener.rs b/lib/naxum/src/fair/listener.rs new file mode 100644 index 0000000000..f45304b517 --- /dev/null +++ b/lib/naxum/src/fair/listener.rs @@ -0,0 +1,226 @@ +use std::{ + hash::Hash, + time::Duration, +}; + +use async_nats::jetstream::consumer::pull::Config as ConsumerConfig; +use futures::StreamExt as _; +use telemetry_utils::monotonic; +use thiserror::Error; +use tokio::sync::mpsc; +use tokio_util::sync::CancellationToken; +use tracing::{ + debug, + error, + info, + warn, +}; + +use super::{ + config::FairSchedulingConfig, + stream::KeyReady, +}; + +const MAX_RETRY_DELAY: Duration = Duration::from_secs(60); +const BASE_RETRY_DELAY: Duration = Duration::from_secs(1); + +enum ConsumerCreationResult { + Success, + ChannelClosed, + CreationFailed, +} + +#[derive(Debug, Error)] +pub enum TaskListenerError { + #[error("failed to create task listener consumer: {0}")] + CreateConsumer(#[source] async_nats::jetstream::stream::ConsumerError), + #[error("failed to get messages stream: {0}")] + Messages(#[source] async_nats::jetstream::consumer::StreamError), + #[error("failed to get/create key consumer: {0}")] + GetKeyConsumer(#[source] async_nats::jetstream::stream::ConsumerError), + #[error("failed to get key consumer messages: {0}")] + KeyConsumerMessages(#[source] async_nats::jetstream::consumer::StreamError), +} + +pub fn spawn_task_listener( + config: FairSchedulingConfig, + consumer_tx: mpsc::Sender>, + shutdown: CancellationToken, +) -> tokio::task::JoinHandle<()> +where + K: Clone + Eq + Hash + Send + Sync + 'static, +{ + tokio::spawn(async move { + let mut retry_count = 0u32; + + loop { + tokio::select! { + _ = shutdown.cancelled() => { + debug!("task listener shutdown"); + break; + } + result = run_task_listener(&config, consumer_tx.clone(), shutdown.clone()) => { + match result { + Ok(()) => { + debug!("task listener completed"); + break; + } + Err(e) => { + retry_count += 1; + let delay = std::cmp::min( + BASE_RETRY_DELAY * 2u32.saturating_pow(retry_count.saturating_sub(1)), + MAX_RETRY_DELAY, + ); + error!(error = %e, retry_count, "task listener error, retrying after {}s", delay.as_secs()); + monotonic!(naxum.fair_scheduler.task_listener_restarts = 1); + tokio::time::sleep(delay).await; + } + } + } + } + } + }) +} + +async fn run_task_listener( + config: &FairSchedulingConfig, + consumer_tx: mpsc::Sender>, + shutdown: CancellationToken, +) -> Result<(), TaskListenerError> +where + K: Clone + Eq + Hash + Send + Sync + 'static, +{ + let consumer = config + .tasks_stream + .create_consumer(ConsumerConfig { + durable_name: Some(config.task_listener_name.clone()), + filter_subject: config.tasks_filter_subject.clone(), + ..Default::default() + }) + .await + .map_err(TaskListenerError::CreateConsumer)?; + + let mut messages = consumer + .messages() + .await + .map_err(TaskListenerError::Messages)?; + + loop { + tokio::select! { + _ = shutdown.cancelled() => { + debug!("task listener shutting down"); + break; + } + msg = messages.next() => { + let Some(msg) = msg else { + debug!("task stream ended"); + break; + }; + + let msg = match msg { + Ok(m) => m, + Err(e) => { + warn!(error = %e, "error receiving task message"); + continue; + } + }; + + let subject = msg.subject.as_str(); + monotonic!(naxum.fair_scheduler.task_notifications = 1); + + let Some(key) = (config.key_extractor)(subject, config.subject_prefix.as_deref(), &config.service_name) else { + warn!(%subject, "could not extract key"); + let _ = msg.ack().await; // Best effort ack for invalid message, ignore errors + continue; + }; + + // Always create/refresh the consumer for this key. The FairSchedulerStream + // handles deduplication - if a consumer already exists for this key, the new + // one replaces it. This ensures work is never missed when a workspace becomes + // active again after its previous consumer was exhausted. + let result = create_and_send_consumer(config, &key, &consumer_tx).await; + + // Ack on success or channel closed; don't ack on failure (let NATS redeliver) + let should_ack = !matches!(result, ConsumerCreationResult::CreationFailed); + + if should_ack { + if let Err(e) = msg.ack().await { + warn!(error = %e, "failed to ack task message"); + monotonic!(naxum.fair_scheduler.task_ack_errors = 1); + } + } + + // If channel closed, break the loop + if matches!(result, ConsumerCreationResult::ChannelClosed) { + break; + } + } + } + } + + Ok(()) +} + +async fn create_and_send_consumer( + config: &FairSchedulingConfig, + key: &K, + consumer_tx: &mpsc::Sender>, +) -> ConsumerCreationResult +where + K: Clone + Eq + Hash, +{ + let consumer_name = (config.consumer_name_fn)(key, &config.service_name); + let filter_subject = (config.consumer_filter_subject_fn)( + config.subject_prefix.as_deref(), + key, + &config.service_name, + ); + + debug!(%consumer_name, %filter_subject, "creating key consumer"); + + let consumer = match config + .requests_stream + .get_or_create_consumer( + &consumer_name, + ConsumerConfig { + durable_name: Some(consumer_name.clone()), + filter_subject, + inactive_threshold: config.inactive_threshold, + ..Default::default() + }, + ) + .await + { + Ok(c) => c, + Err(e) => { + warn!(error = %e, "failed to create key consumer"); + monotonic!(naxum.fair_scheduler.consumer_create_errors = 1); + return ConsumerCreationResult::CreationFailed; + } + }; + + let key_messages = match consumer.messages().await { + Ok(m) => m, + Err(e) => { + warn!(error = %e, "failed to get key consumer messages"); + monotonic!(naxum.fair_scheduler.consumer_create_errors = 1); + return ConsumerCreationResult::CreationFailed; + } + }; + + monotonic!(naxum.fair_scheduler.consumers_created = 1); + + if consumer_tx + .send(KeyReady { + key: key.clone(), + messages: key_messages, + }) + .await + .is_err() + { + debug!("consumer_tx closed"); + return ConsumerCreationResult::ChannelClosed; + } + + ConsumerCreationResult::Success +} diff --git a/lib/naxum/src/fair/mod.rs b/lib/naxum/src/fair/mod.rs new file mode 100644 index 0000000000..1a500785a3 --- /dev/null +++ b/lib/naxum/src/fair/mod.rs @@ -0,0 +1,18 @@ +// Fair scheduling for consuming from multiple NATS JetStream consumers. + +mod config; +mod listener; +mod scheduler; +mod stream; + +pub use config::FairSchedulingConfig; +pub use listener::{ + TaskListenerError, + spawn_task_listener, +}; +pub use scheduler::Scheduler; +pub use stream::{ + FairSchedulerError, + FairSchedulerStream, + KeyReady, +}; diff --git a/lib/naxum/src/fair/scheduler.rs b/lib/naxum/src/fair/scheduler.rs new file mode 100644 index 0000000000..b9976bd3cb --- /dev/null +++ b/lib/naxum/src/fair/scheduler.rs @@ -0,0 +1,443 @@ +use std::{ + collections::{ + HashMap, + HashSet, + }, + hash::Hash, +}; + +use tracing::{ + trace, + warn, +}; + +/// Fair scheduler that selects keys with lowest message count. +/// +/// This implements fair queuing (not weighted) where each key gets +/// equal opportunity to process messages. The key with the fewest +/// processed messages is always selected next, ensuring balanced +/// distribution across all active keys. +#[derive(Debug)] +pub struct Scheduler { + message_counts: HashMap, + keys_with_work: HashSet, + max_processed: u64, + stale_threshold: u64, +} + +impl Default for Scheduler +where + K: Clone + Eq + Hash, +{ + fn default() -> Self { + Self::new() + } +} + +impl Scheduler +where + K: Clone + Eq + Hash, +{ + pub fn new() -> Self { + Self { + message_counts: HashMap::new(), + keys_with_work: HashSet::new(), + max_processed: 0, + stale_threshold: 100, + } + } + + pub fn with_stale_threshold(mut self, threshold: u64) -> Self { + self.stale_threshold = threshold; + self + } + + /// Marks key as having work. Resets message count if stale. + pub fn notify_has_work(&mut self, key: K) { + let count = self + .message_counts + .entry(key.clone()) + .or_insert(self.max_processed); + + // Reset count if this key is too far behind (stale) + if self.max_processed.saturating_sub(*count) > self.stale_threshold { + *count = self.max_processed; + } + + self.keys_with_work.insert(key); + } + + pub fn notify_no_work(&mut self, key: &K) { + self.keys_with_work.remove(key); + } + + pub fn keys_with_work_count(&self) -> usize { + self.keys_with_work.len() + } + + pub fn has_work(&self) -> bool { + !self.keys_with_work.is_empty() + } + + /// Returns key with lowest message count (for fairness). + pub fn next(&self) -> Option { + self.keys_with_work + .iter() + .min_by_key(|key| self.message_counts.get(*key).copied().unwrap_or(0)) + .cloned() + } + + /// Increments message count after processing a message. + pub fn complete(&mut self, key: &K, messages_processed: u64) { + if let Some(count) = self.message_counts.get_mut(key) { + *count = count.saturating_add(messages_processed); + self.max_processed = self.max_processed.max(*count); + trace!( + key_count = *count, + max_processed = self.max_processed, + "message count updated" + ); + } else { + // This indicates a programming error: complete() was called for a key + // that was never registered via notify_has_work(). + warn!("complete() called for unknown key - this may indicate a bug"); + debug_assert!( + false, + "complete() called for key not in message_counts - was notify_has_work() called first?" + ); + } + } + + pub fn remove(&mut self, key: &K) { + self.message_counts.remove(key); + self.keys_with_work.remove(key); + } + + pub fn iter_keys_with_work(&self) -> impl Iterator { + self.keys_with_work.iter() + } + + pub fn message_count(&self, key: &K) -> Option { + self.message_counts.get(key).copied() + } + + pub fn max_processed(&self) -> u64 { + self.max_processed + } +} + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + + use super::super::Scheduler; + + // Tests for Scheduler + #[test] + fn test_basic_scheduling() { + let mut scheduler: Scheduler<&str> = Scheduler::new(); + scheduler.notify_has_work("a"); + scheduler.notify_has_work("b"); + scheduler.notify_has_work("c"); + assert!(scheduler.has_work()); + assert_eq!(scheduler.keys_with_work_count(), 3); + } + + #[test] + fn test_fairness_with_multiple_keys() { + let mut scheduler: Scheduler<&str> = Scheduler::new(); + scheduler.notify_has_work("a"); + scheduler.notify_has_work("b"); + scheduler.notify_has_work("c"); + + let mut counts = HashMap::new(); + counts.insert("a", 0); + counts.insert("b", 0); + counts.insert("c", 0); + + // Process 30 messages + for _ in 0..30 { + let next = scheduler.next().unwrap(); + *counts.get_mut(next).unwrap() += 1; + scheduler.complete(&next, 1); + // Keep all keys active + scheduler.notify_has_work("a"); + scheduler.notify_has_work("b"); + scheduler.notify_has_work("c"); + } + + // Each key should get exactly equal processing + assert_eq!(counts["a"], 10); + assert_eq!(counts["b"], 10); + assert_eq!(counts["c"], 10); + } + + #[test] + fn test_scheduler_selects_lowest_message_count() { + let mut scheduler: Scheduler<&str> = Scheduler::new(); + + // Add all keys first + scheduler.notify_has_work("a"); + scheduler.notify_has_work("b"); + scheduler.notify_has_work("c"); + + // Now set up different message counts by completing work + scheduler.complete(&"a", 10); + scheduler.complete(&"b", 5); + // c remains at 0 + + // Check message counts + assert_eq!(scheduler.message_count(&"a"), Some(10)); + assert_eq!(scheduler.message_count(&"b"), Some(5)); + assert_eq!(scheduler.message_count(&"c"), Some(0)); + + // Should select c (lowest count = 0) + assert_eq!(scheduler.next(), Some("c")); + scheduler.complete(&"c", 6); + + // Now c has 6, b has 5, a has 10 - should select b + assert_eq!(scheduler.next(), Some("b")); + scheduler.complete(&"b", 2); + + // Now b has 7, c has 6, a has 10 - should select c + assert_eq!(scheduler.next(), Some("c")); + } + + #[test] + fn test_message_count_cleanup_on_remove() { + let mut scheduler: Scheduler<&str> = Scheduler::new(); + + // Add key and build up message count + scheduler.notify_has_work("workspace1"); + scheduler.complete(&"workspace1", 100); + assert_eq!(scheduler.message_count(&"workspace1"), Some(100)); + + // Remove key + scheduler.remove(&"workspace1"); + + // Verify message count is gone + assert_eq!(scheduler.message_count(&"workspace1"), None); + assert!(!scheduler.has_work()); + + // Re-add key, verify starts at max_processed + scheduler.notify_has_work("workspace1"); + assert_eq!(scheduler.message_count(&"workspace1"), Some(100)); // max_processed is 100 + } + + #[test] + fn test_scheduler_state_consistency() { + let mut scheduler: Scheduler<&str> = Scheduler::new(); + + // Add keys + scheduler.notify_has_work("a"); + scheduler.notify_has_work("b"); + + // Both should be in keys_with_work + assert_eq!(scheduler.keys_with_work_count(), 2); + assert!(scheduler.message_count(&"a").is_some()); + assert!(scheduler.message_count(&"b").is_some()); + + // Remove from work + scheduler.notify_no_work(&"a"); + assert_eq!(scheduler.keys_with_work_count(), 1); + // Message count should still exist + assert!(scheduler.message_count(&"a").is_some()); + + // Remove entirely + scheduler.remove(&"b"); + assert_eq!(scheduler.keys_with_work_count(), 0); + assert!(scheduler.message_count(&"b").is_none()); + } + + #[test] + fn test_rapid_add_remove_same_key() { + let mut scheduler: Scheduler<&str> = Scheduler::new(); + + for i in 0..10 { + scheduler.notify_has_work("key"); + scheduler.complete(&"key", 1); + if i % 2 == 0 { + scheduler.notify_no_work(&"key"); + } + } + + // Should handle rapid changes without issues + assert!(scheduler.message_count(&"key").is_some()); + assert_eq!(scheduler.message_count(&"key"), Some(10)); + } + + #[test] + fn test_scheduler_with_many_keys() { + let mut scheduler: Scheduler = Scheduler::new(); + + // Add 100 keys + for i in 0..100 { + let key = format!("key_{i}"); + scheduler.notify_has_work(key); + } + + assert_eq!(scheduler.keys_with_work_count(), 100); + + // Process some messages + for _ in 0..200 { + let next = scheduler.next().unwrap(); + scheduler.complete(&next, 1); + scheduler.notify_has_work(next); + } + + // Verify all keys still tracked + assert_eq!(scheduler.keys_with_work_count(), 100); + } + + #[test] + fn test_stale_key_reset() { + let mut scheduler: Scheduler<&str> = Scheduler::new().with_stale_threshold(100); + + scheduler.notify_has_work("a"); + scheduler.complete(&"a", 500); + scheduler.notify_no_work(&"a"); + + // Add new key, advances max_processed + scheduler.notify_has_work("b"); + assert_eq!(scheduler.max_processed(), 500); + + // Re-add stale key + scheduler.notify_has_work("a"); + // Should reset to max_processed due to staleness + assert_eq!(scheduler.message_count(&"a"), Some(500)); + } + + #[test] + fn test_no_work() { + let scheduler: Scheduler<&str> = Scheduler::new(); + assert!(!scheduler.has_work()); + assert_eq!(scheduler.next(), None); + } + + #[test] + fn test_remove_key() { + let mut scheduler: Scheduler<&str> = Scheduler::new(); + scheduler.notify_has_work("a"); + scheduler.notify_has_work("b"); + scheduler.complete(&"a", 10); + scheduler.remove(&"a"); + + // Use public API to verify state + assert_eq!(scheduler.keys_with_work_count(), 1); + assert!(scheduler.message_count(&"a").is_none()); + assert_eq!(scheduler.next(), Some("b")); + } + + #[test] + fn test_interleaved_processing() { + let mut scheduler: Scheduler<&str> = Scheduler::new(); + scheduler.notify_has_work("a"); + scheduler.notify_has_work("b"); + + let mut a_count = 0; + let mut b_count = 0; + + for _ in 0..20 { + let next = scheduler.next().unwrap(); + if next == "a" { + a_count += 1; + } else { + b_count += 1; + } + scheduler.complete(&next, 1); + scheduler.notify_has_work("a"); + scheduler.notify_has_work("b"); + } + + assert_eq!(a_count, 10); + assert_eq!(b_count, 10); + } + + #[test] + fn test_consumer_replacement_preserves_message_count() { + let mut scheduler: Scheduler<&str> = Scheduler::new(); + + // Add key and build up message count + scheduler.notify_has_work("ws1"); + scheduler.complete(&"ws1", 50); + assert_eq!(scheduler.message_count(&"ws1"), Some(50)); + + // Simulate consumer replacement by just notifying again (new behavior) + scheduler.notify_has_work("ws1"); + + // Message count should be preserved at 50 + assert_eq!(scheduler.message_count(&"ws1"), Some(50)); + } + + #[test] + fn test_keys_with_work_and_message_counts_consistency() { + let mut scheduler: Scheduler<&str> = Scheduler::new(); + + // Add some keys + scheduler.notify_has_work("a"); + scheduler.notify_has_work("b"); + scheduler.notify_has_work("c"); + + // All keys should have message counts + assert!(scheduler.message_count(&"a").is_some()); + assert!(scheduler.message_count(&"b").is_some()); + assert!(scheduler.message_count(&"c").is_some()); + + // Remove work notification + scheduler.notify_no_work(&"b"); + + // Key should not be in work set but message count persists + assert_eq!(scheduler.keys_with_work_count(), 2); + assert!(scheduler.message_count(&"b").is_some()); + + // Complete removal + scheduler.remove(&"c"); + + // Should be gone from both + assert_eq!(scheduler.keys_with_work_count(), 1); + assert!(scheduler.message_count(&"c").is_none()); + } + + // Tests that verify behavior we'll change in refactoring + #[test] + fn test_replacement_behavior_preserves_fairness() { + // After refactor: replacement preserves message count for fairness + let mut scheduler: Scheduler<&str> = Scheduler::new(); + + // Add two workspaces + scheduler.notify_has_work("ws1"); + scheduler.notify_has_work("ws2"); + + // ws1 processes more work (scheduler should alternate due to fairness) + let mut ws1_count = 0; + for _ in 0..10 { + let next = scheduler.next().unwrap(); + if next == "ws1" { + ws1_count += 1; + } + scheduler.complete(&next, 1); + scheduler.notify_has_work("ws1"); + scheduler.notify_has_work("ws2"); + } + + // Both should have processed equal work due to fairness + assert_eq!(ws1_count, 5); + assert_eq!(scheduler.message_count(&"ws1"), Some(5)); + assert_eq!(scheduler.message_count(&"ws2"), Some(5)); + + // Process one from ws1 to make it have higher message count + assert!(scheduler.next().is_some()); + scheduler.complete(&"ws1", 1); + + // Now ws1 has count 6, ws2 has 5 + assert_eq!(scheduler.message_count(&"ws1"), Some(6)); + assert_eq!(scheduler.message_count(&"ws2"), Some(5)); + + // Simulate ws1 consumer replacement + scheduler.notify_has_work("ws1"); + scheduler.notify_has_work("ws2"); + + // Message count preserved - ws2 should be selected next (lower count) + assert_eq!(scheduler.next(), Some("ws2")); + } +} diff --git a/lib/naxum/src/fair/stream.rs b/lib/naxum/src/fair/stream.rs new file mode 100644 index 0000000000..80dbf14c6f --- /dev/null +++ b/lib/naxum/src/fair/stream.rs @@ -0,0 +1,249 @@ +// Fair scheduler stream for balanced message processing across keys. + +use std::{ + collections::HashMap, + hash::Hash, + pin::Pin, + task::{ + Context, + Poll, + }, +}; + +use async_nats::jetstream::{ + Message, + consumer::pull::Stream as MessageStream, +}; +use futures::Stream; +use telemetry_utils::{ + counter, + monotonic, +}; +use thiserror::Error; +use tokio::sync::mpsc; +use tokio_util::sync::CancellationToken; +use tracing::{ + info, + trace, +}; + +use super::scheduler::Scheduler; + +/// Notification that a key has work available. +pub struct KeyReady { + pub key: K, + pub messages: MessageStream, +} + +#[derive(Debug, Error)] +pub enum FairSchedulerError { + #[error("consumer stream error: {0}")] + ConsumerStream(#[source] async_nats::jetstream::consumer::pull::MessagesError), +} + +pub struct FairSchedulerStream { + scheduler: Scheduler, + consumers: HashMap>>, + consumer_rx: mpsc::Receiver>, + shutdown: CancellationToken, +} + +impl FairSchedulerStream +where + K: Clone + Eq + Hash, +{ + pub fn new(consumer_rx: mpsc::Receiver>, shutdown: CancellationToken) -> Self { + Self { + scheduler: Scheduler::new(), + consumers: HashMap::new(), + consumer_rx, + shutdown, + } + } + + fn drain_new_consumers(&mut self) { + while let Ok(ready) = self.consumer_rx.try_recv() { + trace!("received key consumer for draining"); + self.insert_consumer(ready.key, ready.messages); + } + } + + fn insert_consumer(&mut self, key: K, messages: MessageStream) { + if let Some(old_consumer) = self.consumers.insert(key.clone(), Box::pin(messages)) { + // Replacing an existing consumer - clean up the old one + trace!("replacing existing consumer, notifying has work"); + drop(old_consumer); + self.scheduler.notify_has_work(key.clone()); + } else { + // New consumer - increment the metric and notify scheduler + trace!("inserting new consumer, notifying has work"); + counter!(naxum.fair_scheduler.active_keys = 1); + self.scheduler.notify_has_work(key); + } + } + + fn cleanup(&mut self, key: K) { + self.scheduler.notify_no_work(&key); + self.consumers.remove(&key); + counter!(naxum.fair_scheduler.active_keys = -1); + monotonic!(naxum.fair_scheduler.key_empty = 1); + } +} + +impl Stream for FairSchedulerStream +where + K: Clone + Eq + Hash + Unpin, +{ + type Item = Result; + + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let this = self.get_mut(); + + if this.shutdown.is_cancelled() { + return Poll::Ready(None); + } + + loop { + // Drain all available consumers from the channel + this.drain_new_consumers(); + + if let Some(key) = this.scheduler.next() { + let Some(consumer) = this.consumers.get_mut(&key) else { + this.scheduler.notify_no_work(&key); + continue; + }; + + trace!("scheduler selected key, polling consumer"); + match consumer.as_mut().poll_next(cx) { + Poll::Ready(Some(Ok(msg))) => { + // Got message - update state and return + this.scheduler.complete(&key, 1); + monotonic!(naxum.fair_scheduler.messages_processed = 1); + return Poll::Ready(Some(Ok(msg))); + } + Poll::Ready(Some(Err(e))) => { + return Poll::Ready(Some(Err(FairSchedulerError::ConsumerStream(e)))); + } + Poll::Ready(None) => { + // Consumer exhausted - clean up and try next + trace!("consumer exhausted, marking no work"); + this.cleanup(key); + continue; + } + Poll::Pending => { + // Consumer pending but not idle yet - wake ourselves to retry + // This ensures we keep polling even if the JetStream consumer + // doesn't properly propagate wakeups + cx.waker().wake_by_ref(); + return Poll::Pending; + } + } + } else { + // No keys with work - wait for new consumers on the channel + trace!("no keys with work, waiting for consumer notifications"); + match Pin::new(&mut this.consumer_rx).poll_recv(cx) { + Poll::Ready(Some(ready)) => { + trace!("received new consumer notification"); + this.insert_consumer(ready.key, ready.messages); + // Loop to try polling the new consumer + continue; + } + Poll::Ready(None) => { + trace!("consumer channel closed, terminating stream"); + return Poll::Ready(None); + } + Poll::Pending => { + trace!("no new consumers, returning pending"); + return Poll::Pending; + } + } + } + } + } +} + +#[cfg(test)] +mod tests { + use std::time::Duration; + + use futures::StreamExt; + use tokio::sync::mpsc; + use tokio_util::sync::CancellationToken; + + use super::super::FairSchedulerStream; + + #[tokio::test] + async fn test_fair_stream_shutdown_cancels_immediately() { + let (_consumer_tx, consumer_rx) = mpsc::channel(10); + let shutdown = CancellationToken::new(); + + let mut stream = FairSchedulerStream::::new(consumer_rx, shutdown.clone()); + + // Cancel immediately + shutdown.cancel(); + + // Stream should return None (terminated) + let result = stream.next().await; + assert!(result.is_none(), "Stream should terminate on shutdown"); + } + + #[tokio::test] + async fn test_fair_stream_handles_no_consumers() { + let (_consumer_tx, consumer_rx) = mpsc::channel(10); + let shutdown = CancellationToken::new(); + + let mut stream = FairSchedulerStream::::new(consumer_rx, shutdown.clone()); + + // Try to poll with no consumers (should be pending) + tokio::select! { + _ = stream.next() => { + panic!("Stream should not yield without consumers"); + } + _ = tokio::time::sleep(Duration::from_millis(100)) => { + // Expected: stream is waiting for consumers + } + } + + shutdown.cancel(); + } + + #[tokio::test] + async fn test_fair_stream_returns_none_when_consumer_channel_closes() { + let (consumer_tx, consumer_rx) = mpsc::channel(10); + let shutdown = CancellationToken::new(); + + let mut stream = FairSchedulerStream::::new(consumer_rx, shutdown.clone()); + + // Close the consumer channel (no consumers will ever arrive) + drop(consumer_tx); + + // Stream should return None when channel is closed + let result = stream.next().await; + assert!( + result.is_none(), + "Stream should terminate when consumer channel closes" + ); + } + + #[tokio::test] + async fn test_fair_stream_shutdown_vs_channel_close() { + // Test that shutdown takes precedence over channel operations + let (consumer_tx, consumer_rx) = mpsc::channel(10); + let shutdown = CancellationToken::new(); + + let mut stream = FairSchedulerStream::::new(consumer_rx, shutdown.clone()); + + // Cancel shutdown before closing channel + shutdown.cancel(); + + // Even though channel is still open, shutdown should terminate stream + let result = stream.next().await; + assert!( + result.is_none(), + "Shutdown should terminate stream immediately" + ); + + // Clean up + drop(consumer_tx); + } +} diff --git a/lib/naxum/src/lib.rs b/lib/naxum/src/lib.rs index 2041dcdd4f..5700085dab 100644 --- a/lib/naxum/src/lib.rs +++ b/lib/naxum/src/lib.rs @@ -6,6 +6,7 @@ mod cancellation; mod error; pub mod error_handling; pub mod extract; +pub mod fair; pub mod handler; mod json; mod make_service; @@ -13,6 +14,7 @@ mod message; pub mod middleware; pub mod response; pub mod serve; +mod serve_builder; mod service_ext; pub use async_nats::StatusCode; @@ -36,10 +38,12 @@ pub use self::{ MessageHead, }, serve::{ + SemaphoreMode, serve, serve_with_external_semaphore, serve_with_incoming_limit, }, + serve_builder::ServeBuilder, service_ext::ServiceExt, }; diff --git a/lib/naxum/src/serve.rs b/lib/naxum/src/serve.rs index 345f2a216f..ad036b2d4e 100644 --- a/lib/naxum/src/serve.rs +++ b/lib/naxum/src/serve.rs @@ -52,7 +52,7 @@ use crate::{ }; #[derive(Debug, Clone)] -enum SemaphoreMode { +pub enum SemaphoreMode { /// Permits are released when the request completes (standard behavior). Internal(Arc), /// Permits are forgotten after acquisition; owner restores via `add_permits()`. @@ -72,6 +72,8 @@ where stream, make_service, semaphore: None, + task_listener_handle: None, + shutdown_token: None, _service_marker: PhantomData, _stream_error_marker: PhantomData, _request_marker: PhantomData, @@ -97,6 +99,8 @@ where semaphore: limit .into() .map(|l| SemaphoreMode::Internal(Arc::new(Semaphore::new(l)))), + task_listener_handle: None, + shutdown_token: None, _service_marker: PhantomData, _stream_error_marker: PhantomData, _request_marker: PhantomData, @@ -122,6 +126,8 @@ where stream, make_service, semaphore: Some(SemaphoreMode::External(semaphore)), + task_listener_handle: None, + shutdown_token: None, _service_marker: PhantomData, _stream_error_marker: PhantomData, _request_marker: PhantomData, @@ -130,12 +136,14 @@ where #[must_use = "futures must be awaited or polled"] pub struct Serve { - stream: T, - make_service: M, - semaphore: Option, - _service_marker: PhantomData, - _stream_error_marker: PhantomData, - _request_marker: PhantomData, + pub(crate) stream: T, + pub(crate) make_service: M, + pub(crate) semaphore: Option, + pub(crate) task_listener_handle: Option>, + pub(crate) shutdown_token: Option, + pub(crate) _service_marker: PhantomData, + pub(crate) _stream_error_marker: PhantomData, + pub(crate) _request_marker: PhantomData, } impl fmt::Debug for Serve @@ -158,6 +166,8 @@ impl Serve { stream: self.stream, make_service: self.make_service, semaphore: self.semaphore, + task_listener_handle: self.task_listener_handle, + shutdown_token: self.shutdown_token, signal, _service_marker: PhantomData, _stream_error_marker: PhantomData, @@ -172,6 +182,8 @@ pub struct WithGracefulShutdown { stream: T, make_service: M, semaphore: Option, + task_listener_handle: Option>, + shutdown_token: Option, signal: F, _service_marker: PhantomData, _stream_error_marker: PhantomData, @@ -211,6 +223,8 @@ where stream, mut make_service, semaphore, + task_listener_handle, + shutdown_token, signal, .. } = self; @@ -308,6 +322,14 @@ where } } + // If using fair scheduling, cancel the shutdown token and wait for task listener + if let Some(token) = shutdown_token { + token.cancel(); + if let Some(handle) = task_listener_handle { + let _ = handle.await; + } + } + Ok(()) })) } diff --git a/lib/naxum/src/serve_builder.rs b/lib/naxum/src/serve_builder.rs new file mode 100644 index 0000000000..9b788970d1 --- /dev/null +++ b/lib/naxum/src/serve_builder.rs @@ -0,0 +1,207 @@ +use std::{ + convert::Infallible, + error, + fmt, + hash::Hash, + marker::PhantomData, + sync::Arc, +}; + +use async_nats::jetstream::Message as JsMessage; +use futures::Stream; +use tokio::sync::{ + Semaphore, + mpsc::{ + self, + Receiver, + Sender, + }, +}; +use tokio_util::sync::CancellationToken; + +use crate::{ + fair::{ + FairSchedulerError, + FairSchedulerStream, + FairSchedulingConfig, + KeyReady, + spawn_task_listener, + }, + message::{ + Message, + MessageHead, + }, + response::Response, + serve::{ + IncomingMessage, + SemaphoreMode, + Serve, + }, +}; + +pub struct NoFairScheduling; +pub struct WithFairScheduling { + config: FairSchedulingConfig, +} + +/// Builder for configuring the naxum serve loop. +/// +/// Uses type-state pattern to ensure correct configuration at compile time. +/// +/// # Example +/// +/// ```ignore +/// // Simple case - no fair scheduling +/// ServeBuilder::new(app) +/// .with_incoming_limit(100) +/// .serve(stream) +/// .with_graceful_shutdown(shutdown_signal) +/// .await?; +/// +/// // With fair scheduling +/// ServeBuilder::new(app) +/// .with_external_semaphore(semaphore) +/// .with_fair_scheduling(config) +/// .serve() +/// .with_graceful_shutdown(shutdown_signal) +/// .await?; +/// ``` +pub struct ServeBuilder { + make_service: M, + semaphore: Option, + fair_scheduling: FS, + _marker: PhantomData, +} + +impl fmt::Debug for ServeBuilder +where + M: fmt::Debug, +{ + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("ServeBuilder") + .field("make_service", &self.make_service) + .finish_non_exhaustive() + } +} + +impl ServeBuilder { + /// Creates a new serve builder with the given service. + pub fn new(make_service: M) -> Self { + Self { + make_service, + semaphore: None, + fair_scheduling: NoFairScheduling, + _marker: PhantomData, + } + } + + /// Sets a limit on the number of concurrent incoming messages. + /// + /// This creates an internal semaphore that limits concurrency. + /// Permits are released when message processing completes. + pub fn with_incoming_limit(mut self, limit: usize) -> Self { + self.semaphore = Some(SemaphoreMode::Internal(Arc::new(Semaphore::new(limit)))); + self + } + + /// Uses an external semaphore for concurrency control. + /// + /// Permits are "forgotten" after acquisition - the semaphore owner + /// is responsible for restoring permits via `add_permits()`. + pub fn with_external_semaphore(mut self, semaphore: Arc) -> Self { + self.semaphore = Some(SemaphoreMode::External(semaphore)); + self + } + + /// Serves messages from the given stream. + /// + /// Use this for the standard case without fair scheduling. + pub fn serve(self, stream: T) -> Serve + where + M: for<'a> tower::Service, Error = Infallible, Response = S>, + S: tower::Service, Response = Response, Error = Infallible> + + Clone + + Send + + 'static, + S::Future: Send, + T: Stream>, + E: error::Error, + R: MessageHead, + { + Serve { + stream, + make_service: self.make_service, + semaphore: self.semaphore, + task_listener_handle: None, + shutdown_token: None, + _service_marker: PhantomData, + _stream_error_marker: PhantomData, + _request_marker: PhantomData, + } + } + + /// Configures fair scheduling. + /// + /// Fair scheduling ensures that no single key (e.g., workspace) can + /// monopolize processing capacity. + /// + /// An internal shutdown token will be created to coordinate between the + /// fair scheduler stream and task listener. The token will be cancelled + /// automatically during graceful shutdown. + pub fn with_fair_scheduling( + self, + config: FairSchedulingConfig, + ) -> ServeBuilder> + where + K: Clone + Eq + Hash + Send + Sync + 'static, + { + ServeBuilder { + make_service: self.make_service, + semaphore: self.semaphore, + fair_scheduling: WithFairScheduling { config }, + _marker: PhantomData, + } + } +} + +impl ServeBuilder> +where + K: Clone + Eq + Hash + Send + Sync + Unpin + 'static, +{ + /// Creates the fair scheduler and returns a Serve instance. + /// + /// This spawns the task listener in the background and creates the fair scheduler stream. + /// An internal shutdown token is created to coordinate between components and will be + /// cancelled automatically during graceful shutdown. + pub fn serve(self) -> Serve, FairSchedulerError, JsMessage> + where + M: for<'a> tower::Service, Error = Infallible, Response = S>, + S: tower::Service, Response = Response, Error = Infallible> + + Clone + + Send + + 'static, + S::Future: Send, + { + let WithFairScheduling { config } = self.fair_scheduling; + + let shutdown = CancellationToken::new(); + + // Channel for task listener to send key consumers to fair scheduler + let (consumer_tx, consumer_rx): (Sender>, Receiver>) = + mpsc::channel(64); + + let task_listener_handle = spawn_task_listener(config, consumer_tx, shutdown.clone()); + let stream = FairSchedulerStream::new(consumer_rx, shutdown.clone()); + + Serve { + stream, + make_service: self.make_service, + semaphore: self.semaphore, + task_listener_handle: Some(task_listener_handle), + shutdown_token: Some(shutdown), + _service_marker: PhantomData, + _stream_error_marker: PhantomData, + _request_marker: PhantomData, + } + } +} diff --git a/lib/naxum/tests/integration/fair.rs b/lib/naxum/tests/integration/fair.rs new file mode 100644 index 0000000000..fde3c3def3 --- /dev/null +++ b/lib/naxum/tests/integration/fair.rs @@ -0,0 +1,172 @@ +use std::time::Duration; + +use naxum::fair::{ + FairSchedulingConfig, + spawn_task_listener, +}; +use si_data_nats::jetstream; +use tokio::sync::mpsc; +use tokio_util::sync::CancellationToken; + +use crate::helpers::{ + SERVICE_NAME, + create_test_streams, + nats, + nats_prefix, +}; + +#[tokio::test] +async fn stress_test_high_volume_single_workspace() { + // Stress test: 100 messages in rapid succession for a single workspace + let prefix = nats_prefix(); + let nats_client = nats(prefix.clone()).await; + let (tasks_stream, requests_stream) = create_test_streams(&prefix, &nats_client).await; + + let config = FairSchedulingConfig::for_workspace_partitioning( + SERVICE_NAME, + tasks_stream, + requests_stream.clone(), + Some(prefix.clone()), + ); + + let (consumer_tx, mut consumer_rx) = mpsc::channel(10); + let shutdown = CancellationToken::new(); + + let _listener_handle = spawn_task_listener(config, consumer_tx, shutdown.clone()); + + let js_context = jetstream::new(nats_client.clone()); + + // Publish task notification + js_context + .publish( + format!("{prefix}.{SERVICE_NAME}.tasks.workspace-1"), + "notification".into(), + ) + .await + .expect("failed to publish"); + + // Wait for consumer creation + let ready = tokio::time::timeout(Duration::from_secs(5), consumer_rx.recv()) + .await + .expect("timeout") + .expect("consumer channel closed"); + + let mut stream = ready.messages; + + // Publish 100 messages as fast as possible + for i in 0..100 { + let subject = format!("{prefix}.{SERVICE_NAME}.requests.workspace-1.test"); + js_context + .publish(subject, format!("message-{i}").into()) + .await + .expect("failed to publish") + .await + .expect("failed to ack publish"); + } + + // Consume all 100 messages + use futures::StreamExt; + for i in 0..100 { + let msg = tokio::time::timeout(Duration::from_secs(10), stream.next()) + .await + .unwrap_or_else(|_| panic!("timeout waiting for message {i}")) + .expect("stream ended") + .expect("stream error"); + + msg.ack().await.expect("failed to ack"); + } + + shutdown.cancel(); +} + +#[tokio::test] +async fn stress_test_many_workspaces_concurrent() { + // Stress test: 20 workspaces, 10 messages each, sent concurrently + let prefix = nats_prefix(); + let nats_client = nats(prefix.clone()).await; + let (tasks_stream, requests_stream) = create_test_streams(&prefix, &nats_client).await; + + let config = FairSchedulingConfig::for_workspace_partitioning( + SERVICE_NAME, + tasks_stream, + requests_stream.clone(), + Some(prefix.clone()), + ); + + let (consumer_tx, mut consumer_rx) = mpsc::channel(50); + let shutdown = CancellationToken::new(); + + let _listener_handle = spawn_task_listener(config, consumer_tx, shutdown.clone()); + + let js_context = jetstream::new(nats_client.clone()); + + // Publish task notifications for 20 workspaces + let workspace_count = 20; + for i in 0..workspace_count { + js_context + .publish( + format!("{prefix}.{SERVICE_NAME}.tasks.workspace-{i}"), + "notification".into(), + ) + .await + .expect("failed to publish"); + } + + // Collect all consumers + let mut consumers = Vec::new(); + for _ in 0..workspace_count { + let ready = tokio::time::timeout(Duration::from_secs(10), consumer_rx.recv()) + .await + .expect("timeout") + .expect("consumer channel closed"); + consumers.push(ready); + } + + // Publish 10 messages to each workspace concurrently + let mut publish_handles = vec![]; + for i in 0..workspace_count { + let js = js_context.clone(); + let p = prefix.clone(); + let handle = tokio::spawn(async move { + for j in 0..10 { + let subject = format!("{p}.{SERVICE_NAME}.requests.workspace-{i}.test"); + js.publish(subject, format!("ws-{i}-msg-{j}").into()) + .await + .expect("failed to publish") + .await + .expect("failed to ack publish"); + } + }); + publish_handles.push(handle); + } + + // Wait for all publishes to complete + for handle in publish_handles { + handle.await.expect("publish task failed"); + } + + // Consume all messages from all workspaces concurrently + use futures::StreamExt; + let mut consume_handles = vec![]; + for ready in consumers { + let mut stream = ready.messages; + let handle = tokio::spawn(async move { + for _ in 0..10 { + let msg = tokio::time::timeout(Duration::from_secs(10), stream.next()) + .await + .expect("timeout") + .expect("stream ended") + .expect("stream error"); + msg.ack().await.expect("failed to ack"); + } + }); + consume_handles.push(handle); + } + + // Wait for all consumers to finish + for handle in consume_handles { + handle.await.expect("consume task failed"); + } + + shutdown.cancel(); +} diff --git a/lib/naxum/tests/integration/helpers/mod.rs b/lib/naxum/tests/integration/helpers/mod.rs new file mode 100644 index 0000000000..8bf06905ea --- /dev/null +++ b/lib/naxum/tests/integration/helpers/mod.rs @@ -0,0 +1,68 @@ +use std::env; + +use si_data_nats::{ + NatsClient, + NatsConfig, + async_nats::jetstream::stream::{ + Config as StreamConfig, + Stream, + }, + jetstream, +}; +use uuid::Uuid; + +pub const SERVICE_NAME: &str = "test-service"; + +pub fn nats_config(subject_prefix: String) -> NatsConfig { + let mut config = NatsConfig::default(); + #[allow(clippy::disallowed_methods)] + if let Ok(value) = env::var("SI_TEST_NATS_URL") { + config.url = value; + } + config.subject_prefix = Some(subject_prefix); + config +} + +pub async fn nats(subject_prefix: String) -> NatsClient { + NatsClient::new(&nats_config(subject_prefix)) + .await + .expect("failed to connect to NATS") +} + +pub fn nats_prefix() -> String { + #[allow(clippy::disallowed_methods)] + let test_name = env::var("NEXTEST_TEST_NAME").unwrap_or_else(|_| { + std::thread::current() + .name() + .unwrap_or("unknown") + .to_string() + }); + format!("test-{}-{}", Uuid::new_v4(), test_name) +} + +pub async fn create_test_streams(prefix: &str, nats_client: &NatsClient) -> (Stream, Stream) { + let js_context = jetstream::new(nats_client.clone()); + + let tasks_stream_name = format!("{prefix}-tasks-test"); + let requests_stream_name = format!("{prefix}-requests-test"); + + let tasks_stream = js_context + .get_or_create_stream(StreamConfig { + name: tasks_stream_name.clone(), + subjects: vec![format!("{}.{}.tasks.*", prefix, SERVICE_NAME)], + ..Default::default() + }) + .await + .expect("failed to create tasks stream"); + + let requests_stream = js_context + .get_or_create_stream(StreamConfig { + name: requests_stream_name.clone(), + subjects: vec![format!("{}.{}.requests.>", prefix, SERVICE_NAME)], + ..Default::default() + }) + .await + .expect("failed to create requests stream"); + + (tasks_stream, requests_stream) +} diff --git a/lib/naxum/tests/integration/mod.rs b/lib/naxum/tests/integration/mod.rs new file mode 100644 index 0000000000..b3433cec5e --- /dev/null +++ b/lib/naxum/tests/integration/mod.rs @@ -0,0 +1,552 @@ +#![allow(clippy::disallowed_methods)] + +mod fair; +mod helpers; + +use std::time::Duration; + +use helpers::{ + SERVICE_NAME, + create_test_streams, + nats, + nats_prefix, +}; +use naxum::fair::{ + FairSchedulingConfig, + spawn_task_listener, +}; +use si_data_nats::{ + async_nats::jetstream::{ + consumer::PullConsumer, + stream::Stream, + }, + jetstream, +}; +use test_log::test; +use tokio::sync::mpsc; +use tokio_util::sync::CancellationToken; + +async fn get_consumer_with_retry(stream: &Stream, consumer_name: &str) -> PullConsumer { + let mut attempts = 0; + loop { + match stream.get_consumer(consumer_name).await { + Ok(c) => break c, + Err(_) if attempts < 20 => { + attempts += 1; + tokio::time::sleep(Duration::from_millis(100)).await; + } + Err(e) => { + panic!("failed to get consumer '{consumer_name}' after {attempts} attempts: {e}") + } + } + } +} + +#[test(tokio::test)] +async fn listener_creates_consumers_on_task_notifications() { + let prefix = nats_prefix(); + let nats_client = nats(prefix.clone()).await; + let (tasks_stream, requests_stream) = create_test_streams(&prefix, &nats_client).await; + + let config = FairSchedulingConfig::for_workspace_partitioning( + SERVICE_NAME, + tasks_stream.clone(), + requests_stream.clone(), + Some(prefix.clone()), + ); + + let (consumer_tx, mut consumer_rx) = mpsc::channel(10); + let shutdown = CancellationToken::new(); + + let _listener_handle = spawn_task_listener(config, consumer_tx, shutdown.clone()); + + // Publish task notification for workspace-1 + let js_context = jetstream::new(nats_client.clone()); + js_context + .publish( + format!("{prefix}.{SERVICE_NAME}.tasks.workspace-1"), + "task notification".into(), + ) + .await + .expect("failed to publish task notification"); + + // Wait for consumer to be created and sent + let ready = tokio::time::timeout(Duration::from_secs(5), consumer_rx.recv()) + .await + .expect("timeout waiting for consumer") + .expect("consumer channel closed"); + + assert_eq!(ready.key, "workspace-1"); + + // Verify consumer was created in NATS + let consumer_name = format!("{SERVICE_NAME}-ws-workspace-1"); + let _consumer = get_consumer_with_retry(&requests_stream, &consumer_name).await; + + shutdown.cancel(); +} + +#[test(tokio::test)] +async fn listener_handles_multiple_workspaces() { + let prefix = nats_prefix(); + let nats_client = nats(prefix.clone()).await; + let (tasks_stream, requests_stream) = create_test_streams(&prefix, &nats_client).await; + + let config = FairSchedulingConfig::for_workspace_partitioning( + SERVICE_NAME, + tasks_stream.clone(), + requests_stream.clone(), + Some(prefix.clone()), + ); + + let (consumer_tx, mut consumer_rx) = mpsc::channel(10); + let shutdown = CancellationToken::new(); + + let _listener_handle = spawn_task_listener(config, consumer_tx, shutdown.clone()); + + let js_context = jetstream::new(nats_client.clone()); + + // Publish notifications for multiple workspaces + for workspace in ["ws-a", "ws-b", "ws-c"] { + js_context + .publish( + format!("{prefix}.{SERVICE_NAME}.tasks.{workspace}"), + "task notification".into(), + ) + .await + .expect("failed to publish task notification"); + } + + // Collect all consumers + let mut received_keys = Vec::new(); + for _ in 0..3 { + let ready = tokio::time::timeout(Duration::from_secs(5), consumer_rx.recv()) + .await + .expect("timeout waiting for consumer") + .expect("consumer channel closed"); + received_keys.push(ready.key); + } + + received_keys.sort(); + assert_eq!(received_keys, vec!["ws-a", "ws-b", "ws-c"]); + + shutdown.cancel(); +} + +#[test(tokio::test)] +async fn listener_acks_message_after_successful_consumer_creation() { + let prefix = nats_prefix(); + let nats_client = nats(prefix.clone()).await; + let (tasks_stream, requests_stream) = create_test_streams(&prefix, &nats_client).await; + + let config = FairSchedulingConfig::for_workspace_partitioning( + SERVICE_NAME, + tasks_stream.clone(), + requests_stream.clone(), + Some(prefix.clone()), + ); + + let (consumer_tx, mut consumer_rx) = mpsc::channel(10); + let shutdown = CancellationToken::new(); + + let _listener_handle = spawn_task_listener(config, consumer_tx, shutdown.clone()); + + let js_context = jetstream::new(nats_client.clone()); + js_context + .publish( + format!("{prefix}.{SERVICE_NAME}.tasks.workspace-1"), + "task notification".into(), + ) + .await + .expect("failed to publish"); + + // Wait for consumer + tokio::time::timeout(Duration::from_secs(5), consumer_rx.recv()) + .await + .expect("timeout") + .expect("consumer channel closed"); + + // Give time for ack to process + tokio::time::sleep(Duration::from_millis(500)).await; + + // Verify message was acknowledged (no pending messages) + let mut task_listener_consumer = + get_consumer_with_retry(&tasks_stream, &format!("{SERVICE_NAME}-task-listener")).await; + + let info = task_listener_consumer + .info() + .await + .expect("failed to get consumer info"); + + assert_eq!(info.num_pending, 0, "message should be acknowledged"); + + shutdown.cancel(); +} + +#[test(tokio::test)] +async fn listener_shuts_down_gracefully() { + let prefix = nats_prefix(); + let nats_client = nats(prefix.clone()).await; + let (tasks_stream, requests_stream) = create_test_streams(&prefix, &nats_client).await; + + let config = FairSchedulingConfig::for_workspace_partitioning( + SERVICE_NAME, + tasks_stream, + requests_stream, + Some(prefix.clone()), + ); + + let (consumer_tx, _consumer_rx) = mpsc::channel(10); + let shutdown = CancellationToken::new(); + + let listener_handle = spawn_task_listener(config, consumer_tx, shutdown.clone()); + + tokio::time::sleep(Duration::from_millis(100)).await; + + shutdown.cancel(); + + tokio::time::timeout(Duration::from_secs(5), listener_handle) + .await + .expect("listener should shut down within timeout") + .expect("listener task should complete without error"); +} + +#[test(tokio::test)] +async fn listener_stops_when_consumer_channel_closed() { + let prefix = nats_prefix(); + let nats_client = nats(prefix.clone()).await; + let (tasks_stream, requests_stream) = create_test_streams(&prefix, &nats_client).await; + + let config = FairSchedulingConfig::for_workspace_partitioning( + SERVICE_NAME, + tasks_stream, + requests_stream, + Some(prefix.clone()), + ); + + let (consumer_tx, consumer_rx) = mpsc::channel(10); + let shutdown = CancellationToken::new(); + + let listener_handle = spawn_task_listener(config, consumer_tx, shutdown.clone()); + + tokio::time::sleep(Duration::from_millis(100)).await; + + // Drop the receiver to close the channel + drop(consumer_rx); + + let js_context = jetstream::new(nats_client.clone()); + js_context + .publish( + format!("{prefix}.{SERVICE_NAME}.tasks.workspace-1"), + "task notification".into(), + ) + .await + .expect("failed to publish"); + + // Listener should stop when it tries to send to closed channel + tokio::time::timeout(Duration::from_secs(5), listener_handle) + .await + .expect("listener should stop within timeout") + .expect("listener task should complete without error"); + + shutdown.cancel(); +} + +#[test(tokio::test)] +async fn listener_replaces_consumer_on_duplicate_notification() { + let prefix = nats_prefix(); + let nats_client = nats(prefix.clone()).await; + let (tasks_stream, requests_stream) = create_test_streams(&prefix, &nats_client).await; + + let config = FairSchedulingConfig::for_workspace_partitioning( + SERVICE_NAME, + tasks_stream, + requests_stream, + Some(prefix.clone()), + ); + + let (consumer_tx, mut consumer_rx) = mpsc::channel(10); + let shutdown = CancellationToken::new(); + + let _listener_handle = spawn_task_listener(config, consumer_tx, shutdown.clone()); + + let js_context = jetstream::new(nats_client.clone()); + + // Send first notification + js_context + .publish( + format!("{prefix}.{SERVICE_NAME}.tasks.workspace-1"), + "first notification".into(), + ) + .await + .expect("failed to publish"); + + let _first = tokio::time::timeout(Duration::from_secs(5), consumer_rx.recv()) + .await + .expect("timeout") + .expect("consumer channel closed"); + + // Send second notification for same workspace + js_context + .publish( + format!("{prefix}.{SERVICE_NAME}.tasks.workspace-1"), + "second notification".into(), + ) + .await + .expect("failed to publish"); + + let second = tokio::time::timeout(Duration::from_secs(5), consumer_rx.recv()) + .await + .expect("timeout") + .expect("consumer channel closed"); + + assert_eq!(second.key, "workspace-1"); + + shutdown.cancel(); +} + +#[test(tokio::test)] +async fn listener_skips_messages_with_invalid_subjects() { + let prefix = nats_prefix(); + let nats_client = nats(prefix.clone()).await; + let (tasks_stream, requests_stream) = create_test_streams(&prefix, &nats_client).await; + + let config = FairSchedulingConfig::for_workspace_partitioning( + SERVICE_NAME, + tasks_stream.clone(), + requests_stream, + Some(prefix.clone()), + ); + + let (consumer_tx, mut consumer_rx) = mpsc::channel(10); + let shutdown = CancellationToken::new(); + + let _listener_handle = spawn_task_listener(config, consumer_tx, shutdown.clone()); + + let js_context = jetstream::new(nats_client.clone()); + + // Publish invalid subject (missing workspace id) + js_context + .publish( + format!("{prefix}.{SERVICE_NAME}.tasks"), + "invalid notification".into(), + ) + .await + .expect("failed to publish"); + + // Wait a bit to ensure message is processed + tokio::time::sleep(Duration::from_millis(500)).await; + + // Verify invalid message was acknowledged (not stuck) + let mut task_listener_consumer = + get_consumer_with_retry(&tasks_stream, &format!("{SERVICE_NAME}-task-listener")).await; + + let info = task_listener_consumer + .info() + .await + .expect("failed to get consumer info"); + + assert_eq!( + info.num_pending, 0, + "invalid message should be acknowledged and skipped" + ); + + // Publish valid message to ensure listener is still working + js_context + .publish( + format!("{prefix}.{SERVICE_NAME}.tasks.workspace-1"), + "valid notification".into(), + ) + .await + .expect("failed to publish"); + + let ready = tokio::time::timeout(Duration::from_secs(5), consumer_rx.recv()) + .await + .expect("timeout") + .expect("consumer channel closed"); + + assert_eq!(ready.key, "workspace-1"); + + shutdown.cancel(); +} + +#[test(tokio::test)] +async fn listener_handles_concurrent_task_notifications() { + let prefix = nats_prefix(); + let nats_client = nats(prefix.clone()).await; + let (tasks_stream, requests_stream) = create_test_streams(&prefix, &nats_client).await; + + let config = FairSchedulingConfig::for_workspace_partitioning( + SERVICE_NAME, + tasks_stream, + requests_stream, + Some(prefix.clone()), + ); + + let (consumer_tx, mut consumer_rx) = mpsc::channel(20); + let shutdown = CancellationToken::new(); + + let _listener_handle = spawn_task_listener(config, consumer_tx, shutdown.clone()); + + let js_context = jetstream::new(nats_client.clone()); + + // Send multiple notifications concurrently + let mut publish_tasks = Vec::new(); + for i in 0..10 { + let js = js_context.clone(); + let prefix_clone = prefix.clone(); + publish_tasks.push(tokio::spawn(async move { + js.publish( + format!("{prefix_clone}.{SERVICE_NAME}.tasks.workspace-{i}"), + "concurrent notification".into(), + ) + .await + .expect("failed to publish"); + })); + } + + // Wait for all publishes to complete + for task in publish_tasks { + task.await.expect("publish task failed"); + } + + // Collect all consumers + let mut received_keys = Vec::new(); + for _ in 0..10 { + let ready = tokio::time::timeout(Duration::from_secs(5), consumer_rx.recv()) + .await + .expect("timeout waiting for consumer") + .expect("consumer channel closed"); + received_keys.push(ready.key); + } + + received_keys.sort(); + let expected: Vec = (0..10).map(|i| format!("workspace-{i}")).collect(); + assert_eq!(received_keys, expected); + + shutdown.cancel(); +} + +#[test(tokio::test(flavor = "multi_thread", worker_threads = 2))] +async fn listener_fair_scheduling_across_workspaces() { + let prefix = nats_prefix(); + let nats_client = nats(prefix.clone()).await; + let (tasks_stream, requests_stream) = create_test_streams(&prefix, &nats_client).await; + + let config = FairSchedulingConfig::for_workspace_partitioning( + SERVICE_NAME, + tasks_stream.clone(), + requests_stream.clone(), + Some(prefix.clone()), + ); + + let (consumer_tx, mut consumer_rx) = mpsc::channel(20); + let shutdown = CancellationToken::new(); + + let _listener_handle = spawn_task_listener(config, consumer_tx, shutdown.clone()); + + let js_context = jetstream::new(nats_client.clone()); + + // Create heavy load for workspace-a + for _ in 0..5 { + js_context + .publish( + format!("{prefix}.{SERVICE_NAME}.tasks.workspace-a"), + "workspace-a notification".into(), + ) + .await + .expect("failed to publish"); + } + + // Small delay to let workspace-a start + tokio::time::sleep(Duration::from_millis(100)).await; + + // Create light load for workspace-b + for _ in 0..2 { + js_context + .publish( + format!("{prefix}.{SERVICE_NAME}.tasks.workspace-b"), + "workspace-b notification".into(), + ) + .await + .expect("failed to publish"); + } + + // Collect consumers and verify both workspaces get processed + let mut workspace_a_count = 0; + let mut workspace_b_count = 0; + + for _ in 0..7 { + let ready = tokio::time::timeout(Duration::from_secs(5), consumer_rx.recv()) + .await + .expect("timeout") + .expect("consumer channel closed"); + + if ready.key == "workspace-a" { + workspace_a_count += 1; + } else if ready.key == "workspace-b" { + workspace_b_count += 1; + } + } + + assert_eq!( + workspace_a_count, 5, + "workspace-a should process 5 notifications" + ); + assert_eq!( + workspace_b_count, 2, + "workspace-b should process 2 notifications" + ); + + // Verify consumers were created + let consumer_a_name = format!("{SERVICE_NAME}-ws-workspace-a"); + let consumer_b_name = format!("{SERVICE_NAME}-ws-workspace-b"); + + let _consumer_a = get_consumer_with_retry(&requests_stream, &consumer_a_name).await; + let _consumer_b = get_consumer_with_retry(&requests_stream, &consumer_b_name).await; + + shutdown.cancel(); +} + +#[test(tokio::test)] +async fn listener_processes_messages_in_order() { + let prefix = nats_prefix(); + let nats_client = nats(prefix.clone()).await; + let (tasks_stream, requests_stream) = create_test_streams(&prefix, &nats_client).await; + + let config = FairSchedulingConfig::for_workspace_partitioning( + SERVICE_NAME, + tasks_stream, + requests_stream, + Some(prefix.clone()), + ); + + let (consumer_tx, mut consumer_rx) = mpsc::channel(10); + let shutdown = CancellationToken::new(); + + let _listener_handle = spawn_task_listener(config, consumer_tx, shutdown.clone()); + + let js_context = jetstream::new(nats_client.clone()); + + // Publish multiple notifications in order for the same workspace + for i in 0..5 { + js_context + .publish( + format!("{prefix}.{SERVICE_NAME}.tasks.workspace-1"), + format!("notification-{i}").into(), + ) + .await + .expect("failed to publish"); + } + + // Verify we receive consumers for each notification + for _ in 0..5 { + let ready = tokio::time::timeout(Duration::from_secs(5), consumer_rx.recv()) + .await + .expect("timeout") + .expect("consumer channel closed"); + + assert_eq!(ready.key, "workspace-1"); + } + + shutdown.cancel(); +}