From a4a00087055186095ba3f7e597e746aa8dc80bc6 Mon Sep 17 00:00:00 2001 From: Patrick Lee Scott Date: Mon, 17 Aug 2026 05:03:13 -0500 Subject: [PATCH] feat: add aggregate load suite and bus idle/batch fixes MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Idle wait vs drain is an explicit RunOptions policy so listen() no longer treats an empty recv as EOF. SQL claim peeks before UPDATE; Postgres LISTEN/NOTIFY wakes waiters. Kafka creates subscribed topics before consume, reuses the consumer, and commits one high-water offset per fetched batch. NATS/Kafka expose batch fetch/produce. SQLite/Postgres commits retry BUSY_SNAPSHOT. tests/load is a workspace-excluded suite with applied and pipelined bus cells across memory/sqlite/postgres × direct/http/grpc/bus (nats, kafka, rabbitmq) plus locks and snapshot frequencies. Implements [[tasks/load-test-bench-1]] [[tasks/load-test-bench-2]] [[tasks/load-test-bench-3]] --- Cargo.toml | 2 +- Makefile | 87 +++++ src/bus/in_memory_bus.rs | 24 ++ src/bus/kafka.rs | 299 ++++++++++++--- src/bus/kafka_bus.rs | 26 +- src/bus/mod.rs | 2 +- src/bus/nats.rs | 71 +++- src/bus/postgres_bus.rs | 55 ++- src/bus/run_options.rs | 25 ++ src/bus/runner/receive_loop.rs | 44 ++- src/bus/source.rs | 13 + src/bus/sql_bus_common.rs | 72 +++- src/bus/sqlite_bus.rs | 16 + src/microsvc/workers.rs | 2 +- src/sqlx_repo/repo/commit.rs | 57 ++- tests/kafka_transport/main.rs | 70 ++++ tests/load/.gitignore | 1 + tests/load/Cargo.toml | 40 ++ tests/load/src/bin/load-client.rs | 97 +++++ tests/load/src/bin/load-host.rs | 91 +++++ tests/load/src/bin/load-suite.rs | 178 +++++++++ tests/load/src/client.rs | 213 +++++++++++ tests/load/src/counter.rs | 57 +++ tests/load/src/host.rs | 358 +++++++++++++++++ tests/load/src/invoke.rs | 345 +++++++++++++++++ tests/load/src/kinds.rs | 138 +++++++ tests/load/src/lib.rs | 166 ++++++++ tests/load/src/stats.rs | 72 ++++ tests/load/src/suite.rs | 613 ++++++++++++++++++++++++++++++ 29 files changed, 3128 insertions(+), 106 deletions(-) create mode 100644 tests/load/.gitignore create mode 100644 tests/load/Cargo.toml create mode 100644 tests/load/src/bin/load-client.rs create mode 100644 tests/load/src/bin/load-host.rs create mode 100644 tests/load/src/bin/load-suite.rs create mode 100644 tests/load/src/client.rs create mode 100644 tests/load/src/counter.rs create mode 100644 tests/load/src/host.rs create mode 100644 tests/load/src/invoke.rs create mode 100644 tests/load/src/kinds.rs create mode 100644 tests/load/src/lib.rs create mode 100644 tests/load/src/stats.rs create mode 100644 tests/load/src/suite.rs diff --git a/Cargo.toml b/Cargo.toml index b46a2e6e..132f4b92 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [workspace] members = ["distributed_macros", "distributed_cli"] -exclude = ["tests/e2e-ui"] +exclude = ["tests/e2e-ui", "tests/load"] resolver = "2" [workspace.package] diff --git a/Makefile b/Makefile index 5262d84a..b09867c2 100644 --- a/Makefile +++ b/Makefile @@ -44,6 +44,93 @@ compose-up: compose-down: $(DOCKER_COMPOSE) down $(COMPOSE_DOWN_FLAGS) +LOAD_MANIFEST ?= tests/load/Cargo.toml +LOAD_REPO ?= memory +LOAD_BIND ?= 127.0.0.1:8790 +LOAD_SCENARIO ?= unique-create +LOAD_CONCURRENCY ?= 32 +LOAD_DURATION ?= 15s +LOAD_WARMUP ?= 2s +LOAD_DATABASE_URL ?= $(DATABASE_URL) +LOAD_SQLITE_PATH ?= target/load.sqlite + +LOAD_FEATURES ?= kafka,rabbitmq +LOAD_FILTER ?= +LOAD_SUITE_FLAGS ?= +LOAD_SNAPSHOTS ?= + +.PHONY: load-host load-client load-run load-matrix load-test load-suite + +## Opt-in load harness (not part of `make test`). See tests/load/src/bin/*.rs --help. +load-host: + $(CARGO) run --manifest-path $(LOAD_MANIFEST) --release --bin load-host -- \ + --repo $(LOAD_REPO) --bind $(LOAD_BIND) \ + --database-url $(LOAD_DATABASE_URL) --sqlite-path $(LOAD_SQLITE_PATH) \ + $(if $(LOAD_SNAPSHOTS),--snapshots $(LOAD_SNAPSHOTS),) + +load-client: + $(CARGO) run --manifest-path $(LOAD_MANIFEST) --release --bin load-client -- \ + --url http://$(LOAD_BIND) --scenario $(LOAD_SCENARIO) \ + --concurrency $(LOAD_CONCURRENCY) --duration $(LOAD_DURATION) \ + --warmup $(LOAD_WARMUP) --repo $(LOAD_REPO) \ + $(if $(LOAD_SNAPSHOTS),--snapshots $(LOAD_SNAPSHOTS),) + +## Build, start host, wait for /health, run client, stop host. +load-run: + @set -eu; \ + $(CARGO) build --manifest-path $(LOAD_MANIFEST) --release --bins; \ + host_bin="tests/load/target/release/load-host"; \ + client_bin="tests/load/target/release/load-client"; \ + if [ ! -x "$$host_bin" ]; then host_bin="target/release/load-host"; fi; \ + if [ ! -x "$$client_bin" ]; then client_bin="target/release/load-client"; fi; \ + host_args="--repo $(LOAD_REPO) --bind $(LOAD_BIND) --sqlite-path $(LOAD_SQLITE_PATH)"; \ + if [ -n "$(LOAD_DATABASE_URL)" ]; then host_args="$$host_args --database-url $(LOAD_DATABASE_URL)"; fi; \ + if [ -n "$(LOAD_SNAPSHOTS)" ]; then host_args="$$host_args --snapshots $(LOAD_SNAPSHOTS)"; fi; \ + $$host_bin $$host_args & host_pid=$$!; \ + cleanup() { kill $$host_pid >/dev/null 2>&1 || true; wait $$host_pid >/dev/null 2>&1 || true; }; \ + trap cleanup EXIT INT TERM; \ + ready=0; \ + for _ in 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20; do \ + if curl -sf "http://$(LOAD_BIND)/health" >/dev/null; then ready=1; break; fi; \ + sleep 0.25; \ + done; \ + if [ "$$ready" != "1" ]; then echo "load-host did not become healthy at http://$(LOAD_BIND)/health" >&2; exit 1; fi; \ + client_args="--url http://$(LOAD_BIND) --scenario $(LOAD_SCENARIO) \ + --concurrency $(LOAD_CONCURRENCY) --duration $(LOAD_DURATION) \ + --warmup $(LOAD_WARMUP) --repo $(LOAD_REPO)"; \ + if [ -n "$(LOAD_SNAPSHOTS)" ]; then client_args="$$client_args --snapshots $(LOAD_SNAPSHOTS)"; fi; \ + $$client_bin $$client_args + +## Compare memory, sqlite, and postgres (postgres needs `make compose-up` or a live DATABASE_URL). +load-matrix: + @set -eu; \ + for repo in memory sqlite postgres; do \ + echo "======== $$repo / $(LOAD_SCENARIO) ========"; \ + $(MAKE) load-run LOAD_REPO=$$repo LOAD_SCENARIO=$(LOAD_SCENARIO) \ + LOAD_CONCURRENCY=$(LOAD_CONCURRENCY) LOAD_DURATION=$(LOAD_DURATION) \ + LOAD_WARMUP=$(LOAD_WARMUP) LOAD_BIND=$(LOAD_BIND); \ + done + +load-test: + $(CARGO) test --manifest-path $(LOAD_MANIFEST) + +## Full Counter suite: every dispatch, bus (incl. kafka/rabbitmq), lock, snapshot, scenario. +## make compose-up && make load-suite +## make load-suite LOAD_FILTER=direct LOAD_DURATION=3s +## make load-run LOAD_SNAPSHOTS=10 +load-suite: + DATABASE_URL="$(LOAD_DATABASE_URL)" \ + NATS_URL="$(NATS_URL)" \ + KAFKA_BROKERS="$(KAFKA_BROKERS)" \ + AMQP_URL="$(AMQP_URL)" \ + $(CARGO) run --manifest-path $(LOAD_MANIFEST) --release \ + $(if $(LOAD_FEATURES),--features $(LOAD_FEATURES),) \ + --bin load-suite -- \ + --duration $(LOAD_DURATION) --warmup $(LOAD_WARMUP) \ + --concurrency $(LOAD_CONCURRENCY) \ + $(if $(LOAD_FILTER),--filter $(LOAD_FILTER),) \ + $(LOAD_SUITE_FLAGS) + .PHONY: contracts-check ## Read-only aggregate contract lifecycle check (never writes tracked files). diff --git a/src/bus/in_memory_bus.rs b/src/bus/in_memory_bus.rs index 2a76f974..481eefb3 100644 --- a/src/bus/in_memory_bus.rs +++ b/src/bus/in_memory_bus.rs @@ -13,6 +13,8 @@ use std::collections::{HashMap, VecDeque}; use std::sync::{Arc, Mutex}; +use tokio::sync::Notify; + use super::source::{MessageSource, ReceivedMessage}; use super::{run_source, Bus, BusConsumer, MessageRouter, RunOptions, TransportError}; use super::{Message, OrderedDelivery}; @@ -35,6 +37,7 @@ pub struct InMemoryBus { queues: Queues, topics: Topics, source_epoch: ProjectionEpoch, + wake: Arc, } impl Default for InMemoryBus { @@ -44,6 +47,7 @@ impl Default for InMemoryBus { topics: Topics::default(), source_epoch: ProjectionEpoch::new(format!("instance-{}", uuid::Uuid::now_v7())) .expect("an in-memory bus UUID is a valid projection source epoch"), + wake: Arc::new(Notify::new()), } } } @@ -60,6 +64,7 @@ impl InMemoryBus { .entry(message.name().to_string()) .or_default() .push_back(message); + self.wake.notify_waiters(); Ok(()) } @@ -79,6 +84,7 @@ impl InMemoryBus { .entry(message.name().to_string()) .or_default() .push(message); + self.wake.notify_waiters(); Ok(()) } @@ -138,6 +144,7 @@ impl BusConsumer for InMemoryBus { let source = QueueSource { queues: self.queues.clone(), names, + wake: Arc::clone(&self.wake), }; run_source(router, source, options).await } @@ -153,6 +160,7 @@ impl BusConsumer for InMemoryBus { names, cursors: TopicCursors::default(), source_epoch: self.source_epoch.clone(), + wake: Arc::clone(&self.wake), }; run_source(router, source, options).await } @@ -162,6 +170,7 @@ impl BusConsumer for InMemoryBus { struct QueueSource { queues: Queues, names: Vec, + wake: Arc, } impl MessageSource for QueueSource { @@ -184,6 +193,11 @@ impl MessageSource for QueueSource { } Ok(None) } + + async fn wait(&mut self) -> Result<(), TransportError> { + self.wake.notified().await; + Ok(()) + } } /// Fan-out source over the named retained logs: each `TopicSource` has its own @@ -193,6 +207,7 @@ struct TopicSource { names: Vec, cursors: TopicCursors, source_epoch: ProjectionEpoch, + wake: Arc, } impl MessageSource for TopicSource { @@ -238,6 +253,11 @@ impl MessageSource for TopicSource { } Ok(None) } + + async fn wait(&mut self) -> Result<(), TransportError> { + self.wake.notified().await; + Ok(()) + } } struct TopicSettlement { @@ -375,10 +395,12 @@ mod tests { let mut a = QueueSource { queues: bus.queues.clone(), names: vec!["work".to_string()], + wake: Arc::clone(&bus.wake), }; let mut b = QueueSource { queues: bus.queues.clone(), names: vec!["work".to_string()], + wake: Arc::clone(&bus.wake), }; let mut got = Vec::new(); // Alternate; each pop removes the message (competing). @@ -437,6 +459,7 @@ mod tests { names: vec!["evt".into()], cursors: TopicCursors::default(), source_epoch: bus.source_epoch.clone(), + wake: Arc::clone(&bus.wake), }; let first = block_on(source.recv()).unwrap().unwrap(); @@ -482,6 +505,7 @@ mod tests { names: vec!["evt".into()], cursors: TopicCursors::default(), source_epoch: bus.source_epoch.clone(), + wake: Arc::clone(&bus.wake), }; let received = block_on(source.recv()).unwrap().unwrap(); assert_eq!(received.message().id(), Some("e0")); diff --git a/src/bus/kafka.rs b/src/bus/kafka.rs index cabc3003..38f5e766 100644 --- a/src/bus/kafka.rs +++ b/src/bus/kafka.rs @@ -13,11 +13,15 @@ //! Requires the `kafka` feature (builds `librdkafka` via cmake). Integration- //! tested in `tests/kafka_transport` against a broker (see `compose.yaml`). -use std::sync::Arc; +use std::collections::{HashMap, HashSet, VecDeque}; +use std::sync::{Arc, Mutex}; use std::time::{Duration, Instant}; +use rdkafka::admin::{AdminClient, AdminOptions, NewTopic, TopicReplication}; +use rdkafka::client::DefaultClientContext; use rdkafka::config::ClientConfig; use rdkafka::consumer::{Consumer, StreamConsumer}; +use rdkafka::error::RDKafkaErrorCode; use rdkafka::message::{Header, Headers, OwnedHeaders}; use rdkafka::producer::{FutureProducer, FutureRecord}; use rdkafka::{Message as KafkaMessageTrait, Offset, TopicPartitionList}; @@ -51,6 +55,9 @@ impl KafkaPublisher { .set("bootstrap.servers", brokers) .set("acks", "all") .set("message.timeout.ms", "10000") + .set("linger.ms", "5") + .set("batch.num.messages", "10000") + .set("queue.buffering.max.kbytes", "65536") .create() .map_err(|err| retryable("kafka producer", err))?; Ok(Self::new(producer)) @@ -92,25 +99,143 @@ impl MessagePublisher for KafkaPublisher { .map_err(|(err, _)| retryable("kafka send", err))?; Ok(()) } + + async fn publish_batch(&self, messages: Vec) -> Result<(), TransportError> { + if messages.is_empty() { + return Ok(()); + } + let prepared: Vec<(String, String, Vec, OwnedHeaders)> = messages + .into_iter() + .map(|message| { + let topic = message.name().to_string(); + let key = message.id().unwrap_or(message.name()).to_string(); + let headers = owned_headers(&message); + (topic, key, message.payload, headers) + }) + .collect(); + let mut pending = Vec::with_capacity(prepared.len()); + for (topic, key, payload, headers) in &prepared { + let record = FutureRecord::to(topic) + .payload(payload) + .key(key) + .headers(headers.clone()); + pending.push(self.producer.send(record, self.send_timeout)); + } + for send in pending { + send.await + .map_err(|(err, _)| retryable("kafka send", err))?; + } + Ok(()) + } } /// Consumes a topic with a consumer group, committing offsets on ack. pub struct KafkaSource { - consumer: Arc, + pub(crate) consumer: Arc, fetch_timeout: Duration, + fetch_max: usize, strip_prefix: Option, + buffer: VecDeque, + book: Arc>, +} + +/// Records successful acks and commits one high-water mark per partition +/// after a fetch batch is fully settled (or on nack). +struct KafkaOffsetBook { + consumer: Arc, + high_water: HashMap<(String, i32), i64>, + invalidate: HashSet<(String, i32)>, + unsettled: usize, +} + +impl KafkaOffsetBook { + fn new(consumer: Arc) -> Self { + Self { + consumer, + high_water: HashMap::new(), + invalidate: HashSet::new(), + unsettled: 0, + } + } + + fn begin_batch(&mut self, count: usize) -> Result<(), TransportError> { + self.flush()?; + self.unsettled = count; + Ok(()) + } + + fn record_success( + &mut self, + topic: String, + partition: i32, + offset: i64, + ) -> Result<(), TransportError> { + let next = offset + 1; + self.high_water + .entry((topic, partition)) + .and_modify(|current| *current = (*current).max(next)) + .or_insert(next); + self.unsettled = self.unsettled.saturating_sub(1); + if self.unsettled == 0 { + self.flush()?; + } + Ok(()) + } + + fn record_nack( + &mut self, + topic: String, + partition: i32, + offset: i64, + ) -> Result<(), TransportError> { + if let Some(next) = self.high_water.get_mut(&(topic.clone(), partition)) { + if *next > offset { + *next = offset; + } + } + self.unsettled = 0; + self.invalidate.insert((topic, partition)); + self.flush() + } + + fn take_invalidated(&mut self) -> HashSet<(String, i32)> { + std::mem::take(&mut self.invalidate) + } + + fn flush(&mut self) -> Result<(), TransportError> { + if self.high_water.is_empty() { + return Ok(()); + } + let mut tpl = TopicPartitionList::new(); + for ((topic, partition), next) in self.high_water.drain() { + tpl.add_partition_offset(&topic, partition, Offset::Offset(next)) + .map_err(|err| retryable("kafka offset", err))?; + } + self.consumer + .commit(&tpl, rdkafka::consumer::CommitMode::Async) + .map_err(|err| retryable("kafka commit", err)) + } } impl KafkaSource { /// Wrap an existing subscribed consumer. pub fn new(consumer: Arc) -> Self { Self { - consumer, fetch_timeout: Duration::from_secs(5), + fetch_max: 32, strip_prefix: None, + buffer: VecDeque::new(), + book: Arc::new(Mutex::new(KafkaOffsetBook::new(Arc::clone(&consumer)))), + consumer, } } + /// How many extra records to drain after the first wait. + pub fn with_fetch_max(mut self, max: usize) -> Self { + self.fetch_max = max.max(1); + self + } + /// How long `recv` waits for a record before returning `Ok(None)`. pub fn with_fetch_timeout(mut self, timeout: Duration) -> Self { self.fetch_timeout = timeout; @@ -127,16 +252,25 @@ impl KafkaSource { /// Connect a consumer (group `group_id`, auto-commit off, earliest reset) and /// subscribe to `topics`. + /// + /// Topics are created before subscribe. A Kafka consumer that names a topic + /// that does not exist yet is not assigned that topic until a later metadata + /// refresh (often minutes). `listen` on several command names would then + /// consume the first produced name and stall on every later name — the + /// load-suite hot-increment cell after a successful initialize. pub async fn connect( brokers: &str, group_id: &str, topics: &[&str], ) -> Result { + ensure_topics(brokers, topics).await?; let consumer: StreamConsumer = ClientConfig::new() .set("bootstrap.servers", brokers) .set("group.id", group_id) .set("enable.auto.commit", "false") .set("auto.offset.reset", "earliest") + .set("allow.auto.create.topics", "true") + .set("topic.metadata.refresh.interval.ms", "1000") .create() .map_err(|err| retryable("kafka consumer", err))?; consumer @@ -144,22 +278,33 @@ impl KafkaSource { .map_err(|err| retryable("kafka subscribe", err))?; Ok(Self::new(Arc::new(consumer))) } -} - -impl MessageSource for KafkaSource { - type Received = KafkaReceived; - fn transport_name(&self) -> &'static str { - "kafka" + async fn fill_buffer(&mut self, first_timeout: Duration) -> Result { + if !self.buffer.is_empty() { + return Ok(true); + } + let Some(first) = self.poll_one(first_timeout).await? else { + return Ok(false); + }; + self.buffer.push_back(first); + while self.buffer.len() < self.fetch_max { + match self.poll_one(Duration::from_millis(1)).await? { + Some(next) => self.buffer.push_back(next), + None => break, + } + } + self.book + .lock() + .map_err(|_| TransportError::retryable("kafka offset book poisoned"))? + .begin_batch(self.buffer.len())?; + Ok(true) } - async fn recv(&mut self) -> Result, TransportError> { - // Poll within the fetch-timeout budget. Kafka surfaces transient broker - // transport/coordination errors (normal during group bootstrap and - // rebalances) as recv errors; the client recovers, so we retry until a - // message arrives or the budget elapses (drain → None), rather than - // ending the run on a transient hiccup. - let deadline = Instant::now() + self.fetch_timeout; + async fn poll_one(&self, timeout: Duration) -> Result, TransportError> { + if timeout.is_zero() { + return Ok(None); + } + let deadline = Instant::now() + timeout; loop { let remaining = deadline.saturating_duration_since(Instant::now()); if remaining.is_zero() { @@ -170,12 +315,12 @@ impl MessageSource for KafkaSource { return Ok(Some(KafkaReceived::from_borrowed( &borrowed, self.consumer.clone(), + Arc::clone(&self.book), self.strip_prefix.as_deref(), ))); } Ok(Err(_transient)) => { - // Back off briefly, then retry within the remaining budget. - tokio::time::sleep(Duration::from_millis(100)).await; + tokio::time::sleep(Duration::from_millis(20)).await; } Err(_elapsed) => return Ok(None), } @@ -183,9 +328,74 @@ impl MessageSource for KafkaSource { } } +impl MessageSource for KafkaSource { + type Received = KafkaReceived; + + fn transport_name(&self) -> &'static str { + "kafka" + } + + async fn recv(&mut self) -> Result, TransportError> { + { + let invalidated = self + .book + .lock() + .map_err(|_| TransportError::retryable("kafka offset book poisoned"))? + .take_invalidated(); + if !invalidated.is_empty() { + self.buffer.retain(|message| { + !invalidated.contains(&(message.topic.clone(), message.partition)) + }); + } + } + if !self.fill_buffer(self.fetch_timeout).await? { + return Ok(None); + } + Ok(self.buffer.pop_front()) + } + + async fn wait(&mut self) -> Result<(), TransportError> { + let _ = self.fill_buffer(self.fetch_timeout).await?; + Ok(()) + } +} + +/// Create `topics` if they are missing so a consumer can be assigned immediately. +async fn ensure_topics(brokers: &str, topics: &[&str]) -> Result<(), TransportError> { + if topics.is_empty() { + return Ok(()); + } + let admin: AdminClient = ClientConfig::new() + .set("bootstrap.servers", brokers) + .create() + .map_err(|err| retryable("kafka admin", err))?; + let new_topics: Vec> = topics + .iter() + .map(|topic| NewTopic::new(topic, 1, TopicReplication::Fixed(1))) + .collect(); + let results = admin + .create_topics( + &new_topics, + &AdminOptions::new().operation_timeout(Some(Duration::from_secs(10))), + ) + .await + .map_err(|err| retryable("kafka create topics", err))?; + for result in results { + match result { + Ok(_) => {} + Err((_, RDKafkaErrorCode::TopicAlreadyExists)) => {} + Err((name, code)) => { + return Err(retryable("kafka create topics", format!("{name}: {code}"))); + } + } + } + Ok(()) +} + /// A consumed record plus the means to commit/seek its offset. pub struct KafkaReceived { consumer: Arc, + book: Arc>, topic: String, partition: i32, offset: i64, @@ -196,6 +406,7 @@ impl KafkaReceived { fn from_borrowed( borrowed: &rdkafka::message::BorrowedMessage<'_>, consumer: Arc, + book: Arc>, strip_prefix: Option<&str>, ) -> Self { let payload = borrowed.payload().map(|p| p.to_vec()).unwrap_or_default(); @@ -222,30 +433,13 @@ impl KafkaReceived { ); Self { consumer, + book, topic, partition: borrowed.partition(), offset: borrowed.offset(), message, } } - - /// Commit this record's offset (the next offset to read) without blocking the - /// async runtime. - /// - /// Uses [`CommitMode::Async`]: the offset is handed to librdkafka's background - /// thread and this returns immediately, rather than blocking the tokio worker - /// on a broker round trip as `CommitMode::Sync` does. This keeps at-least-once - /// semantics — the runner already acks only after handler effects complete, so - /// a crash between effect and the async commit landing simply redelivers the - /// record (a duplicate the consumer must tolerate, which it already does). - fn commit_offset(&self) -> Result<(), TransportError> { - let mut tpl = TopicPartitionList::new(); - tpl.add_partition_offset(&self.topic, self.partition, Offset::Offset(self.offset + 1)) - .map_err(|err| retryable("kafka offset", err))?; - self.consumer - .commit(&tpl, rdkafka::consumer::CommitMode::Async) - .map_err(|err| retryable("kafka commit", err)) - } } impl ReceivedMessage for KafkaReceived { @@ -254,18 +448,21 @@ impl ReceivedMessage for KafkaReceived { } async fn ack(self) -> Result<(), TransportError> { - self.commit_offset() + self.book + .lock() + .map_err(|_| TransportError::retryable("kafka offset book poisoned"))? + .record_success(self.topic, self.partition, self.offset) } async fn nack(self, _reason: &str) -> Result<(), TransportError> { - // Do not commit; seek back so this record is re-read (redelivery). - // - // `seek` blocks the calling thread until it takes effect (up to the - // timeout) and rdkafka exposes no async variant, so run it on the - // blocking pool to avoid stalling the tokio worker. The consumer is - // Arc-shared and `seek` takes `&self`, so a clone moves cleanly into the - // blocking task. - let consumer = self.consumer.clone(); + self.book + .lock() + .map_err(|_| TransportError::retryable("kafka offset book poisoned"))? + .record_nack(self.topic.clone(), self.partition, self.offset)?; + // Seek back so this record is re-read. `seek` is blocking; run it off + // the tokio worker. Remaining buffered records for this partition are + // dropped on the next recv. + let consumer = self.consumer; let topic = self.topic; let partition = self.partition; let offset = self.offset; @@ -283,12 +480,16 @@ impl ReceivedMessage for KafkaReceived { } async fn dead_letter(self, _reason: &str) -> Result<(), TransportError> { - // Skip the poison record by committing past it. A DLQ-topic producer is a - // follow-up. - self.commit_offset() + self.book + .lock() + .map_err(|_| TransportError::retryable("kafka offset book poisoned"))? + .record_success(self.topic, self.partition, self.offset) } async fn park(self, _reason: &str) -> Result<(), TransportError> { - self.commit_offset() + self.book + .lock() + .map_err(|_| TransportError::retryable("kafka offset book poisoned"))? + .record_success(self.topic, self.partition, self.offset) } } diff --git a/src/bus/kafka_bus.rs b/src/bus/kafka_bus.rs index d21320e8..dbf97b0c 100644 --- a/src/bus/kafka_bus.rs +++ b/src/bus/kafka_bus.rs @@ -18,11 +18,15 @@ //! //! Requires the `kafka` feature. Integration-tested in `tests/kafka_transport`. +use std::collections::HashMap; use std::future::{Future, IntoFuture}; use std::pin::Pin; use std::sync::Arc; use std::time::Duration; +use rdkafka::consumer::StreamConsumer; +use tokio::sync::Mutex; + use super::kafka::{KafkaPublisher, KafkaSource}; use super::{ run_source, Bus, BusConsumer, BusTopologyConfig, MessagePublisher, MessageRouter, RunOptions, @@ -39,6 +43,8 @@ pub struct KafkaBus { publisher: Arc, topology: BusTopologyConfig, fetch_timeout: Duration, + fetch_max: usize, + consumers: Arc>>>, } /// Awaitable builder returned by [`KafkaBus::connect`]. @@ -77,6 +83,8 @@ impl KafkaBusConnect { publisher: Arc::new(publisher), topology, fetch_timeout: self.fetch_timeout, + fetch_max: 32, + consumers: Arc::new(Mutex::new(HashMap::new())), }) } } @@ -176,9 +184,21 @@ impl KafkaBus { .resolve_consumer_group(router.as_ref(), "kafka")?; let group_id = format!("{namespace}.{group}.{suffix}"); let topic_refs: Vec<&str> = topics.iter().map(String::as_str).collect(); - let source = KafkaSource::connect(&self.brokers, &group_id, &topic_refs) - .await? + let consumer = { + let mut cache = self.consumers.lock().await; + if let Some(existing) = cache.get(&group_id) { + Arc::clone(existing) + } else { + let created = KafkaSource::connect(&self.brokers, &group_id, &topic_refs) + .await? + .consumer; + cache.insert(group_id, Arc::clone(&created)); + created + } + }; + let source = KafkaSource::new(consumer) .with_fetch_timeout(self.fetch_timeout) + .with_fetch_max(self.fetch_max) .with_strip_prefix(prefix); run_source(router, source, options).await } @@ -248,6 +268,8 @@ mod tests { publisher: Arc::new(KafkaPublisher::new(producer)), topology: BusTopologyConfig::default(), fetch_timeout: Duration::from_millis(1), + fetch_max: 32, + consumers: Arc::new(Mutex::new(HashMap::new())), } } diff --git a/src/bus/mod.rs b/src/bus/mod.rs index 75b21e9d..c1c89446 100644 --- a/src/bus/mod.rs +++ b/src/bus/mod.rs @@ -154,7 +154,7 @@ pub use ordered_delivery::OrderedDelivery; pub use postgres_bus::{LogReceived, PostgresBus, QueueReceived}; pub use publisher::MessagePublisher; pub use router::MessageRouter; -pub use run_options::{ConsumerDeliveryMode, InboxHook, NoInbox, RunOptions}; +pub use run_options::{ConsumerDeliveryMode, IdlePolicy, InboxHook, NoInbox, RunOptions}; pub use runner::run_source; pub use source::{MessageSource, ReceivedMessage}; #[cfg(feature = "sqlite")] diff --git a/src/bus/nats.rs b/src/bus/nats.rs index 0e897aba..08ccae83 100644 --- a/src/bus/nats.rs +++ b/src/bus/nats.rs @@ -10,6 +10,7 @@ //! Requires the `nats` feature. Integration-tested in `tests/nats_transport` //! against a JetStream-enabled server (see `compose.yaml`). +use std::collections::VecDeque; use std::time::Duration; use async_nats::jetstream::consumer::pull::Config as PullConfig; @@ -102,7 +103,9 @@ impl MessagePublisher for NatsPublisher { pub struct NatsJetStreamSource { consumer: Consumer, fetch_timeout: Duration, + fetch_max: usize, strip_prefix: Option, + buffer: VecDeque, } impl NatsJetStreamSource { @@ -111,10 +114,18 @@ impl NatsJetStreamSource { Self { consumer, fetch_timeout: Duration::from_millis(500), + fetch_max: 32, strip_prefix: None, + buffer: VecDeque::new(), } } + /// How many messages to pull per fetch. + pub fn with_fetch_max(mut self, max: usize) -> Self { + self.fetch_max = max.max(1); + self + } + /// How long `recv` waits for a message before returning `Ok(None)`. pub fn with_fetch_timeout(mut self, timeout: Duration) -> Self { self.fetch_timeout = timeout; @@ -132,6 +143,42 @@ impl NatsJetStreamSource { self } + async fn fill(&mut self) -> Result<(), TransportError> { + // First message waits up to fetch_timeout. Extra records use a short + // expire so applied (RPC) cells do not stall waiting for a full batch. + self.pull_into(1, self.fetch_timeout).await?; + if self.buffer.is_empty() { + return Ok(()); + } + let extra = self.fetch_max.saturating_sub(1); + if extra > 0 { + self.pull_into(extra, Duration::from_millis(10)).await?; + } + Ok(()) + } + + async fn pull_into(&mut self, max: usize, expires: Duration) -> Result<(), TransportError> { + if max == 0 { + return Ok(()); + } + let mut batch = self + .consumer + .batch() + .max_messages(max) + .expires(expires) + .messages() + .await + .map_err(|err| retryable("nats fetch", err))?; + while let Some(message) = batch.next().await { + let message = message.map_err(|err| retryable("nats batch message", err))?; + self.buffer.push_back(NatsReceived::from_jetstream( + message, + self.strip_prefix.as_deref(), + )); + } + Ok(()) + } + /// Connect to a NATS server URL, then create/open the stream + consumer. pub async fn connect( url: &str, @@ -186,23 +233,17 @@ impl MessageSource for NatsJetStreamSource { } async fn recv(&mut self) -> Result, TransportError> { - let mut batch = self - .consumer - .batch() - .max_messages(1) - .expires(self.fetch_timeout) - .messages() - .await - .map_err(|err| retryable("nats fetch", err))?; + if self.buffer.is_empty() { + self.fill().await?; + } + Ok(self.buffer.pop_front()) + } - match batch.next().await { - Some(Ok(message)) => Ok(Some(NatsReceived::from_jetstream( - message, - self.strip_prefix.as_deref(), - ))), - Some(Err(err)) => Err(retryable("nats batch message", err)), - None => Ok(None), + async fn wait(&mut self) -> Result<(), TransportError> { + if self.buffer.is_empty() { + self.fill().await?; } + Ok(()) } } diff --git a/src/bus/postgres_bus.rs b/src/bus/postgres_bus.rs index a1b57699..3b3fd25b 100644 --- a/src/bus/postgres_bus.rs +++ b/src/bus/postgres_bus.rs @@ -41,7 +41,11 @@ //! [`MessageSource`]: super::MessageSource //! [`run_source`]: super::run_source +use std::sync::Arc; + +use sqlx::postgres::PgListener; use sqlx::{PgConnection, PgPool, Row}; +use tokio::sync::Mutex; use super::sql_bus_common::{ db_err as sql_db_err, message_from_row, metadata_json, validate_log_retry, ClaimedRow, @@ -132,7 +136,10 @@ impl PostgresBus { /// `bus_queue` by message name, so command replicas compete by listening to the /// same registered command names. pub fn new(pool: PgPool) -> Self { - SqlBus::from_dialect(PostgresBusDialect { pool }) + SqlBus::from_dialect(PostgresBusDialect { + pool, + listener: Arc::new(Mutex::new(None)), + }) } /// Build a bus with an explicit group for direct/low-level use. @@ -145,6 +152,7 @@ impl PostgresBus { #[derive(Clone)] pub struct PostgresBusDialect { pool: PgPool, + listener: Arc>>, } impl PostgresBusDialect { @@ -382,7 +390,12 @@ impl SqlBusDialect for PostgresBusDialect { "enqueue", message, ) - .await + .await?; + sqlx::query("SELECT pg_notify('distributed_bus_queue', '')") + .execute(&self.pool) + .await + .map_err(|err| db_err("notify queue", err))?; + Ok(()) } async fn insert_log( @@ -584,6 +597,44 @@ impl SqlBusDialect for PostgresBusDialect { .collect() } + async fn has_claimable(&self, names: &[String]) -> Result { + let row = sqlx::query( + "SELECT 1 FROM bus_queue \ + WHERE (name = ANY($1) OR name IS NULL) AND available_at <= now() \ + AND (locked_until IS NULL OR locked_until <= now()) \ + LIMIT 1", + ) + .bind(names) + .fetch_optional(&self.pool) + .await + .map_err(|err| db_err("peek claimable", err))?; + Ok(row.is_some()) + } + + async fn listen_wakeup(&self) -> Result<(), TransportError> { + let mut listener = { + let mut slot = self.listener.lock().await; + match slot.take() { + Some(existing) => existing, + None => { + let mut created = PgListener::connect_with(&self.pool) + .await + .map_err(|err| db_err("connect notify listener", err))?; + created + .listen("distributed_bus_queue") + .await + .map_err(|err| db_err("listen distributed_bus_queue", err))?; + created + } + } + }; + let received = listener.recv().await; + *self.listener.lock().await = Some(listener); + received + .map(|_| ()) + .map_err(|err| db_err("recv queue notify", err)) + } + async fn log_read( &self, names: &[String], diff --git a/src/bus/run_options.rs b/src/bus/run_options.rs index 2a98ae45..99a0a6c5 100644 --- a/src/bus/run_options.rs +++ b/src/bus/run_options.rs @@ -60,6 +60,21 @@ impl Default for ConsumerDeliveryMode { #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] pub enum NoInbox {} +/// What [`run_source`](super::run_source) does when `recv` returns `Ok(None)`. +/// +/// Adapters use `None` both for "queue is idle" and for "source closed". +/// Drain (the default) keeps the historical test contract: process what is +/// there and return. Wait is the long-lived consumer: park on +/// [`MessageSource::wait`](super::MessageSource::wait) and recv again. +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub enum IdlePolicy { + /// Stop the run when `recv` returns `Ok(None)`. + #[default] + Drain, + /// Wait for work and recv again. Cancel by dropping the `listen` future. + Wait, +} + /// Cross-cutting execution policy for a transport run. /// /// Generic over the inbox hook type `I`, defaulting to [`NoInbox`] for the @@ -70,6 +85,8 @@ pub struct RunOptions { pub delivery_mode: ConsumerDeliveryMode, /// What the runner does with a permanent handler/transport failure. pub failure_policy: FailurePolicy, + /// Whether an empty `recv` ends the run or waits for more work. + pub idle: IdlePolicy, } impl Default for RunOptions { @@ -78,6 +95,7 @@ impl Default for RunOptions { Self { delivery_mode: ConsumerDeliveryMode::default(), failure_policy: FailurePolicy::default(), + idle: IdlePolicy::Drain, } } } @@ -98,6 +116,7 @@ impl RunOptions { Self { delivery_mode: ConsumerDeliveryMode::Inbox(hook), failure_policy: FailurePolicy::default(), + idle: IdlePolicy::Drain, } } @@ -107,6 +126,12 @@ impl RunOptions { self } + /// Keep listening when the source is idle instead of ending the run. + pub fn wait_when_idle(mut self) -> Self { + self.idle = IdlePolicy::Wait; + self + } + /// Whether this run dispatches directly without an inbox. pub fn is_idempotent(&self) -> bool { matches!(self.delivery_mode, ConsumerDeliveryMode::Idempotent) diff --git a/src/bus/runner/receive_loop.rs b/src/bus/runner/receive_loop.rs index c04ce842..eabbdc90 100644 --- a/src/bus/runner/receive_loop.rs +++ b/src/bus/runner/receive_loop.rs @@ -2,7 +2,9 @@ use std::future::Future; use std::sync::Arc; use crate::bus::source::{MessageSource, ReceivedMessage}; -use crate::bus::{FailureAction, MessageRouter, RunOptions, TransportError, TransportErrorKind}; +use crate::bus::{ + FailureAction, IdlePolicy, MessageRouter, RunOptions, TransportError, TransportErrorKind, +}; use crate::bus::{Message, MessageKind}; /// Run the receive loop for a direct transport source. @@ -49,8 +51,14 @@ where let transport = source.transport_name(); loop { - let Some(received) = recv_next(&mut source, service, transport).await? else { - break; + let Some(received) = recv_next(&mut source, service, transport, options.idle).await? else { + match options.idle { + IdlePolicy::Drain => break, + IdlePolicy::Wait => { + source.wait().await?; + continue; + } + } }; // A delivery the transport could not decode is a permanent failure: it @@ -273,17 +281,27 @@ async fn recv_next( source: &mut S, service: Option<&str>, transport: &str, + idle: IdlePolicy, ) -> Result, TransportError> { - match source.recv().await { - Ok(received) => Ok(received), - Err(error) => { - record_transport_failure( - service, - transport, - error.kind(), - crate::telemetry::failure_action::RECV_ERROR, - ); - Err(error) + loop { + match source.recv().await { + Ok(received) => return Ok(received), + Err(error) => { + record_transport_failure( + service, + transport, + error.kind(), + crate::telemetry::failure_action::RECV_ERROR, + ); + // A long-lived consumer treats a transient recv (SQLite BUSY, + // brief broker blip) as idle: wait and try again. Drain keeps + // the historical contract that recv errors end the run. + if idle == IdlePolicy::Wait && error.kind().is_retryable() { + source.wait().await?; + continue; + } + return Err(error); + } } } } diff --git a/src/bus/source.rs b/src/bus/source.rs index a96c10bb..601adf29 100644 --- a/src/bus/source.rs +++ b/src/bus/source.rs @@ -37,6 +37,19 @@ pub trait MessageSource: Send { fn recv( &mut self, ) -> impl Future, TransportError>> + Send + '_; + + /// Park until `recv` might return work. Used only when + /// [`IdlePolicy::Wait`](super::IdlePolicy::Wait) is set. + /// + /// The default sleeps a short interval so an adapter that cannot wait + /// does not busy-spin. Adapters with a wakeup (notify, LISTEN, broker + /// long-poll) should override this. + fn wait(&mut self) -> impl Future> + Send + '_ { + async { + tokio::time::sleep(std::time::Duration::from_millis(25)).await; + Ok(()) + } + } } /// A message received from a transport, plus the means to settle it. diff --git a/src/bus/sql_bus_common.rs b/src/bus/sql_bus_common.rs index 5c3a689e..24ccbc25 100644 --- a/src/bus/sql_bus_common.rs +++ b/src/bus/sql_bus_common.rs @@ -37,6 +37,8 @@ use std::sync::atomic::{AtomicI64, Ordering}; use std::sync::Arc; use std::time::Duration; +use tokio::sync::Notify; + use sqlx::{ColumnIndex, Decode, Row, Type}; use crate::projection_protocol::{ProjectionEpoch, ProjectionSource}; @@ -304,6 +306,22 @@ pub trait SqlBusDialect: Clone + Send + Sync + 'static { limit: i64, ) -> impl Future, TransportError>> + Send; + /// Cheap read: is there at least one claimable row? Empty must not take a + /// writer lock — an idle supervisor must not `UPDATE` the queue file. + fn has_claimable( + &self, + names: &[String], + ) -> impl Future> + Send; + + /// Block until another process may have enqueued work. Combined with the + /// in-process `Notify` in [`SqlBus`]. The default is a short sleep. + fn listen_wakeup(&self) -> impl Future> + Send { + async { + tokio::time::sleep(Duration::from_millis(50)).await; + Ok(()) + } + } + /// Read up to `limit` `bus_log` entries past `consumer`'s offset, in `seq` /// order, whose `name` matches one of `names` **or is NULL** (surfaced, not /// silently skipped, so the failure policy advances the offset past poison @@ -354,6 +372,7 @@ pub struct SqlBus { topology: BusTopologyConfig, lease: Duration, source_epoch: Option, + wake: Arc, } impl SqlBus { @@ -363,6 +382,7 @@ impl SqlBus { topology: BusTopologyConfig::default(), lease: DEFAULT_LEASE, source_epoch: None, + wake: Arc::new(Notify::new()), } } @@ -445,14 +465,18 @@ impl SqlBus { impl Bus for SqlBus { async fn send_message(&self, message: Message) -> Result<(), TransportError> { - self.dialect.insert_queue(&message).await + self.dialect.insert_queue(&message).await?; + self.wake.notify_waiters(); + Ok(()) } async fn publish_message(&self, message: Message) -> Result<(), TransportError> { let epoch_candidate = self.source_epoch.clone().unwrap_or_else(fresh_log_epoch); self.dialect .insert_log(&message, &epoch_candidate, self.source_epoch.as_ref()) - .await + .await?; + self.wake.notify_waiters(); + Ok(()) } } @@ -472,6 +496,7 @@ impl BusConsumer for SqlBus { names, lease_secs: self.lease.as_secs_f64(), buffer: VecDeque::new(), + wake: Arc::clone(&self.wake), }; run_source(router, source, options).await } @@ -502,6 +527,7 @@ impl BusConsumer for SqlBus { last_delivered: None, settled_seq: Arc::new(AtomicI64::new(0)), source_epoch, + wake: Arc::clone(&self.wake), }; run_source(router, source, options).await } @@ -520,6 +546,7 @@ struct SqlQueueSource { names: Vec, lease_secs: f64, buffer: VecDeque, + wake: Arc, } impl MessageSource for SqlQueueSource { @@ -531,6 +558,9 @@ impl MessageSource for SqlQueueSource { async fn recv(&mut self) -> Result, TransportError> { if self.buffer.is_empty() { + if !self.dialect.has_claimable(&self.names).await? { + return Ok(None); + } let mut claimed = self .dialect .claim(&self.names, self.lease_secs, SOURCE_BATCH) @@ -545,6 +575,27 @@ impl MessageSource for SqlQueueSource { claim_token: claimed.claim_token, })) } + + async fn wait(&mut self) -> Result<(), TransportError> { + // Register before peeking so a send that races the empty-queue check + // still wakes this waiter instead of dropping the notify. + let notified = self.wake.notified(); + tokio::pin!(notified); + notified.as_mut().enable(); + if self.dialect.has_claimable(&self.names).await? { + return Ok(()); + } + tokio::select! { + _ = notified => {} + result = self.dialect.listen_wakeup() => { + if let Err(error) = result { + eprintln!("{} bus wakeup: {error}", B::BACKEND); + tokio::time::sleep(Duration::from_millis(50)).await; + } + } + } + Ok(()) + } } /// A claimed `bus_queue` row: `ack` deletes it (done); `nack` makes it available @@ -612,6 +663,7 @@ struct SqlLogSource { /// Highest `seq` settled forward by this source's handles. settled_seq: Arc, source_epoch: ProjectionEpoch, + wake: Arc, } impl MessageSource for SqlLogSource { @@ -670,6 +722,22 @@ impl MessageSource for SqlLogSource { ordered, })) } + + async fn wait(&mut self) -> Result<(), TransportError> { + let notified = self.wake.notified(); + tokio::pin!(notified); + notified.as_mut().enable(); + tokio::select! { + _ = notified => {} + result = self.dialect.listen_wakeup() => { + if let Err(error) = result { + eprintln!("{} bus wakeup: {error}", B::BACKEND); + tokio::time::sleep(Duration::from_millis(50)).await; + } + } + } + Ok(()) + } } /// A `bus_log` entry: `ack` advances this consumer's offset to its `seq` (the diff --git a/src/bus/sqlite_bus.rs b/src/bus/sqlite_bus.rs index 8df2da90..d5437011 100644 --- a/src/bus/sqlite_bus.rs +++ b/src/bus/sqlite_bus.rs @@ -605,6 +605,22 @@ impl SqlBusDialect for SqliteBusDialect { .collect() } + async fn has_claimable(&self, names: &[String]) -> Result { + let mut query = QueryBuilder::::new("SELECT 1 FROM bus_queue WHERE "); + push_name_filter(&mut query, names); + query.push( + " AND available_at <= unixepoch('now','subsec') \ + AND (locked_until IS NULL OR locked_until <= unixepoch('now','subsec')) \ + LIMIT 1", + ); + let row = query + .build() + .fetch_optional(&self.pool) + .await + .map_err(|err| db_err("peek claimable", err))?; + Ok(row.is_some()) + } + async fn log_read( &self, names: &[String], diff --git a/src/microsvc/workers.rs b/src/microsvc/workers.rs index 532337f7..d0bfc33e 100644 --- a/src/microsvc/workers.rs +++ b/src/microsvc/workers.rs @@ -51,7 +51,7 @@ where tokio::spawn(async move { loop { let service = build_service(); - match service.run(RunOptions::idempotent()).await { + match service.run(RunOptions::idempotent().wait_when_idle()).await { Ok(()) => tokio::time::sleep(Duration::from_millis(25)).await, Err(e) => { eprintln!("consumer: {e}"); diff --git a/src/sqlx_repo/repo/commit.rs b/src/sqlx_repo/repo/commit.rs index d0e13f1e..f9d682f7 100644 --- a/src/sqlx_repo/repo/commit.rs +++ b/src/sqlx_repo/repo/commit.rs @@ -49,7 +49,7 @@ where async fn commit_sqlx_batch<'a, DB>( repository: &'a SqlxRepository, - batch: CommitBatch<'a>, + batch: &mut CommitBatch<'a>, mut completion: Option, direct_projection: Option, ) -> Result<(), CommandLedgerError> @@ -127,11 +127,11 @@ where insert_outbox_messages_in_tx(&mut tx, &batch.outbox_messages).await?; let mut changed_tables = std::collections::BTreeSet::new(); - for plan in batch.read_model_plans { + for plan in &batch.read_model_plans { for mutation in &plan.mutations { changed_tables.insert(mutation.table_name().to_string()); } - apply_read_model_write_plan_in_tx(&mut tx, plan) + apply_read_model_write_plan_in_tx(&mut tx, plan.clone()) .await .map_err(RepositoryError::from)?; } @@ -154,10 +154,10 @@ where changed_tables.insert(PROJECTION_CHANGE_NOTIFY_TABLE.to_string()); } - for write in batch.snapshots { + for write in &batch.snapshots { match write { SnapshotWrite::Save { identity, record } => { - save_snapshot_in_tx(&mut tx, &identity, record).await?; + save_snapshot_in_tx(&mut tx, identity, record.clone()).await?; } } } @@ -187,7 +187,7 @@ where tables: changed_tables, }); } - for stream in batch.streams { + for stream in &mut batch.streams { stream.entity.mark_committed(); } Ok(()) @@ -211,16 +211,30 @@ where { fn commit_batch<'a>( &'a self, - batch: CommitBatch<'a>, + mut batch: CommitBatch<'a>, ) -> impl Future> + Send + 'a { async move { - match commit_sqlx_batch(self, batch, None, None).await { - Ok(()) => Ok(()), - Err(CommandLedgerError::Storage(error)) => Err(error), - Err(error) => Err(RepositoryError::Model(format!( - "unexpected command ledger error in ordinary commit: {error}" - ))), + let mut delay = std::time::Duration::from_millis(5); + for attempt in 0..8 { + match commit_sqlx_batch(self, &mut batch, None, None).await { + Ok(()) => return Ok(()), + Err(CommandLedgerError::Storage(error)) + if error.is_retryable() && attempt + 1 < 8 => + { + tokio::time::sleep(delay).await; + delay = (delay * 2).min(std::time::Duration::from_millis(80)); + } + Err(CommandLedgerError::Storage(error)) => return Err(error), + Err(error) => { + return Err(RepositoryError::Model(format!( + "unexpected command ledger error in ordinary commit: {error}" + ))); + } + } } + Err(RepositoryError::Model( + "exhausted retryable storage retries in ordinary commit".into(), + )) } } } @@ -243,14 +257,17 @@ where { fn commit_causal_batch<'a>( &'a self, - batch: CausalCommitBatch<'a>, + mut batch: CausalCommitBatch<'a>, ) -> impl Future> + Send + 'a { - commit_sqlx_batch( - self, - batch.domain, - Some(batch.completion), - batch.direct_projection, - ) + async move { + commit_sqlx_batch( + self, + &mut batch.domain, + Some(batch.completion), + batch.direct_projection, + ) + .await + } } } diff --git a/tests/kafka_transport/main.rs b/tests/kafka_transport/main.rs index 93a26af1..0d377b7c 100644 --- a/tests/kafka_transport/main.rs +++ b/tests/kafka_transport/main.rs @@ -203,6 +203,76 @@ async fn bus_listen_shared_group_consumes_each_command_once() { ); } +/// Subscribe to two command topics, produce only the second. The unused topic +/// must exist (or be created) before subscribe, or Kafka never assigns it and +/// the second command times out — the load-suite hot-increment failure mode. +#[tokio::test] +async fn listen_receives_later_command_topic_after_first_name_is_idle() { + let Some(brokers) = brokers() else { return }; + let ns = unique("ns"); + let bus = KafkaBus::connect(&brokers) + .group("counters") + .namespace(&ns) + .with_fetch_timeout(Duration::from_secs(3)) + .await + .expect("connect"); + + let handled = Arc::new(Mutex::new(Vec::::new())); + let rec = handled.clone(); + let rec_inc = handled.clone(); + let service = Arc::new( + Service::new().routes( + Routes::new() + .with_dependencies(()) + .command("counter.initialize") + .handle(move |ctx: &Context<()>| { + rec.lock() + .unwrap() + .push(ctx.message().id().unwrap_or_default().to_string()); + async move { Ok(json!({})) } + }) + .command("counter.increment") + .handle(move |ctx: &Context<()>| { + rec_inc + .lock() + .unwrap() + .push(ctx.message().id().unwrap_or_default().to_string()); + async move { Ok(json!({})) } + }), + ), + ); + + let listener = bus.clone(); + let listen = tokio::spawn(async move { + listener + .listen(service, RunOptions::idempotent().wait_when_idle()) + .await + }); + tokio::time::sleep(Duration::from_secs(2)).await; + + bus.send_message( + Message::new("counter.increment", MessageKind::Command, b"{}".to_vec()).with_id("inc-1"), + ) + .await + .expect("send increment"); + + let deadline = tokio::time::Instant::now() + Duration::from_secs(10); + loop { + if handled.lock().unwrap().iter().any(|id| id == "inc-1") { + break; + } + if tokio::time::Instant::now() >= deadline { + listen.abort(); + panic!( + "increment command was not consumed; handled={:?}", + handled.lock().unwrap() + ); + } + tokio::time::sleep(Duration::from_millis(50)).await; + } + listen.abort(); +} + /// Build a namespaced `KafkaBus` for `group` (empty `group` = no group). async fn kafka_bus(brokers: &str, ns: &str, group: &str) -> KafkaBus { let builder = KafkaBus::connect(brokers).namespace(ns); diff --git a/tests/load/.gitignore b/tests/load/.gitignore new file mode 100644 index 00000000..ea8c4bf7 --- /dev/null +++ b/tests/load/.gitignore @@ -0,0 +1 @@ +/target diff --git a/tests/load/Cargo.toml b/tests/load/Cargo.toml new file mode 100644 index 00000000..73269b94 --- /dev/null +++ b/tests/load/Cargo.toml @@ -0,0 +1,40 @@ +[package] +name = "distributed_load" +version = "0.1.0" +edition = "2021" +license = "MIT" +publish = false +description = "Opt-in aggregate-command load host and client. Not part of the default workspace test suite." + +[[bin]] +name = "load-host" +path = "src/bin/load-host.rs" + +[[bin]] +name = "load-client" +path = "src/bin/load-client.rs" + +[[bin]] +name = "load-suite" +path = "src/bin/load-suite.rs" + +[features] +default = [] +kafka = ["distributed/kafka"] +rabbitmq = ["distributed/rabbitmq"] + +[dependencies] +axum = "0.8" +distributed = { path = "../..", features = ["http", "grpc", "sqlite", "postgres", "nats"] } +serde = { version = "1", features = ["derive"] } +serde_json = "1" +tokio = { version = "1", features = ["rt-multi-thread", "macros", "net", "time", "sync", "signal"] } +tokio-stream = "0.1" +tonic = { version = "0.14", default-features = false, features = ["transport"] } +reqwest = { version = "0.13", default-features = false, features = ["json", "rustls"] } +sqlx = { version = "0.9", default-features = false, features = ["runtime-tokio", "sqlite", "postgres"] } +uuid = { version = "1", features = ["v7"] } + +[dev-dependencies] +http-body-util = "0.1" +tower = { version = "0.5", features = ["util"] } diff --git a/tests/load/src/bin/load-client.rs b/tests/load/src/bin/load-client.rs new file mode 100644 index 00000000..10f30dc7 --- /dev/null +++ b/tests/load/src/bin/load-client.rs @@ -0,0 +1,97 @@ +//! Drive the Counter HTTP API and print a JSON throughput/latency report. +//! +//! ```text +//! cargo run --manifest-path tests/load/Cargo.toml --release --bin load-client -- \ +//! --url http://127.0.0.1:8790 --scenario unique-create --concurrency 32 --duration 15s +//! ``` + +use std::time::Duration; + +use distributed_load::{run_client, ClientConfig, Scenario}; + +#[tokio::main] +async fn main() -> Result<(), Box> { + let config = parse_args(std::env::args().skip(1))?; + let report = run_client(config).await?; + println!("{}", serde_json::to_string_pretty(&report)?); + if report.err > 0 && report.ok == 0 { + std::process::exit(1); + } + Ok(()) +} + +fn parse_args(args: impl IntoIterator) -> Result { + let mut config = ClientConfig::default(); + let mut args = args.into_iter(); + while let Some(arg) = args.next() { + match arg.as_str() { + "--help" | "-h" => { + print_help(); + std::process::exit(0); + } + "--url" => config.url = expect_value(&arg, args.next())?, + "--scenario" => config.scenario = Scenario::parse(&expect_value(&arg, args.next())?)?, + "--concurrency" => config.concurrency = parse_usize(&expect_value(&arg, args.next())?)?, + "--duration" => config.duration = parse_duration(&expect_value(&arg, args.next())?)?, + "--warmup" => config.warmup = parse_duration(&expect_value(&arg, args.next())?)?, + "--repo" => config.repo = Some(expect_value(&arg, args.next())?), + "--snapshots" => { + let n: u64 = expect_value(&arg, args.next())? + .parse() + .map_err(|_| "--snapshots requires an integer > 0")?; + if n == 0 { + return Err("snapshot frequency must be > 0".into()); + } + config.snapshot_frequency = Some(n); + } + other => return Err(format!("unknown argument {other}")), + } + } + Ok(config) +} + +fn expect_value(flag: &str, value: Option) -> Result { + value.ok_or_else(|| format!("{flag} requires a value")) +} + +fn parse_usize(raw: &str) -> Result { + raw.parse() + .map_err(|_| format!("not a positive integer: {raw}")) +} + +fn parse_duration(raw: &str) -> Result { + if let Some(secs) = raw.strip_suffix('s') { + return parse_f64_secs(secs); + } + if let Some(mins) = raw.strip_suffix('m') { + let mins: f64 = mins.parse().map_err(|_| format!("not a duration: {raw}"))?; + return Ok(Duration::from_secs_f64(mins * 60.0)); + } + parse_f64_secs(raw) +} + +fn parse_f64_secs(raw: &str) -> Result { + let secs: f64 = raw.parse().map_err(|_| format!("not a duration: {raw}"))?; + if secs < 0.0 { + return Err("duration must be >= 0".into()); + } + Ok(Duration::from_secs_f64(secs)) +} + +fn print_help() { + eprintln!( + "\ +load-client — HTTP load driver for load-host + +Options: + --url URL Host base URL (default: http://127.0.0.1:8790) + --scenario unique-create|hot-increment + unique-create: new counter id per request (no contention) + hot-increment: one id, every request increments it + --concurrency N Parallel workers (default: 32) + --duration 15s Measured window after warmup (default: 15s) + --warmup 2s Unmeasured warmup (default: 2s) + --repo NAME Copied into the JSON report only + --snapshots N Copied into the JSON report (host --snapshots N)" + ); +} diff --git a/tests/load/src/bin/load-host.rs b/tests/load/src/bin/load-host.rs new file mode 100644 index 00000000..fc344fe6 --- /dev/null +++ b/tests/load/src/bin/load-host.rs @@ -0,0 +1,91 @@ +//! Serve the Counter aggregate HTTP API for load tests. +//! +//! ```text +//! cargo run --manifest-path tests/load/Cargo.toml --release --bin load-host -- \ +//! --repo memory --bind 127.0.0.1:8790 +//! ``` + +use std::path::PathBuf; + +use distributed_load::host::{bind_listener, serve_listener, CounterService, HostConfig}; +use distributed_load::{LockKind, RepoKind}; + +#[tokio::main] +async fn main() -> Result<(), Box> { + let config = parse_args(std::env::args().skip(1))?; + let host = CounterService::start(&config).await?; + let listener = bind_listener(&config.bind).await?; + let addr = listener.local_addr()?; + eprintln!( + "load-host repo={} lock={} snapshots={} commands=[{}, {}] listening on http://{addr}", + config.repo.as_str(), + config.lock.as_str(), + config + .snapshot_frequency + .map(|n| n.to_string()) + .unwrap_or_else(|| "off".into()), + distributed_load::host::INITIALIZE, + distributed_load::host::INCREMENT, + ); + serve_listener(host.service, listener).await?; + Ok(()) +} + +fn parse_args(args: impl IntoIterator) -> Result { + let mut config = HostConfig::default(); + let mut args = args.into_iter(); + while let Some(arg) = args.next() { + match arg.as_str() { + "--help" | "-h" => { + print_help(); + std::process::exit(0); + } + "--repo" => { + let value = expect_value(&arg, args.next())?; + config.repo = RepoKind::parse(&value)?; + } + "--lock" => { + let value = expect_value(&arg, args.next())?; + config.lock = LockKind::parse(&value)?; + } + "--bind" => config.bind = expect_value(&arg, args.next())?, + "--database-url" => config.database_url = Some(expect_value(&arg, args.next())?), + "--sqlite-path" => config.sqlite_path = PathBuf::from(expect_value(&arg, args.next())?), + "--snapshots" => { + let n: u64 = expect_value(&arg, args.next())? + .parse() + .map_err(|_| "--snapshots requires an integer > 0")?; + if n == 0 { + return Err("snapshot frequency must be > 0".into()); + } + config.snapshot_frequency = Some(n); + } + other => return Err(format!("unknown argument {other}")), + } + } + Ok(config) +} + +fn expect_value(flag: &str, value: Option) -> Result { + value.ok_or_else(|| format!("{flag} requires a value")) +} + +fn print_help() { + eprintln!( + "\ +load-host — Counter aggregate HTTP command API + +Options: + --repo memory|sqlite|postgres Persistence backend (default: memory) + --lock memory|sqlite|postgres Queued lock manager (default: memory) + --snapshots N Enable with_snapshots(N); omit for no snapshots + --bind HOST:PORT Listen address (default: 127.0.0.1:8790) + --database-url URL Postgres URL (default: $DATABASE_URL or compose default) + --sqlite-path PATH SQLite file (default: target/load.sqlite); recreated on start + +Commands: + POST /counter.initialize {{\"id\":\"...\"}} + POST /counter.increment {{\"id\":\"...\",\"amount\":1}} + GET /health" + ); +} diff --git a/tests/load/src/bin/load-suite.rs b/tests/load/src/bin/load-suite.rs new file mode 100644 index 00000000..6c973611 --- /dev/null +++ b/tests/load/src/bin/load-suite.rs @@ -0,0 +1,178 @@ +//! Run the Counter aggregate load matrix. +//! +//! ```text +//! cargo run --manifest-path tests/load/Cargo.toml --release --bin load-suite -- \ +//! --duration 5s --concurrency 16 +//! ``` + +use std::time::Duration; + +use distributed_load::{default_cells, run_suite, Scenario, SuiteConfig}; + +#[tokio::main] +async fn main() -> Result<(), Box> { + let (config, filter, snapshots_only) = parse_args(std::env::args().skip(1))?; + let mut cells = default_cells(&config); + if snapshots_only { + cells.retain(|cell| cell.snapshot_frequency.is_some()); + } + if let Some(filter) = filter { + cells.retain(|cell| cell.name().contains(&filter)); + } + if cells.is_empty() { + return Err("no cells matched".into()); + } + eprintln!("running {} cells", cells.len()); + let outcomes = run_suite(config, cells).await; + println!("{}", serde_json::to_string_pretty(&outcomes)?); + if outcomes.iter().any(|o| o.error.is_some()) { + std::process::exit(1); + } + Ok(()) +} + +fn parse_args( + args: impl IntoIterator, +) -> Result<(SuiteConfig, Option, bool), String> { + let mut config = SuiteConfig::default(); + let mut filter = None; + let mut snapshots_only = false; + let mut args = args.into_iter(); + while let Some(arg) = args.next() { + match arg.as_str() { + "--help" | "-h" => { + print_help(); + std::process::exit(0); + } + "--duration" => config.duration = parse_duration(&expect_value(&arg, args.next())?)?, + "--warmup" => config.warmup = parse_duration(&expect_value(&arg, args.next())?)?, + "--concurrency" => { + config.concurrency = expect_value(&arg, args.next())? + .parse() + .map_err(|_| "concurrency must be an integer")? + } + "--database-url" => config.database_url = Some(expect_value(&arg, args.next())?), + "--no-locks" => config.include_locks = false, + "--locks-only" => { + config.include_locks = true; + } + "--no-snapshots" => config.include_snapshots = false, + "--snapshots-only" => { + config.include_snapshots = true; + snapshots_only = true; + } + "--snapshot-frequency" => { + config.snapshot_frequencies = parse_frequencies(&expect_value(&arg, args.next())?)?; + config.include_snapshots = true; + } + "--snapshot-dispatch" => { + config.snapshot_dispatches = + parse_snapshot_dispatches(&expect_value(&arg, args.next())?)?; + config.include_snapshots = true; + } + "--no-external" => config.include_external = false, + "--scenario" => { + config.scenarios = vec![Scenario::parse(&expect_value(&arg, args.next())?)?]; + } + "--both-scenarios" => { + config.scenarios = vec![Scenario::UniqueCreate, Scenario::HotIncrement]; + } + "--filter" => filter = Some(expect_value(&arg, args.next())?), + other => return Err(format!("unknown argument {other}")), + } + } + Ok((config, filter, snapshots_only)) +} + +fn parse_snapshot_dispatches(raw: &str) -> Result, String> { + use distributed_load::DispatchKind; + let mut out = Vec::new(); + for part in raw.split(',') { + let kind = DispatchKind::parse(part.trim())?; + if kind == DispatchKind::Bus { + return Err( + "snapshot cells do not overlay bus (bus cells have their own pairing); \ + use --snapshot-dispatch direct,http,grpc" + .into(), + ); + } + if !out.contains(&kind) { + out.push(kind); + } + } + if out.is_empty() { + return Err("--snapshot-dispatch needs at least one of direct,http,grpc".into()); + } + Ok(out) +} + +fn parse_frequencies(raw: &str) -> Result, String> { + let freqs: Result, _> = raw + .split(',') + .map(|part| { + let n: u64 = part + .trim() + .parse() + .map_err(|_| format!("not a snapshot frequency: {part}"))?; + if n == 0 { + return Err("snapshot frequency must be > 0".to_string()); + } + Ok(n) + }) + .collect(); + let freqs = freqs?; + if freqs.is_empty() { + return Err("--snapshot-frequency needs at least one value".into()); + } + Ok(freqs) +} + +fn expect_value(flag: &str, value: Option) -> Result { + value.ok_or_else(|| format!("{flag} requires a value")) +} + +fn parse_duration(raw: &str) -> Result { + if let Some(secs) = raw.strip_suffix('s') { + return Ok(Duration::from_secs_f64( + secs.parse().map_err(|_| format!("not a duration: {raw}"))?, + )); + } + if let Some(mins) = raw.strip_suffix('m') { + let mins: f64 = mins.parse().map_err(|_| format!("not a duration: {raw}"))?; + return Ok(Duration::from_secs_f64(mins * 60.0)); + } + Ok(Duration::from_secs_f64( + raw.parse().map_err(|_| format!("not a duration: {raw}"))?, + )) +} + +fn print_help() { + eprintln!( + "\ +load-suite — Counter aggregate matrix (direct, http, grpc, bus, locks, snapshots) + +Options: + --duration 5s Measured window (default: 5s) + --warmup 1s Unmeasured warmup (default: 1s) + --concurrency N Workers per cell (default: 16) + --database-url URL Postgres URL (default: $DATABASE_URL) + --scenario unique-create|hot-increment + --both-scenarios Both scenarios (default) + --no-locks Skip suite 1a lock-manager cells + --no-snapshots Skip snapshot-frequency overlay + --snapshots-only Only snapshot-frequency cells + --snapshot-frequency 1,10,100 + Frequencies for with_snapshots(n) (default: 1,10,100) + --snapshot-dispatch direct,http,grpc + Overlay snapshot frequencies on these ingress modes + (default: direct,http,grpc; buses are always overlayed) + --no-external Skip NATS/Kafka/RabbitMQ even if env is set + --filter SUBSTRING Keep cells whose name contains SUBSTRING + +Env: + DATABASE_URL postgres://sourced:sourced@localhost:5432/distributed + NATS_URL nats://localhost:4222 + KAFKA_BROKERS 127.0.0.1:9092 (requires --features kafka) + AMQP_URL amqp://guest:guest@localhost:5672/%2f (requires --features rabbitmq)" + ); +} diff --git a/tests/load/src/client.rs b/tests/load/src/client.rs new file mode 100644 index 00000000..bf2c55e0 --- /dev/null +++ b/tests/load/src/client.rs @@ -0,0 +1,213 @@ +//! Load driver over any [`Invoker`]. + +use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; +use std::sync::Arc; +use std::time::{Duration, Instant}; + +use tokio::sync::Mutex; + +use crate::invoke::{command_body, setup_hot_id, BusInvoker, Invoker}; +use crate::stats::{summarize, RunReport}; + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum Scenario { + UniqueCreate, + HotIncrement, +} + +impl Scenario { + pub fn parse(raw: &str) -> Result { + match raw { + "unique-create" | "create" => Ok(Self::UniqueCreate), + "hot-increment" | "increment" => Ok(Self::HotIncrement), + other => Err(format!( + "unknown --scenario {other:?} (expected unique-create or hot-increment)" + )), + } + } + + pub fn as_str(self) -> &'static str { + match self { + Self::UniqueCreate => "unique-create", + Self::HotIncrement => "hot-increment", + } + } +} + +#[derive(Clone, Debug)] +pub struct ClientConfig { + pub url: String, + pub scenario: Scenario, + pub concurrency: usize, + pub duration: Duration, + pub warmup: Duration, + pub repo: Option, + pub dispatch: Option, + pub bus: Option, + pub lock: Option, + pub snapshot_frequency: Option, + pub cell: Option, + pub pipelined: bool, +} + +impl Default for ClientConfig { + fn default() -> Self { + Self { + url: "http://127.0.0.1:8790".into(), + scenario: Scenario::UniqueCreate, + concurrency: 32, + duration: Duration::from_secs(15), + warmup: Duration::from_secs(2), + repo: None, + dispatch: None, + bus: None, + lock: None, + snapshot_frequency: None, + cell: None, + pipelined: false, + } + } +} + +pub async fn run_client( + config: ClientConfig, +) -> Result> { + let client = reqwest::Client::builder() + .pool_max_idle_per_host(config.concurrency.max(1)) + .build()?; + let invoker = Invoker::Http { + client, + base: config.url.trim_end_matches('/').to_string(), + }; + run_invoker(invoker, config).await +} + +pub async fn run_invoker( + invoker: Invoker, + config: ClientConfig, +) -> Result> { + if config.concurrency == 0 { + return Err("concurrency must be >= 1".into()); + } + let hot_id = if config.scenario == Scenario::HotIncrement { + Some(setup_hot_id(&invoker).await?) + } else { + None + }; + + let measuring = Arc::new(AtomicBool::new(false)); + let stop = Arc::new(AtomicBool::new(false)); + let ok = Arc::new(AtomicU64::new(0)); + let err = Arc::new(AtomicU64::new(0)); + let samples = Arc::new(Mutex::new(Vec::::new())); + + let mut workers = Vec::with_capacity(config.concurrency); + for _ in 0..config.concurrency { + let invoker = invoker.clone(); + let scenario = config.scenario; + let hot_id = hot_id.clone(); + let measuring = Arc::clone(&measuring); + let stop = Arc::clone(&stop); + let ok = Arc::clone(&ok); + let err = Arc::clone(&err); + let samples = Arc::clone(&samples); + workers.push(tokio::spawn(async move { + while !stop.load(Ordering::Relaxed) { + let (command, body) = command_body(scenario, hot_id.as_deref()); + let started = Instant::now(); + let result = invoker.invoke(&command, body).await; + let elapsed_ms = started.elapsed().as_secs_f64() * 1000.0; + if !measuring.load(Ordering::Relaxed) { + continue; + } + match result { + Ok(()) => { + ok.fetch_add(1, Ordering::Relaxed); + samples.lock().await.push(elapsed_ms); + } + Err(_) => { + err.fetch_add(1, Ordering::Relaxed); + } + } + } + })); + } + + tokio::time::sleep(config.warmup).await; + let pipelined_baseline = if config.pipelined { + if let Invoker::Bus(bus) = &invoker { + Some(( + bus.applied_ok.load(Ordering::Relaxed), + bus.applied_err.load(Ordering::Relaxed), + )) + } else { + None + } + } else { + None + }; + measuring.store(true, Ordering::Relaxed); + let started = Instant::now(); + tokio::time::sleep(config.duration).await; + let elapsed = started.elapsed(); + stop.store(true, Ordering::Relaxed); + for worker in workers { + let _ = worker.await; + } + + let (ok, err) = if let Some((baseline_ok, baseline_err)) = pipelined_baseline { + if let Invoker::Bus(bus) = &invoker { + drain_pipeline(bus, Duration::from_secs(8)).await; + ( + bus.applied_ok + .load(Ordering::Relaxed) + .saturating_sub(baseline_ok), + bus.applied_err + .load(Ordering::Relaxed) + .saturating_sub(baseline_err), + ) + } else { + (ok.load(Ordering::Relaxed), err.load(Ordering::Relaxed)) + } + } else { + (ok.load(Ordering::Relaxed), err.load(Ordering::Relaxed)) + }; + let mut samples = samples.lock().await; + let latency_ms = summarize(&mut samples); + let secs = elapsed.as_secs_f64().max(0.000_001); + Ok(RunReport { + cell: config.cell, + scenario: config.scenario.as_str().into(), + url: config.url, + repo: config.repo, + dispatch: config.dispatch, + bus: config.bus, + lock: config.lock, + snapshot_frequency: config.snapshot_frequency, + concurrency: config.concurrency, + duration_secs: (secs * 1000.0).round() / 1000.0, + warmup_secs: config.warmup.as_secs_f64(), + ok, + err, + throughput_rps: (ok as f64 / secs * 10.0).round() / 10.0, + latency_ms, + }) +} + +/// Wait until applied counts stop moving, so pipelined enqueue is not +/// under-counted just because the measure window ended. +async fn drain_pipeline(bus: &BusInvoker, budget: Duration) { + let deadline = Instant::now() + budget; + let mut last = bus.applied_ok.load(Ordering::Relaxed) + bus.applied_err.load(Ordering::Relaxed); + let mut last_change = Instant::now(); + while Instant::now() < deadline { + tokio::time::sleep(Duration::from_millis(25)).await; + let now = bus.applied_ok.load(Ordering::Relaxed) + bus.applied_err.load(Ordering::Relaxed); + if now != last { + last = now; + last_change = Instant::now(); + } else if last_change.elapsed() >= Duration::from_secs(1) { + break; + } + } +} diff --git a/tests/load/src/counter.rs b/tests/load/src/counter.rs new file mode 100644 index 00000000..a6a7bd3a --- /dev/null +++ b/tests/load/src/counter.rs @@ -0,0 +1,57 @@ +//! Minimal Counter aggregate used as the load-test write-path fixture. + +use distributed::{sourced, Entity, OutboxMessage, Snapshot, SourcedResult}; +use serde::{Deserialize, Serialize}; + +#[derive(Default, Snapshot)] +pub struct Counter { + pub entity: Entity, + pub value: i64, +} + +#[sourced(entity, aggregate_type = "load.counter")] +impl Counter { + #[event("initialized")] + pub fn create(&mut self, id: String) { + self.entity.set_id(&id); + self.value = 0; + } + + #[event("incremented")] + pub fn increment(&mut self, amount: i64) { + self.value += amount; + } +} + +#[derive(Serialize, Deserialize)] +pub struct CreateCounter { + pub id: String, +} + +#[derive(Serialize, Deserialize)] +pub struct IncrementCounter { + pub id: String, + pub amount: i64, +} + +#[derive(Serialize)] +struct CounterState<'a> { + id: &'a str, + value: i64, +} + +pub fn counter_state_message(counter: &Counter, event_type: &str) -> SourcedResult { + OutboxMessage::encode_for_entity( + format!( + "{}:{event_type}:{}", + counter.entity.id(), + counter.entity.version() + ), + event_type, + &CounterState { + id: counter.entity.id(), + value: counter.value, + }, + &counter.entity, + ) +} diff --git a/tests/load/src/host.rs b/tests/load/src/host.rs new file mode 100644 index 00000000..5c466e89 --- /dev/null +++ b/tests/load/src/host.rs @@ -0,0 +1,358 @@ +//! Counter service builder: persistence + lock manager + HTTP helpers. + +use std::path::{Path, PathBuf}; +use std::sync::Arc; + +use distributed::microsvc::{self, Context, HandlerError, HasOutboxStore, Routes, Service}; +use distributed::{ + AggregateBuilder, AggregateRepository, GetStream, InMemoryLockManager, InMemoryRepository, + LockManager, PostgresLockManager, PostgresRepository, Queueable, QueuedRepository, + SnapshotStore, SqliteLockManager, SqliteRepository, TransactionalCommit, +}; +use serde_json::{json, Value}; +use tokio::net::TcpListener; + +use crate::counter::{counter_state_message, Counter, CreateCounter, IncrementCounter}; +use crate::kinds::{compatible_lock, LockKind, RepoKind}; + +pub const INITIALIZE: &str = "counter.initialize"; +pub const INCREMENT: &str = "counter.increment"; + +type CounterRepo = AggregateRepository, Counter>; + +#[derive(Clone, Debug)] +pub struct HostConfig { + pub repo: RepoKind, + pub lock: LockKind, + pub bind: String, + pub database_url: Option, + pub sqlite_path: PathBuf, + /// `None` disables snapshot caching. `Some(n)` calls `with_snapshots(n)`. + pub snapshot_frequency: Option, +} + +impl Default for HostConfig { + fn default() -> Self { + Self { + repo: RepoKind::Memory, + lock: LockKind::Memory, + bind: "127.0.0.1:8790".into(), + database_url: None, + sqlite_path: PathBuf::from("target/load.sqlite"), + snapshot_frequency: None, + } + } +} + +pub struct BuiltService { + pub service: Arc, + pub sqlite: Option, + pub postgres: Option, +} + +pub struct CounterService { + pub service: Arc, + pub repo: RepoKind, +} + +impl CounterService { + pub async fn start( + config: &HostConfig, + ) -> Result> { + let built = build_service(config).await?; + Ok(Self { + service: built.service, + repo: config.repo, + }) + } +} + +pub fn build_service_from_sqlite( + inner: SqliteRepository, + lock: LockKind, + snapshot_frequency: Option, +) -> Result> { + compatible_lock(RepoKind::Sqlite, lock)?; + let service = match lock { + LockKind::Memory => Arc::new(service_for( + inner.clone(), + InMemoryLockManager::new(), + snapshot_frequency, + )), + LockKind::Sqlite => Arc::new(service_for( + inner.clone(), + SqliteLockManager::new(inner.pool().clone()), + snapshot_frequency, + )), + LockKind::Postgres => { + return Err("lock=postgres is not valid with repo=sqlite".into()); + } + }; + Ok(BuiltService { + service, + sqlite: Some(inner), + postgres: None, + }) +} + +pub async fn build_service( + config: &HostConfig, +) -> Result> { + compatible_lock(config.repo, config.lock)?; + match config.repo { + RepoKind::Memory => { + let inner = InMemoryRepository::new(); + Ok(BuiltService { + service: Arc::new(service_for( + inner, + InMemoryLockManager::new(), + config.snapshot_frequency, + )), + sqlite: None, + postgres: None, + }) + } + RepoKind::Sqlite => { + let pool_size = if config.lock == LockKind::Sqlite { + 4 + } else { + 1 + }; + let inner = connect_sqlite(&config.sqlite_path, pool_size).await?; + let service = match config.lock { + LockKind::Memory => Arc::new(service_for( + inner.clone(), + InMemoryLockManager::new(), + config.snapshot_frequency, + )), + LockKind::Sqlite => Arc::new(service_for( + inner.clone(), + SqliteLockManager::new(inner.pool().clone()), + config.snapshot_frequency, + )), + LockKind::Postgres => unreachable!("compatible_lock rejects this"), + }; + Ok(BuiltService { + service, + sqlite: Some(inner), + postgres: None, + }) + } + RepoKind::Postgres => { + let url = postgres_url(config); + let inner = PostgresRepository::connect_and_migrate(&url).await?; + let service = match config.lock { + LockKind::Memory => Arc::new(service_for( + inner.clone(), + InMemoryLockManager::new(), + config.snapshot_frequency, + )), + LockKind::Postgres => Arc::new(service_for( + inner.clone(), + PostgresLockManager::new(inner.pool().clone()), + config.snapshot_frequency, + )), + LockKind::Sqlite => unreachable!("compatible_lock rejects this"), + }; + Ok(BuiltService { + service, + sqlite: None, + postgres: Some(inner), + }) + } + } +} + +pub fn postgres_url(config: &HostConfig) -> String { + config + .database_url + .clone() + .or_else(|| std::env::var("DATABASE_URL").ok()) + .unwrap_or_else(|| "postgres://sourced:sourced@localhost:5432/distributed".into()) +} + +pub fn sqlite_pool_size(lock: LockKind, needs_sql_bus: bool) -> u32 { + if lock == LockKind::Sqlite || needs_sql_bus { + 4 + } else { + 1 + } +} + +pub async fn connect_sqlite( + path: &Path, + pool_size: u32, +) -> Result> { + use sqlx::sqlite::{ + SqliteConnectOptions, SqliteJournalMode, SqlitePoolOptions, SqliteSynchronous, + }; + use std::str::FromStr; + + if let Some(parent) = path.parent() { + if !parent.as_os_str().is_empty() { + std::fs::create_dir_all(parent)?; + } + } + let _ = std::fs::remove_file(path); + let _ = std::fs::remove_file(path.with_extension("sqlite-wal")); + let _ = std::fs::remove_file(path.with_extension("sqlite-shm")); + + let url = format!("sqlite:{}?mode=rwc", path.display()); + let options = SqliteConnectOptions::from_str(&url)? + .create_if_missing(true) + .journal_mode(SqliteJournalMode::Wal) + .synchronous(SqliteSynchronous::Normal) + .busy_timeout(std::time::Duration::from_secs(5)); + let pool = SqlitePoolOptions::new() + .max_connections(pool_size) + .connect_with(options) + .await?; + let repo = SqliteRepository::new(pool); + repo.migrate().await?; + Ok(repo) +} + +fn service_for(inner: R, locks: L, snapshot_frequency: Option) -> Service +where + R: Clone + + GetStream + + TransactionalCommit + + HasOutboxStore + + SnapshotStore + + Send + + Sync + + 'static, + L: LockManager + Send + Sync + 'static, + QueuedRepository: Clone + + GetStream + + TransactionalCommit + + HasOutboxStore + + SnapshotStore + + Send + + Sync + + 'static, +{ + Service::new() + .named("load-counter") + .with_http_command_routes() + .routes(counter_routes(inner, locks, snapshot_frequency)) +} + +fn counter_repo(inner: R, locks: L, snapshot_frequency: Option) -> CounterRepo +where + R: Clone + GetStream + TransactionalCommit + SnapshotStore + Send + Sync + 'static, + L: LockManager + Send + Sync + 'static, + QueuedRepository: + Clone + GetStream + TransactionalCommit + SnapshotStore + Send + Sync + 'static, +{ + let repo = inner.queued_with(locks).aggregate::(); + match snapshot_frequency { + Some(frequency) => repo.with_snapshots(frequency), + None => repo, + } +} + +fn counter_routes( + inner: R, + locks: L, + snapshot_frequency: Option, +) -> Routes> +where + R: Clone + + GetStream + + TransactionalCommit + + HasOutboxStore + + SnapshotStore + + Send + + Sync + + 'static, + L: LockManager + Send + Sync + 'static, + QueuedRepository: Clone + + GetStream + + TransactionalCommit + + HasOutboxStore + + SnapshotStore + + Send + + Sync + + 'static, +{ + Routes::new() + .with_repo(counter_repo(inner, locks, snapshot_frequency)) + .command(INITIALIZE) + .guarded(guard_create::, handle_create::) + .command(INCREMENT) + .guarded(guard_increment::, handle_increment::) +} + +fn guard_create(ctx: &Context>) -> bool { + ctx.has_fields(&["id"]) +} + +fn guard_increment(ctx: &Context>) -> bool { + ctx.has_fields(&["id", "amount"]) +} + +async fn handle_create(ctx: &Context<'_, CounterRepo>) -> Result +where + QueuedRepository: GetStream + TransactionalCommit, +{ + let input = ctx.input::()?; + if ctx.repo().get(&input.id).await?.is_some() { + return Err(HandlerError::Rejected(format!( + "counter {} already exists", + input.id + ))); + } + + let mut counter = Counter::default(); + counter.create(input.id.clone())?; + let message = counter_state_message(&counter, "counter.initialized")?; + ctx.repo().outbox(message).commit(&mut counter).await?; + Ok(json!({ "id": input.id })) +} + +async fn handle_increment(ctx: &Context<'_, CounterRepo>) -> Result +where + QueuedRepository: GetStream + TransactionalCommit, +{ + let input = ctx.input::()?; + let mut counter: Counter = ctx + .repo() + .get(&input.id) + .await? + .ok_or_else(|| HandlerError::NotFound(input.id.clone()))?; + counter.increment(input.amount)?; + let message = counter_state_message(&counter, "counter.incremented")?; + ctx.repo().outbox(message).commit(&mut counter).await?; + Ok(json!({ "id": input.id, "value": counter.value })) +} + +pub async fn bind_listener(addr: &str) -> Result { + TcpListener::bind(addr).await +} + +pub async fn serve_listener( + service: Arc, + listener: TcpListener, +) -> Result<(), std::io::Error> { + axum::serve(listener, microsvc::router(service)).await +} + +pub async fn wait_for_health(base: &str, timeout: std::time::Duration) -> Result<(), String> { + let client = reqwest::Client::new(); + let deadline = tokio::time::Instant::now() + timeout; + let url = format!("{base}/health"); + loop { + if let Ok(resp) = client.get(&url).send().await { + if resp.status().is_success() { + return Ok(()); + } + } + if tokio::time::Instant::now() >= deadline { + return Err(format!( + "health check did not succeed at {url} within {timeout:?}" + )); + } + tokio::time::sleep(std::time::Duration::from_millis(50)).await; + } +} diff --git a/tests/load/src/invoke.rs b/tests/load/src/invoke.rs new file mode 100644 index 00000000..de0670b3 --- /dev/null +++ b/tests/load/src/invoke.rs @@ -0,0 +1,345 @@ +//! Command invokers: direct dispatch, HTTP, and bus send+completion. + +use std::collections::HashMap; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::{Arc, Mutex}; +use std::time::Duration; + +#[cfg(feature = "kafka")] +use distributed::bus::KafkaBus; +#[cfg(feature = "rabbitmq")] +use distributed::bus::RabbitBus; +use distributed::bus::{ + Bus, BusConsumer, InMemoryBus, Message, MessageKind, MessageRouter, NatsBus, OrderedDelivery, + PostgresBus, RunOptions, SqliteBus, SubscriptionPlan, TransportError, +}; +use distributed::microsvc::grpc::{CommandServiceClient, GrpcRequest}; +use distributed::microsvc::{Service, Session}; +use serde_json::Value; +use tokio::sync::{oneshot, Notify}; +use tonic::transport::Channel; +use uuid::Uuid; + +use crate::host::{INCREMENT, INITIALIZE}; + +pub type InvokeError = String; + +#[derive(Clone)] +pub enum Invoker { + Direct(Arc), + Http { + client: reqwest::Client, + base: String, + }, + Grpc(CommandServiceClient), + Bus(BusInvoker), +} + +impl Invoker { + pub async fn invoke(&self, command: &str, body: Value) -> Result<(), InvokeError> { + match self { + Self::Direct(service) => service + .dispatch(command, body, Session::new()) + .await + .map(|_| ()) + .map_err(|e| e.to_string()), + Self::Http { client, base } => { + let resp = client + .post(format!("{base}/{command}")) + .json(&body) + .send() + .await + .map_err(|e| e.to_string())?; + if resp.status().is_success() { + Ok(()) + } else { + Err(format!( + "{} {}", + resp.status(), + resp.text().await.unwrap_or_default() + )) + } + } + Self::Grpc(client) => { + let mut client = client.clone(); + let resp = client + .dispatch(GrpcRequest { + command: command.to_string(), + input: body.to_string(), + session_variables: Default::default(), + }) + .await + .map_err(|e| e.to_string())? + .into_inner(); + if resp.status == 200 { + Ok(()) + } else { + Err(format!("grpc status {} {}", resp.status, resp.body)) + } + } + Self::Bus(bus) => bus.invoke(command, body).await, + } + } +} + +#[derive(Clone)] +pub struct BusInvoker { + live: LiveBus, + pending: CompletionMap, + notify: Arc, + pipelined: bool, + pub applied_ok: Arc, + pub applied_err: Arc, +} + +type CompletionMap = Arc>>>>; + +impl BusInvoker { + async fn invoke(&self, command: &str, body: Value) -> Result<(), InvokeError> { + let id = Uuid::now_v7().to_string(); + let rx = if self.pipelined { + None + } else { + let (tx, rx) = oneshot::channel(); + self.pending + .lock() + .expect("completion map") + .insert(id.clone(), tx); + Some(rx) + }; + let payload = serde_json::to_vec(&body).map_err(|e| e.to_string())?; + let message = Message::new(command, MessageKind::Command, payload).with_id(id); + if let Err(e) = self.live.send(message).await { + return Err(e.to_string()); + } + self.notify.notify_waiters(); + if self.pipelined { + return Ok(()); + } + let rx = rx.expect("applied mode registers a completion channel"); + match tokio::time::timeout(Duration::from_secs(15), rx).await { + Ok(Ok(result)) => result, + Ok(Err(_)) => Err("bus completion dropped".into()), + Err(_) => Err("bus completion timed out".into()), + } + } +} + +#[derive(Clone)] +enum LiveBus { + Memory(InMemoryBus), + Sqlite(SqliteBus), + Postgres(PostgresBus), + Nats(NatsBus), + #[cfg(feature = "kafka")] + Kafka(KafkaBus), + #[cfg(feature = "rabbitmq")] + Rabbit(std::sync::Arc), +} + +impl LiveBus { + async fn send(&self, message: Message) -> Result<(), TransportError> { + match self { + Self::Memory(bus) => bus.send_message(message).await, + Self::Sqlite(bus) => bus.send_message(message).await, + Self::Postgres(bus) => bus.send_message(message).await, + Self::Nats(bus) => bus.send_message(message).await, + #[cfg(feature = "kafka")] + Self::Kafka(bus) => bus.send_message(message).await, + #[cfg(feature = "rabbitmq")] + Self::Rabbit(bus) => bus.send_message(message).await, + } + } + + async fn listen( + &self, + router: Arc, + options: RunOptions, + ) -> Result<(), TransportError> { + match self { + Self::Memory(bus) => bus.listen(router, options).await, + Self::Sqlite(bus) => bus.listen(router, options).await, + Self::Postgres(bus) => bus.listen(router, options).await, + Self::Nats(bus) => bus.listen(router, options).await, + #[cfg(feature = "kafka")] + Self::Kafka(bus) => bus.listen(router, options).await, + #[cfg(feature = "rabbitmq")] + Self::Rabbit(bus) => bus.listen(router, options).await, + } + } +} + +struct CompletingRouter { + service: Arc, + pending: CompletionMap, + applied_ok: Arc, + applied_err: Arc, +} + +impl MessageRouter for CompletingRouter { + fn consumer_group(&self) -> Option<&str> { + self.service.consumer_group() + } + + fn handles(&self, kind: MessageKind, name: &str) -> bool { + self.service.handles_message(kind, name) + } + + fn subscription_plan(&self) -> SubscriptionPlan { + self.service.subscription_plan() + } + + async fn dispatch(&self, message: &Message) -> Result<(), TransportError> { + let result = MessageRouter::dispatch(self.service.as_ref(), message).await; + complete( + &self.pending, + &self.applied_ok, + &self.applied_err, + message, + &result, + ); + result + } + + async fn dispatch_ordered( + &self, + message: &Message, + ordered: Option<&OrderedDelivery>, + ) -> Result<(), TransportError> { + let result = self.service.dispatch_ordered(message, ordered).await; + complete( + &self.pending, + &self.applied_ok, + &self.applied_err, + message, + &result, + ); + result + } +} + +fn complete( + pending: &CompletionMap, + applied_ok: &AtomicU64, + applied_err: &AtomicU64, + message: &Message, + result: &Result<(), TransportError>, +) { + match result { + Ok(()) => { + applied_ok.fetch_add(1, Ordering::Relaxed); + } + Err(_) => { + applied_err.fetch_add(1, Ordering::Relaxed); + } + } + let Some(id) = message.id() else { + return; + }; + if let Some(tx) = pending.lock().expect("completion map").remove(id) { + let mapped = result.as_ref().map(|_| ()).map_err(|e| e.to_string()); + let _ = tx.send(mapped); + } +} + +pub struct BusRuntime { + pub invoker: BusInvoker, + consumer: tokio::task::JoinHandle<()>, +} + +impl BusRuntime { + pub fn start(live: impl Into, service: Arc, pipelined: bool) -> Self { + let live = live.into().0; + let pending: CompletionMap = Arc::new(Mutex::new(HashMap::new())); + let notify = Arc::new(Notify::new()); + let applied_ok = Arc::new(AtomicU64::new(0)); + let applied_err = Arc::new(AtomicU64::new(0)); + let router = Arc::new(CompletingRouter { + service, + pending: Arc::clone(&pending), + applied_ok: Arc::clone(&applied_ok), + applied_err: Arc::clone(&applied_err), + }); + let consumer_bus = live.clone(); + let consumer = tokio::spawn(async move { + if let Err(e) = consumer_bus + .listen(router, RunOptions::idempotent().wait_when_idle()) + .await + { + eprintln!("load-suite bus consumer: {e}"); + } + }); + Self { + invoker: BusInvoker { + live, + pending, + notify, + pipelined, + applied_ok, + applied_err, + }, + consumer, + } + } + + pub fn stop(&self) { + self.consumer.abort(); + self.invoker.notify.notify_waiters(); + } +} + +pub struct LiveBusWrap(LiveBus); + +impl From for LiveBusWrap { + fn from(bus: InMemoryBus) -> Self { + Self(LiveBus::Memory(bus)) + } +} +impl From for LiveBusWrap { + fn from(bus: SqliteBus) -> Self { + Self(LiveBus::Sqlite(bus)) + } +} +impl From for LiveBusWrap { + fn from(bus: PostgresBus) -> Self { + Self(LiveBus::Postgres(bus)) + } +} +impl From for LiveBusWrap { + fn from(bus: NatsBus) -> Self { + Self(LiveBus::Nats(bus)) + } +} +#[cfg(feature = "kafka")] +impl From for LiveBusWrap { + fn from(bus: KafkaBus) -> Self { + Self(LiveBus::Kafka(bus)) + } +} +#[cfg(feature = "rabbitmq")] +impl From for LiveBusWrap { + fn from(bus: RabbitBus) -> Self { + Self(LiveBus::Rabbit(std::sync::Arc::new(bus))) + } +} + +pub async fn setup_hot_id(invoker: &Invoker) -> Result { + let id = format!("hot-{}", Uuid::now_v7()); + invoker + .invoke(INITIALIZE, serde_json::json!({ "id": id })) + .await?; + Ok(id) +} + +pub fn command_body(scenario: crate::client::Scenario, hot_id: Option<&str>) -> (String, Value) { + match scenario { + crate::client::Scenario::UniqueCreate => { + let id = Uuid::now_v7().to_string(); + (INITIALIZE.into(), serde_json::json!({ "id": id })) + } + crate::client::Scenario::HotIncrement => ( + INCREMENT.into(), + serde_json::json!({ "id": hot_id.unwrap_or("hot"), "amount": 1 }), + ), + } +} diff --git a/tests/load/src/kinds.rs b/tests/load/src/kinds.rs new file mode 100644 index 00000000..0a469d85 --- /dev/null +++ b/tests/load/src/kinds.rs @@ -0,0 +1,138 @@ +//! Matrix axes for the aggregate load suite. + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum RepoKind { + Memory, + Sqlite, + Postgres, +} + +impl RepoKind { + pub fn parse(raw: &str) -> Result { + match raw { + "memory" | "mem" | "in-memory" => Ok(Self::Memory), + "sqlite" => Ok(Self::Sqlite), + "postgres" | "pg" | "postgresql" => Ok(Self::Postgres), + other => Err(format!( + "unknown --repo {other:?} (expected memory, sqlite, or postgres)" + )), + } + } + + pub fn as_str(self) -> &'static str { + match self { + Self::Memory => "memory", + Self::Sqlite => "sqlite", + Self::Postgres => "postgres", + } + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum DispatchKind { + Direct, + Http, + Grpc, + Bus, +} + +impl DispatchKind { + pub fn parse(raw: &str) -> Result { + match raw { + "direct" => Ok(Self::Direct), + "http" => Ok(Self::Http), + "grpc" => Ok(Self::Grpc), + "bus" => Ok(Self::Bus), + other => Err(format!( + "unknown --dispatch {other:?} (expected direct, http, grpc, or bus)" + )), + } + } + + pub fn as_str(self) -> &'static str { + match self { + Self::Direct => "direct", + Self::Http => "http", + Self::Grpc => "grpc", + Self::Bus => "bus", + } + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum BusKind { + Memory, + Sqlite, + Postgres, + Nats, + Kafka, + Rabbitmq, +} + +impl BusKind { + pub fn parse(raw: &str) -> Result { + match raw { + "memory" | "mem" | "in-memory" => Ok(Self::Memory), + "sqlite" => Ok(Self::Sqlite), + "postgres" | "pg" => Ok(Self::Postgres), + "nats" => Ok(Self::Nats), + "kafka" => Ok(Self::Kafka), + "rabbitmq" | "amqp" | "rabbit" => Ok(Self::Rabbitmq), + other => Err(format!( + "unknown --bus {other:?} (expected memory, sqlite, postgres, nats, kafka, or rabbitmq)" + )), + } + } + + pub fn as_str(self) -> &'static str { + match self { + Self::Memory => "memory", + Self::Sqlite => "sqlite", + Self::Postgres => "postgres", + Self::Nats => "nats", + Self::Kafka => "kafka", + Self::Rabbitmq => "rabbitmq", + } + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum LockKind { + Memory, + Sqlite, + Postgres, +} + +impl LockKind { + pub fn parse(raw: &str) -> Result { + match raw { + "memory" | "mem" | "in-memory" => Ok(Self::Memory), + "sqlite" => Ok(Self::Sqlite), + "postgres" | "pg" => Ok(Self::Postgres), + other => Err(format!( + "unknown --lock {other:?} (expected memory, sqlite, or postgres)" + )), + } + } + + pub fn as_str(self) -> &'static str { + match self { + Self::Memory => "memory", + Self::Sqlite => "sqlite", + Self::Postgres => "postgres", + } + } +} + +pub fn compatible_lock(repo: RepoKind, lock: LockKind) -> Result<(), String> { + match (repo, lock) { + (_, LockKind::Memory) => Ok(()), + (RepoKind::Sqlite, LockKind::Sqlite) => Ok(()), + (RepoKind::Postgres, LockKind::Postgres) => Ok(()), + (repo, lock) => Err(format!( + "lock={} is not valid with repo={} (sqlite lock needs sqlite repo, postgres lock needs postgres repo)", + lock.as_str(), + repo.as_str() + )), + } +} diff --git a/tests/load/src/lib.rs b/tests/load/src/lib.rs new file mode 100644 index 00000000..38018bd5 --- /dev/null +++ b/tests/load/src/lib.rs @@ -0,0 +1,166 @@ +//! Opt-in aggregate-command load harness. +//! +//! Not a workspace member — run via `--manifest-path tests/load/Cargo.toml`. + +pub mod client; +pub mod counter; +pub mod host; +pub mod invoke; +pub mod kinds; +pub mod stats; +pub mod suite; + +pub use client::{run_client, ClientConfig, Scenario}; +pub use host::{bind_listener, serve_listener, CounterService, HostConfig}; +pub use kinds::{BusKind, DispatchKind, LockKind, RepoKind}; +pub use stats::{percentile_ms, RunReport}; +pub use suite::{default_cells, run_suite, Cell, SuiteConfig}; + +#[cfg(test)] +mod smoke_tests { + use super::*; + use crate::host::{wait_for_health, INCREMENT, INITIALIZE}; + use crate::invoke::Invoker; + use distributed::microsvc::Session; + use serde_json::json; + use std::sync::Arc; + use std::time::Duration; + + #[tokio::test] + async fn memory_host_initialize_and_increment() { + let host = CounterService::start(&HostConfig::default()) + .await + .expect("start memory host"); + let listener = bind_listener("127.0.0.1:0").await.expect("bind"); + let addr = listener.local_addr().expect("addr"); + let base = format!("http://{addr}"); + tokio::spawn(async move { + serve_listener(host.service, listener).await.unwrap(); + }); + wait_for_health(&base, Duration::from_secs(5)) + .await + .expect("health"); + + let client = reqwest::Client::new(); + let created = client + .post(format!("{base}/{INITIALIZE}")) + .json(&json!({ "id": "c1" })) + .send() + .await + .expect("create send"); + assert_eq!(created.status(), 200); + + let incremented = client + .post(format!("{base}/{INCREMENT}")) + .json(&json!({ "id": "c1", "amount": 3 })) + .send() + .await + .expect("increment send"); + assert_eq!(incremented.status(), 200); + let body: serde_json::Value = incremented.json().await.expect("json"); + assert_eq!(body["value"], 3); + } + + #[tokio::test] + async fn memory_direct_dispatch_and_in_memory_bus() { + let host = CounterService::start(&HostConfig::default()) + .await + .expect("start"); + host.service + .dispatch(INITIALIZE, json!({ "id": "direct-1" }), Session::new()) + .await + .expect("direct initialize"); + + let bus = distributed::bus::InMemoryBus::new(); + let runtime = crate::invoke::BusRuntime::start(bus, Arc::clone(&host.service), false); + let invoker = Invoker::Bus(runtime.invoker.clone()); + invoker + .invoke(INITIALIZE, json!({ "id": "bus-1" })) + .await + .expect("bus initialize"); + runtime.stop(); + } + + #[tokio::test] + async fn memory_grpc_initialize_and_increment() { + let host = CounterService::start(&HostConfig::default()) + .await + .expect("start"); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0") + .await + .expect("bind"); + let addr = listener.local_addr().expect("addr"); + let grpc_svc = distributed::microsvc::grpc_server(host.service); + tokio::spawn(async move { + tonic::transport::Server::builder() + .add_service(grpc_svc) + .serve_with_incoming(tokio_stream::wrappers::TcpListenerStream::new(listener)) + .await + .unwrap(); + }); + let endpoint = format!("http://{addr}"); + let mut client = None; + for _ in 0..50 { + if let Ok(c) = + distributed::microsvc::grpc::CommandServiceClient::connect(endpoint.clone()).await + { + client = Some(c); + break; + } + tokio::time::sleep(Duration::from_millis(20)).await; + } + let client = client.expect("grpc connect"); + let invoker = Invoker::Grpc(client); + invoker + .invoke(INITIALIZE, json!({ "id": "grpc-1" })) + .await + .expect("grpc initialize"); + invoker + .invoke(INCREMENT, json!({ "id": "grpc-1", "amount": 4 })) + .await + .expect("grpc increment"); + } + + #[tokio::test] + async fn sqlite_bus_completes_one_initialize() { + let path = + std::env::temp_dir().join(format!("load-sqlite-bus-{}.sqlite", uuid::Uuid::now_v7())); + let inner = crate::host::connect_sqlite(&path, 4).await.expect("sqlite"); + let built = crate::host::build_service_from_sqlite(inner.clone(), LockKind::Memory, None) + .expect("service"); + let plan = built.service.subscription_plan(); + assert!( + plan.commands.iter().any(|n| n == INITIALIZE), + "expected {INITIALIZE} in {:?}", + plan.commands + ); + let bus = distributed::bus::SqliteBus::new(inner.pool().clone()).group("load-test"); + bus.ensure_tables().await.expect("tables"); + let runtime = crate::invoke::BusRuntime::start(bus, Arc::clone(&built.service), false); + let invoker = Invoker::Bus(runtime.invoker.clone()); + invoker + .invoke(INITIALIZE, json!({ "id": "sqlite-bus-1" })) + .await + .expect("bus initialize"); + runtime.stop(); + } + + #[tokio::test] + async fn memory_snapshots_every_event() { + let mut config = HostConfig::default(); + config.snapshot_frequency = Some(1); + let host = CounterService::start(&config).await.expect("start"); + host.service + .dispatch(INITIALIZE, json!({ "id": "snap-1" }), Session::new()) + .await + .expect("initialize"); + host.service + .dispatch( + INCREMENT, + json!({ "id": "snap-1", "amount": 1 }), + Session::new(), + ) + .await + .expect("increment"); + } +} diff --git a/tests/load/src/stats.rs b/tests/load/src/stats.rs new file mode 100644 index 00000000..10059429 --- /dev/null +++ b/tests/load/src/stats.rs @@ -0,0 +1,72 @@ +//! Latency percentile helpers and the JSON run report. + +use serde::Serialize; + +#[derive(Clone, Debug, Serialize)] +pub struct RunReport { + #[serde(skip_serializing_if = "Option::is_none")] + pub cell: Option, + pub scenario: String, + pub url: String, + pub repo: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub dispatch: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub bus: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub lock: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub snapshot_frequency: Option, + pub concurrency: usize, + pub duration_secs: f64, + pub warmup_secs: f64, + pub ok: u64, + pub err: u64, + pub throughput_rps: f64, + pub latency_ms: LatencySummary, +} + +#[derive(Clone, Debug, Serialize)] +pub struct LatencySummary { + pub p50: f64, + pub p95: f64, + pub p99: f64, + pub max: f64, +} + +fn round3(value: f64) -> f64 { + (value * 1000.0).round() / 1000.0 +} + +/// `samples` is a sorted list of latencies in milliseconds. +pub fn percentile_ms(sorted_ms: &[f64], pct: f64) -> f64 { + if sorted_ms.is_empty() { + return 0.0; + } + let rank = ((pct / 100.0) * (sorted_ms.len() as f64 - 1.0)).round() as usize; + round3(sorted_ms[rank.min(sorted_ms.len() - 1)]) +} + +pub fn summarize(samples_ms: &mut [f64]) -> LatencySummary { + samples_ms.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal)); + LatencySummary { + p50: percentile_ms(samples_ms, 50.0), + p95: percentile_ms(samples_ms, 95.0), + p99: percentile_ms(samples_ms, 99.0), + max: round3(samples_ms.last().copied().unwrap_or(0.0)), + } +} + +#[cfg(test)] +mod tests { + use super::percentile_ms; + + #[test] + fn percentile_picks_nearest_rank() { + let samples = [1.0, 2.0, 3.0, 4.0, 5.0]; + assert_eq!(percentile_ms(&samples, 0.0), 1.0); + assert_eq!(percentile_ms(&samples, 50.0), 3.0); + assert_eq!(percentile_ms(&samples, 100.0), 5.0); + assert_eq!(percentile_ms(&[], 50.0), 0.0); + } +} diff --git a/tests/load/src/suite.rs b/tests/load/src/suite.rs new file mode 100644 index 00000000..d6a3f943 --- /dev/null +++ b/tests/load/src/suite.rs @@ -0,0 +1,613 @@ +//! Default aggregate load matrix and runner. + +use std::path::PathBuf; +use std::sync::Arc; +use std::time::Duration; + +use distributed::bus::{InMemoryBus, NatsBus, PostgresBus, SqliteBus}; +use serde::Serialize; +use uuid::Uuid; + +use crate::client::{run_invoker, ClientConfig, Scenario}; +use crate::host::{ + bind_listener, build_service, connect_sqlite, serve_listener, sqlite_pool_size, + wait_for_health, HostConfig, +}; +use crate::invoke::{BusRuntime, Invoker}; +use crate::kinds::{BusKind, DispatchKind, LockKind, RepoKind}; +use distributed::microsvc::grpc::CommandServiceClient; +use tokio::net::TcpListener; +use tokio_stream::wrappers::TcpListenerStream; + +#[derive(Clone, Debug)] +pub struct Cell { + pub repo: RepoKind, + pub dispatch: DispatchKind, + pub bus: Option, + pub lock: LockKind, + pub scenario: Scenario, + pub snapshot_frequency: Option, + pub pipelined: bool, +} + +impl Cell { + fn base( + repo: RepoKind, + dispatch: DispatchKind, + bus: Option, + lock: LockKind, + scenario: Scenario, + ) -> Self { + Self { + repo, + dispatch, + bus, + lock, + scenario, + snapshot_frequency: None, + pipelined: false, + } + } + + pub fn name(&self) -> String { + let bus = self + .bus + .map(|b| format!(" bus={}", b.as_str())) + .unwrap_or_default(); + let snap = match self.snapshot_frequency { + Some(n) => format!(" snap={n}"), + None => " snap=off".into(), + }; + let mode = if self.pipelined { + " mode=pipelined" + } else if self.dispatch == DispatchKind::Bus { + " mode=applied" + } else { + "" + }; + format!( + "repo={} dispatch={}{bus} lock={}{snap}{mode} scenario={}", + self.repo.as_str(), + self.dispatch.as_str(), + self.lock.as_str(), + self.scenario.as_str() + ) + } +} + +#[derive(Clone, Debug)] +pub struct SuiteConfig { + pub duration: Duration, + pub warmup: Duration, + pub concurrency: usize, + pub database_url: Option, + pub include_locks: bool, + pub include_external: bool, + pub include_snapshots: bool, + pub snapshot_frequencies: Vec, + /// Dispatch modes to overlay snapshot frequencies on. Default is direct + /// only so snapshot I/O is not mixed with HTTP/gRPC cost. Pass + /// `--snapshot-dispatch direct,http,grpc` to widen the experiment. + pub snapshot_dispatches: Vec, + pub scenarios: Vec, +} + +impl Default for SuiteConfig { + fn default() -> Self { + Self { + duration: Duration::from_secs(5), + warmup: Duration::from_secs(1), + concurrency: 16, + database_url: None, + include_locks: true, + include_external: true, + include_snapshots: true, + snapshot_frequencies: vec![1, 10, 100], + snapshot_dispatches: vec![DispatchKind::Direct, DispatchKind::Http, DispatchKind::Grpc], + scenarios: vec![Scenario::UniqueCreate, Scenario::HotIncrement], + } + } +} + +#[derive(Clone, Debug, Serialize)] +pub struct CellOutcome { + pub cell: String, + #[serde(skip_serializing_if = "Option::is_none")] + pub skipped: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub error: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub report: Option, +} + +/// Suite 1: persistence × (direct|http|grpc) + paired bus cells, snapshots off. +/// Suite 1a: same persistence × (direct|http|grpc) with matching durable locks. +/// Suite 1b: every suite 1 / 1a / bus cell × snapshot frequencies {1,10,100}. +/// `--snapshot-dispatch` can narrow the ingress overlay; buses always +/// get the same frequencies when snapshots are enabled. +pub fn default_cells(config: &SuiteConfig) -> Vec { + let mut bases = Vec::new(); + for scenario in &config.scenarios { + for repo in [RepoKind::Memory, RepoKind::Sqlite, RepoKind::Postgres] { + for dispatch in [DispatchKind::Direct, DispatchKind::Http, DispatchKind::Grpc] { + bases.push(Cell::base( + repo, + dispatch, + None, + LockKind::Memory, + *scenario, + )); + } + } + bases.extend([ + bus_cell(RepoKind::Memory, BusKind::Memory, *scenario), + bus_cell(RepoKind::Sqlite, BusKind::Sqlite, *scenario), + bus_cell(RepoKind::Postgres, BusKind::Postgres, *scenario), + bus_cell(RepoKind::Postgres, BusKind::Nats, *scenario), + bus_cell(RepoKind::Postgres, BusKind::Kafka, *scenario), + bus_cell(RepoKind::Postgres, BusKind::Rabbitmq, *scenario), + ]); + if config.include_locks { + for dispatch in [DispatchKind::Direct, DispatchKind::Http, DispatchKind::Grpc] { + bases.push(Cell::base( + RepoKind::Sqlite, + dispatch, + None, + LockKind::Sqlite, + *scenario, + )); + bases.push(Cell::base( + RepoKind::Postgres, + dispatch, + None, + LockKind::Postgres, + *scenario, + )); + } + } + } + let mut cells = bases.clone(); + if config.include_snapshots { + for base in &bases { + if !overlays_snapshots(base, config) { + continue; + } + for &frequency in &config.snapshot_frequencies { + let mut cell = base.clone(); + cell.snapshot_frequency = Some(frequency); + cells.push(cell); + } + } + } + let mut pipelined = Vec::new(); + for cell in &cells { + if cell.dispatch == DispatchKind::Bus + && cell.snapshot_frequency.is_none() + && cell.scenario == Scenario::UniqueCreate + && !cell.pipelined + { + let mut next = cell.clone(); + next.pipelined = true; + pipelined.push(next); + } + } + cells.extend(pipelined); + cells +} + +fn overlays_snapshots(cell: &Cell, config: &SuiteConfig) -> bool { + if cell.dispatch == DispatchKind::Bus { + return true; + } + config.snapshot_dispatches.contains(&cell.dispatch) +} + +fn bus_cell(repo: RepoKind, bus: BusKind, scenario: Scenario) -> Cell { + Cell::base( + repo, + DispatchKind::Bus, + Some(bus), + LockKind::Memory, + scenario, + ) +} + +pub async fn run_suite(config: SuiteConfig, cells: Vec) -> Vec { + let mut outcomes = Vec::with_capacity(cells.len()); + for cell in cells { + let name = cell.name(); + eprintln!("======== {name} ========"); + if let Some(reason) = skip_reason(&cell, &config) { + eprintln!("skip: {reason}"); + outcomes.push(CellOutcome { + cell: name, + skipped: Some(reason), + error: None, + report: None, + }); + continue; + } + match run_cell(&config, &cell).await { + Ok(report) => { + eprintln!( + "ok={} err={} rps={} p50={} p99={}", + report.ok, + report.err, + report.throughput_rps, + report.latency_ms.p50, + report.latency_ms.p99 + ); + outcomes.push(CellOutcome { + cell: name, + skipped: None, + error: None, + report: Some(report), + }); + } + Err(e) => { + eprintln!("error: {e}"); + outcomes.push(CellOutcome { + cell: name, + skipped: None, + error: Some(e), + report: None, + }); + } + } + } + outcomes +} + +fn skip_reason(cell: &Cell, config: &SuiteConfig) -> Option { + if matches!(cell.repo, RepoKind::Postgres) + || matches!(cell.bus, Some(BusKind::Postgres)) + || matches!(cell.lock, LockKind::Postgres) + { + // Still try; run_cell will fail with a connect error if postgres is down. + } + if matches!(cell.bus, Some(BusKind::Nats)) { + if !config.include_external { + return Some("external buses disabled (--no-external)".into()); + } + if std::env::var("NATS_URL").is_err() { + return Some("NATS_URL is unset".into()); + } + } + if matches!(cell.bus, Some(BusKind::Kafka)) { + if !cfg!(feature = "kafka") { + return Some("rebuild with --features kafka".into()); + } + if !config.include_external { + return Some("external buses disabled (--no-external)".into()); + } + if std::env::var("KAFKA_BROKERS").is_err() { + return Some("KAFKA_BROKERS is unset".into()); + } + } + if matches!(cell.bus, Some(BusKind::Rabbitmq)) { + if !cfg!(feature = "rabbitmq") { + return Some("rebuild with --features rabbitmq".into()); + } + if !config.include_external { + return Some("external buses disabled (--no-external)".into()); + } + if std::env::var("AMQP_URL").is_err() { + return Some("AMQP_URL is unset".into()); + } + } + None +} + +async fn run_cell(config: &SuiteConfig, cell: &Cell) -> Result { + let sqlite_path = PathBuf::from(format!( + "target/load-{}-{}.sqlite", + cell.repo.as_str(), + Uuid::now_v7() + )); + let needs_sql_bus = matches!(cell.bus, Some(BusKind::Sqlite)); + let host = HostConfig { + repo: cell.repo, + lock: cell.lock, + bind: "127.0.0.1:0".into(), + database_url: config.database_url.clone(), + sqlite_path: sqlite_path.clone(), + snapshot_frequency: cell.snapshot_frequency, + }; + + let built = if cell.repo == RepoKind::Sqlite && needs_sql_bus { + // Rebuild sqlite with a wider pool so the bus consumer can share it. + let inner = connect_sqlite(&sqlite_path, sqlite_pool_size(cell.lock, true)) + .await + .map_err(|e| e.to_string())?; + let mut rebuilt = host.clone(); + rebuilt.lock = cell.lock; + // connect_sqlite already migrated; build_service would wipe the file. + crate::host::build_service_from_sqlite(inner, cell.lock, cell.snapshot_frequency) + .map_err(|e| e.to_string())? + } else { + build_service(&host).await.map_err(|e| e.to_string())? + }; + + let mut client = client_config(config, cell); + if matches!(cell.bus, Some(BusKind::Sqlite)) { + // Concurrent claim+insert+commit on one SQLite file collapses into + // SQLITE_BUSY. One worker is the honest local-SQLite bus number + // for both applied and pipelined cells. + client.concurrency = 1; + } + match cell.dispatch { + DispatchKind::Direct => { + let invoker = Invoker::Direct(Arc::clone(&built.service)); + run_invoker(invoker, client) + .await + .map_err(|e| e.to_string()) + } + DispatchKind::Http => { + let listener = bind_listener("127.0.0.1:0") + .await + .map_err(|e| e.to_string())?; + let addr = listener.local_addr().map_err(|e| e.to_string())?; + let base = format!("http://{addr}"); + let service = Arc::clone(&built.service); + let server = tokio::spawn(async move { + let _ = serve_listener(service, listener).await; + }); + wait_for_health(&base, Duration::from_secs(5)).await?; + let mut client = client; + client.url = base.clone(); + let invoker = Invoker::Http { + client: reqwest::Client::new(), + base, + }; + let report = run_invoker(invoker, client) + .await + .map_err(|e| e.to_string()); + server.abort(); + report + } + DispatchKind::Grpc => { + let (server, grpc_client, endpoint) = start_grpc(Arc::clone(&built.service)).await?; + let mut client = client; + client.url = endpoint; + let invoker = Invoker::Grpc(grpc_client); + let report = run_invoker(invoker, client) + .await + .map_err(|e| e.to_string()); + server.abort(); + report + } + DispatchKind::Bus => { + let bus_kind = cell.bus.ok_or("bus dispatch requires --bus")?; + let (wrap, runtime) = start_bus(config, cell, &host, &built, bus_kind).await?; + let invoker = Invoker::Bus(runtime.invoker.clone()); + let report = run_invoker(invoker, client) + .await + .map_err(|e| e.to_string()); + runtime.stop(); + drop(wrap); + report + } + } +} + +fn client_config(config: &SuiteConfig, cell: &Cell) -> ClientConfig { + ClientConfig { + url: String::new(), + scenario: cell.scenario, + concurrency: config.concurrency, + duration: config.duration, + warmup: config.warmup, + repo: Some(cell.repo.as_str().into()), + dispatch: Some(cell.dispatch.as_str().into()), + bus: cell.bus.map(|b| b.as_str().into()), + lock: Some(cell.lock.as_str().into()), + snapshot_frequency: cell.snapshot_frequency, + cell: Some(cell.name()), + pipelined: cell.pipelined, + } +} + +async fn start_bus( + _config: &SuiteConfig, + cell: &Cell, + _host: &HostConfig, + built: &crate::host::BuiltService, + bus_kind: BusKind, +) -> Result<(BusToken, BusRuntime), String> { + let namespace = format!("load-{}", Uuid::now_v7()); + match bus_kind { + BusKind::Memory => { + let bus = InMemoryBus::new(); + let runtime = + BusRuntime::start(bus.clone(), Arc::clone(&built.service), cell.pipelined); + Ok((BusToken::Memory(bus), runtime)) + } + BusKind::Sqlite => { + let repo = built + .sqlite + .clone() + .ok_or("sqlite bus requires sqlite repo")?; + let bus = SqliteBus::new(repo.pool().clone()).group("load-counter"); + bus.ensure_tables().await.map_err(|e| e.to_string())?; + let runtime = + BusRuntime::start(bus.clone(), Arc::clone(&built.service), cell.pipelined); + Ok((BusToken::Sqlite(bus), runtime)) + } + BusKind::Postgres => { + let repo = built + .postgres + .clone() + .ok_or("postgres bus requires postgres repo")?; + let bus = PostgresBus::new(repo.pool().clone()).group("load-counter"); + bus.ensure_tables().await.map_err(|e| e.to_string())?; + // The compose database is shared across cells and leftover + // `counter.*` queue rows starve applied oneshots (the consumer + // drains stale initialize commands instead of this cell's). + sqlx::query("DELETE FROM bus_queue WHERE name = ANY($1)") + .bind(&[crate::host::INITIALIZE, crate::host::INCREMENT] as &[&str]) + .execute(repo.pool()) + .await + .map_err(|e| e.to_string())?; + let runtime = + BusRuntime::start(bus.clone(), Arc::clone(&built.service), cell.pipelined); + Ok((BusToken::Postgres(bus), runtime)) + } + BusKind::Nats => { + let url = std::env::var("NATS_URL").unwrap_or_else(|_| "nats://localhost:4222".into()); + let bus = NatsBus::connect(&url) + .group("load-counter") + .namespace(&namespace) + .with_fetch_timeout(Duration::from_millis(200)) + .await + .map_err(|e| e.to_string())?; + let runtime = + BusRuntime::start(bus.clone(), Arc::clone(&built.service), cell.pipelined); + // Give JetStream consumers a moment to bind. + tokio::time::sleep(Duration::from_millis(200)).await; + Ok((BusToken::Nats(bus), runtime)) + } + BusKind::Kafka => { + #[cfg(feature = "kafka")] + { + let brokers = + std::env::var("KAFKA_BROKERS").unwrap_or_else(|_| "127.0.0.1:9092".into()); + let bus = distributed::bus::KafkaBus::connect(&brokers) + .group("load-counter") + .namespace(&namespace) + .with_fetch_timeout(Duration::from_secs(2)) + .await + .map_err(|e| e.to_string())?; + let runtime = + BusRuntime::start(bus.clone(), Arc::clone(&built.service), cell.pipelined); + tokio::time::sleep(Duration::from_secs(2)).await; + Ok((BusToken::Kafka(bus), runtime)) + } + #[cfg(not(feature = "kafka"))] + { + let _ = (built, namespace); + Err("rebuild with --features kafka".into()) + } + } + BusKind::Rabbitmq => { + #[cfg(feature = "rabbitmq")] + { + let url = std::env::var("AMQP_URL") + .unwrap_or_else(|_| "amqp://guest:guest@localhost:5672/%2f".into()); + let bus = distributed::bus::RabbitBus::connect(&url) + .group("load-counter") + .namespace(&namespace) + .await + .map_err(|e| e.to_string())?; + let runtime = BusRuntime::start(bus, Arc::clone(&built.service), cell.pipelined); + tokio::time::sleep(Duration::from_millis(200)).await; + Ok((BusToken::Rabbitmq, runtime)) + } + #[cfg(not(feature = "rabbitmq"))] + { + let _ = (built, namespace); + Err("rebuild with --features rabbitmq".into()) + } + } + } +} + +async fn start_grpc( + service: Arc, +) -> Result< + ( + tokio::task::JoinHandle>, + CommandServiceClient, + String, + ), + String, +> { + let listener = TcpListener::bind("127.0.0.1:0") + .await + .map_err(|e| e.to_string())?; + let addr = listener.local_addr().map_err(|e| e.to_string())?; + let grpc_svc = distributed::microsvc::grpc_server(service); + let server = tokio::spawn(async move { + tonic::transport::Server::builder() + .add_service(grpc_svc) + .serve_with_incoming(TcpListenerStream::new(listener)) + .await + }); + let endpoint = format!("http://{addr}"); + let deadline = tokio::time::Instant::now() + Duration::from_secs(5); + let client = loop { + match CommandServiceClient::connect(endpoint.clone()).await { + Ok(client) => break client, + Err(e) => { + if tokio::time::Instant::now() >= deadline { + server.abort(); + return Err(format!("gRPC connect to {endpoint} failed: {e}")); + } + tokio::time::sleep(Duration::from_millis(50)).await; + } + } + }; + Ok((server, client, endpoint)) +} + +/// Keeps the bus alive for the cell. +pub enum BusToken { + Memory(InMemoryBus), + Sqlite(SqliteBus), + Postgres(PostgresBus), + Nats(NatsBus), + #[cfg(feature = "kafka")] + Kafka(distributed::bus::KafkaBus), + #[cfg(feature = "rabbitmq")] + Rabbitmq, +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::client::Scenario; + + #[test] + fn default_matrix_covers_every_requested_axis() { + let cells = default_cells(&SuiteConfig::default()); + assert!(cells + .iter() + .any(|c| c.bus == Some(BusKind::Rabbitmq) && c.scenario == Scenario::HotIncrement)); + assert!(cells + .iter() + .any(|c| c.bus == Some(BusKind::Kafka) && c.snapshot_frequency == Some(1))); + assert!(cells.iter().any(|c| { + c.dispatch == DispatchKind::Http + && c.lock == LockKind::Postgres + && c.snapshot_frequency == Some(100) + && c.scenario == Scenario::HotIncrement + })); + assert!(cells.iter().any(|c| { + c.dispatch == DispatchKind::Grpc + && c.snapshot_frequency == Some(10) + && c.repo == RepoKind::Sqlite + })); + assert!(cells.iter().any(|c| { + c.pipelined && c.bus == Some(BusKind::Nats) && c.scenario == Scenario::UniqueCreate + })); + } + + #[test] + fn snapshot_dispatch_can_narrow_ingress_but_still_overlays_buses() { + let config = SuiteConfig { + include_locks: false, + include_snapshots: true, + snapshot_frequencies: vec![100], + snapshot_dispatches: vec![DispatchKind::Direct], + scenarios: vec![Scenario::HotIncrement], + ..SuiteConfig::default() + }; + let snaps: Vec<_> = default_cells(&config) + .into_iter() + .filter(|cell| cell.snapshot_frequency.is_some()) + .collect(); + assert!(snaps.iter().any(|c| c.dispatch == DispatchKind::Direct)); + assert!(snaps.iter().any(|c| c.dispatch == DispatchKind::Bus)); + assert!(snaps.iter().all(|c| c.dispatch != DispatchKind::Http)); + assert!(snaps.iter().all(|c| c.snapshot_frequency == Some(100))); + } +}