feat: aggregate load suite and bus idle/batch fixes

Cargo.toml

@@ -1,6 +1,11 @@
 [workspace]
 members = ["distributed_macros", "distributed_cli"]
-exclude = ["tests/e2e-ui", "tests/celld/worker", "tests/celld/relay-worker"]
+exclude = [
+    "tests/e2e-ui",
+    "tests/celld/worker",
+    "tests/celld/relay-worker",
+    "tests/load",
+]
 resolver = "2"
 
 [workspace.package]

Makefile

@@ -51,6 +51,94 @@ 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) \
+		$(if $(LOAD_DATABASE_URL),--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).

src/bus/in_memory_bus.rs

@@ -14,6 +14,7 @@ use std::collections::{HashMap, VecDeque};
 use std::sync::{Arc, Mutex};
 
 use super::source::{MessageSource, ReceivedMessage};
+use super::wake::Notify;
 use super::{run_source, Bus, BusConsumer, MessageRouter, RunOptions, TransportError};
 use super::{Message, OrderedDelivery};
 use crate::projection_protocol::{ProjectionEpoch, ProjectionSource};
@@ -35,6 +36,7 @@ pub struct InMemoryBus {
     queues: Queues,
     topics: Topics,
     source_epoch: ProjectionEpoch,
+    wake: Arc<Notify>,
 }
 
 impl Default for InMemoryBus {
@@ -44,6 +46,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 +63,7 @@ impl InMemoryBus {
             .entry(message.name().to_string())
             .or_default()
             .push_back(message);
+        self.wake.notify_waiters();
         Ok(())
     }
 
@@ -79,6 +83,7 @@ impl InMemoryBus {
             .entry(message.name().to_string())
             .or_default()
             .push(message);
+        self.wake.notify_waiters();
         Ok(())
     }
 
@@ -138,6 +143,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 +159,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 +169,17 @@ impl BusConsumer for InMemoryBus {
 struct QueueSource {
     queues: Queues,
     names: Vec<String>,
+    wake: Arc<Notify>,
+}
+
+impl QueueSource {
+    fn has_available(&self) -> Result<bool, TransportError> {
+        let queues = self.queues.lock().map_err(|_| lock_poisoned("queue"))?;
+        Ok(self
+            .names
+            .iter()
+            .any(|name| queues.get(name).is_some_and(|queue| !queue.is_empty())))
+    }
 }
 
 impl MessageSource for QueueSource {
@@ -184,6 +202,15 @@ impl MessageSource for QueueSource {
         }
         Ok(None)
     }
+
+    async fn wait(&mut self) -> Result<(), TransportError> {
+        let notified = self.wake.notified();
+        if self.has_available()? {
+            return Ok(());
+        }
+        notified.await;
+        Ok(())
+    }
 }
 
 /// Fan-out source over the named retained logs: each `TopicSource` has its own
@@ -193,6 +220,22 @@ struct TopicSource {
     names: Vec<String>,
     cursors: TopicCursors,
     source_epoch: ProjectionEpoch,
+    wake: Arc<Notify>,
+}
+
+impl TopicSource {
+    fn has_available(&self) -> Result<bool, TransportError> {
+        let topics = self.topics.lock().map_err(|_| lock_poisoned("topic"))?;
+        let cursors = self
+            .cursors
+            .lock()
+            .map_err(|_| lock_poisoned("topic cursor"))?;
+        Ok(self.names.iter().any(|name| {
+            topics
+                .get(name)
+                .is_some_and(|log| cursors.get(name).copied().unwrap_or_default() < log.len())
+        }))
+    }
 }
 
 impl MessageSource for TopicSource {
@@ -238,6 +281,15 @@ impl MessageSource for TopicSource {
         }
         Ok(None)
     }
+
+    async fn wait(&mut self) -> Result<(), TransportError> {
+        let notified = self.wake.notified();
+        if self.has_available()? {
+            return Ok(());
+        }
+        notified.await;
+        Ok(())
+    }
 }
 
 struct TopicSettlement {
@@ -301,6 +353,7 @@ mod tests {
     use crate::bus::{Handlers, MessageKind};
     use crate::trace_context::{CAUSATION_ID, TRACEPARENT};
     use std::future::Future;
+    use std::time::Duration;
 
     fn block_on<F: Future>(future: F) -> F::Output {
         use std::ptr;
@@ -375,10 +428,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).
@@ -423,6 +478,54 @@ mod tests {
         assert_eq!(b_ids, vec!["e0", "e1", "e2"]);
     }
 
+    #[tokio::test]
+    async fn one_publish_wakes_every_waiting_topic_subscriber() {
+        let bus = InMemoryBus::new();
+        let first = recorder();
+        let second = recorder();
+        let first_task = {
+            let bus = bus.clone();
+            let service = event_service(first.clone());
+            tokio::spawn(async move {
+                bus.subscribe(service, RunOptions::idempotent().wait_when_idle())
+                    .await
+            })
+        };
+        let second_task = {
+            let bus = bus.clone();
+            let service = event_service(second.clone());
+            tokio::spawn(async move {
+                bus.subscribe(service, RunOptions::idempotent().wait_when_idle())
+                    .await
+            })
+        };
+        tokio::time::timeout(Duration::from_secs(1), async {
+            while bus.wake.waiter_count() < 2 {
+                tokio::task::yield_now().await;
+            }
+        })
+        .await
+        .expect("both subscribers registered their idle waiters");
+
+        bus.publish_message(
+            Message::new("evt", MessageKind::Event, b"{}".to_vec()).with_id("broadcast"),
+        )
+        .await
+        .unwrap();
+        tokio::time::timeout(Duration::from_secs(1), async {
+            while first.lock().unwrap().is_empty() || second.lock().unwrap().is_empty() {
+                tokio::task::yield_now().await;
+            }
+        })
+        .await
+        .expect("both waiting subscribers handled the publish");
+
+        assert_eq!(first.lock().unwrap().as_slice(), ["broadcast"]);
+        assert_eq!(second.lock().unwrap().as_slice(), ["broadcast"]);
+        first_task.abort();
+        second_task.abort();
+    }
+
     #[test]
     fn topic_nack_redelivers_exact_gap_free_position_and_ack_advances() {
         let bus = InMemoryBus::new();
@@ -437,6 +540,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 +586,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"));