diff --git a/.github/workflows/bump-and-release.yml b/.github/workflows/bump-and-release.yml index a6c2460..c30beae 100644 --- a/.github/workflows/bump-and-release.yml +++ b/.github/workflows/bump-and-release.yml @@ -37,7 +37,7 @@ jobs: outputs: tag: ${{ steps.tag.outputs.tag }} steps: - - uses: actions/checkout@v6 + - uses: actions/checkout@v7 with: # PAT so the bump commit can be pushed to protected master. # GITHUB_TOKEN (github-actions[bot]) cannot bypass branch rules. @@ -120,7 +120,7 @@ jobs: if: ${{ !inputs.beta || inputs.force_release }} runs-on: ubuntu-latest steps: - - uses: actions/checkout@v6 + - uses: actions/checkout@v7 with: ref: ${{ needs.bump.outputs.tag }} diff --git a/.github/workflows/checks.yml b/.github/workflows/checks.yml index 40e4d04..f5c7c2b 100644 --- a/.github/workflows/checks.yml +++ b/.github/workflows/checks.yml @@ -19,19 +19,19 @@ jobs: fmt: runs-on: ubuntu-latest steps: - - uses: actions/checkout@v6 + - uses: actions/checkout@v7 - run: cargo fmt --all --check clippy: runs-on: ubuntu-latest steps: - - uses: actions/checkout@v6 + - uses: actions/checkout@v7 - uses: Swatinem/rust-cache@v2 - run: cargo clippy --all-targets --all-features -- -D warnings test: runs-on: ubuntu-latest steps: - - uses: actions/checkout@v6 + - uses: actions/checkout@v7 - uses: Swatinem/rust-cache@v2 - run: cargo test --all-features diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 6b3b01d..bfdea76 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -16,32 +16,21 @@ concurrency: cancel-in-progress: true jobs: - # fmt, clippy and doc run on the dev toolchain, auto-installed by rustup from + # fmt and clippy run on the dev toolchain, auto-installed by rustup from # rust-toolchain.toml. Only the test matrix pins explicit versions. fmt: runs-on: ubuntu-latest steps: - - uses: actions/checkout@v6 + - uses: actions/checkout@v7 - run: cargo fmt --all --check clippy: runs-on: ubuntu-latest steps: - - uses: actions/checkout@v6 + - uses: actions/checkout@v7 - uses: Swatinem/rust-cache@v2 - run: cargo clippy --all-targets --all-features -- -D warnings - doc: - runs-on: ubuntu-latest - # The crate denies broken intra-doc links; this is the only place rustdoc - # runs, so enforce that (and every other rustdoc warning) here. - env: - RUSTDOCFLAGS: -D warnings - steps: - - uses: actions/checkout@v6 - - uses: Swatinem/rust-cache@v2 - - run: cargo doc --no-deps --all-features - test: needs: [fmt, clippy] strategy: @@ -73,7 +62,7 @@ jobs: run: | pacman -Syu --noconfirm git curl base-devel - - uses: actions/checkout@v6 + - uses: actions/checkout@v7 # Explicit per-version install: RUSTUP_TOOLCHAIN overrides # rust-toolchain.toml, and the distro containers have no preinstalled Rust. @@ -96,13 +85,13 @@ jobs: ci-success: name: CI success if: ${{ always() }} - needs: [fmt, clippy, doc, test] + needs: [fmt, clippy, test] runs-on: ubuntu-latest steps: - name: Check job results run: | for r in "${{ needs.fmt.result }}" "${{ needs.clippy.result }}" \ - "${{ needs.doc.result }}" "${{ needs.test.result }}"; do + "${{ needs.test.result }}"; do case "$r" in success | skipped) ;; *) echo "required job result: $r"; exit 1 ;; diff --git a/.github/workflows/publish-crates.yml b/.github/workflows/publish-crates.yml index bf20c01..7049a86 100644 --- a/.github/workflows/publish-crates.yml +++ b/.github/workflows/publish-crates.yml @@ -21,7 +21,7 @@ jobs: if: startsWith(github.ref, 'refs/tags/') runs-on: ubuntu-latest steps: - - uses: actions/checkout@v6 + - uses: actions/checkout@v7 - uses: Swatinem/rust-cache@v2 - name: Publish diff --git a/CHANGELOG.md b/CHANGELOG.md index 4518650..b9f56b0 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,3 +5,9 @@ All notable changes to this project are documented in this file. The format follows Keep a Changelog, and the project uses Semantic Versioning. ## [Unreleased] + +### Added + +- Core public types: the `BatchCollector` trait, `BatchLoaderConfig`, and + `Error`. The runtime is not included yet - `BatchLoader` arrives in a later + release. diff --git a/Cargo.lock b/Cargo.lock index 3559188..df60c7a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -7,6 +7,7 @@ name = "carpool" version = "0.0.1" dependencies = [ "metrics", + "thiserror", "tokio", "tracing", ] @@ -83,6 +84,26 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "thiserror" +version = "2.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4288b5bcbc7920c07a1149a35cf9590a2aa808e0bc1eafaade0b80947865fbc4" +dependencies = [ + "thiserror-impl", +] + +[[package]] +name = "thiserror-impl" +version = "2.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ebc4ee7f67670e9b64d05fa4253e753e016c6c95ff35b89b7941d6b856dec1d5" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "tokio" version = "1.52.3" diff --git a/Cargo.toml b/Cargo.toml index 8a9d866..2f6ed5a 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -11,7 +11,8 @@ keywords = ["async", "batch", "dataloader", "dedup", "tokio"] categories = ["asynchronous", "concurrency"] [dependencies] -tokio = { version = "1", default-features = false, features = ["sync", "time"] } +tokio = { version = "1", default-features = false, features = ["sync", "time", "macros", "rt"] } +thiserror = "2" tracing = { version = "0.1", optional = true } metrics = { version = "0.24", optional = true } @@ -21,4 +22,4 @@ tracing = ["dep:tracing"] metrics = ["dep:metrics"] [dev-dependencies] -tokio = { version = "1", default-features = false, features = ["macros", "rt"] } +tokio = { version = "1", default-features = false, features = ["macros", "rt", "test-util"] } diff --git a/examples/dedup_window.rs b/examples/dedup_window.rs new file mode 100644 index 0000000..c3d59e6 --- /dev/null +++ b/examples/dedup_window.rs @@ -0,0 +1,74 @@ +use std::collections::HashMap; +use std::convert::Infallible; +use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; + +use carpool::{BatchCollector, BatchLoader, BatchLoaderConfig}; + +// Records how many times it is actually called, so dedup shows up as an +// observable fact, not just as equal results. Output depends on the key alone, +// so it does not matter which duplicate became the batch representative. +#[derive(Clone)] +struct CountingSquares { + calls: Arc, +} + +impl BatchCollector for CountingSquares { + type Input = u64; + type Output = u64; + type Key = u64; + type Error = Infallible; + + // Fold many inputs onto four keys so duplicates dominate the window. + fn key(&self, input: &u64) -> u64 { + *input % 4 + } + + async fn load(&self, batch: HashMap) -> Result, Infallible> { + self.calls.fetch_add(1, Ordering::SeqCst); + Ok(batch.into_keys().map(|k| (k, k * k)).collect()) + } +} + +#[tokio::main(flavor = "current_thread")] +async fn main() { + let calls = Arc::new(AtomicUsize::new(0)); + let loader = BatchLoader::spawn( + CountingSquares { + calls: calls.clone(), + }, + BatchLoaderConfig::default(), + ); + + // A hundred concurrent loads, but only four distinct keys. + let handles: Vec<_> = (0..100u64) + .map(|i| { + let loader = loader.clone(); + tokio::spawn(async move { loader.load(i).await }) + }) + .collect(); + + let mut results = Vec::new(); + for handle in handles { + results.push(handle.await.expect("task joins").expect("load succeeds")); + } + + let downstream_calls = calls.load(Ordering::SeqCst); + println!("loads: {}", results.len()); + println!("distinct keys: 4"); + println!("downstream calls: {downstream_calls}"); + + // The invariant, not the line order: one call served the whole window, and + // every load got the square of its key. + assert_eq!(results.len(), 100); + assert_eq!( + downstream_calls, 1, + "100 loads over 4 keys collapse to one downstream call" + ); + for (i, result) in results.iter().enumerate() { + let key = i as u64 % 4; + assert_eq!(*result, key * key, "each load gets the square of its key"); + } + + println!("dedup confirmed: 100 loads served by {downstream_calls} downstream call"); +} diff --git a/examples/quickstart.rs b/examples/quickstart.rs new file mode 100644 index 0000000..b7d4ea5 --- /dev/null +++ b/examples/quickstart.rs @@ -0,0 +1,44 @@ +// Smoke demo: implement BatchCollector, read the default config, +// see that equal inputs map to the same key. +// Run: cargo run --example quickstart +// Expected output: +// config: window=30ms max_batch=1024 timeout=30s concurrency=None max_waiting=None +// key(7) = 7, key(7) = 7 -> equal inputs share one batch entry + +use std::collections::HashMap; + +use carpool::{BatchCollector, BatchLoaderConfig}; + +#[derive(Clone)] +struct SquareLoader; + +impl BatchCollector for SquareLoader { + type Input = u64; + type Output = u64; + type Key = u64; + type Error = std::convert::Infallible; + + fn key(&self, input: &u64) -> u64 { + *input + } + + async fn load(&self, batch: HashMap) -> Result, Self::Error> { + Ok(batch.into_iter().map(|(k, n)| (k, n * n)).collect()) + } +} + +fn main() { + let cfg = BatchLoaderConfig::default(); + println!( + "config: window={:?} max_batch={} timeout={:?} concurrency={:?} max_waiting={:?}", + cfg.window, + cfg.max_batch_size.get(), + cfg.timeout, + cfg.concurrency_limit, + cfg.max_waiting, + ); + + let loader = SquareLoader; + let (a, b) = (loader.key(&7), loader.key(&7)); + println!("key(7) = {a}, key(7) = {b} -> equal inputs share one batch entry"); +} diff --git a/src/collector.rs b/src/collector.rs new file mode 100644 index 0000000..200714e --- /dev/null +++ b/src/collector.rs @@ -0,0 +1,63 @@ +use std::collections::HashMap; +use std::future::Future; +use std::hash::Hash; + +// Clone: loader and dispatcher hold separate clones - +// key() must be pure, shared state behind an Arc. +pub trait BatchCollector: Clone + Send + Sync + 'static { + type Input: Send + 'static; + type Output: Send + Clone + 'static; + type Key: Hash + Eq + Send + Clone + 'static; + type Error: std::error::Error + Send + Sync + Clone + 'static; + + fn key(&self, input: &Self::Input) -> Self::Key; + + // batch arrives deduplicated: inputs sharing a key are interchangeable, + // only one representative reaches load. + // strict contract: every requested key must get a value; + // a missing key is Error::MissingOutput for its waiters, + // an unknown key in the response is Error::ContractViolation + // for the whole batch (it takes precedence when both occur). + // absence semantics live in the implementor's Output type. + fn load( + &self, + batch: HashMap, + ) -> impl Future, Self::Error>> + Send; +} + +#[cfg(test)] +mod tests { + use super::*; + + #[derive(Clone)] + struct SquareLoader; + + impl BatchCollector for SquareLoader { + type Input = u64; + type Output = u64; + type Key = u64; + type Error = std::convert::Infallible; + + fn key(&self, input: &u64) -> u64 { + *input + } + + async fn load(&self, batch: HashMap) -> Result, Self::Error> { + Ok(batch.into_iter().map(|(k, x)| (k, x * x)).collect()) + } + } + + // Spawning the load future proves it is `Send` (the whole point of the + // RPITIT + Send form) and that a plain `async fn` satisfies the bound. + #[tokio::test] + async fn load_future_is_send_and_key_addressed() { + let loader = SquareLoader; + let batch = HashMap::from([(2, 2), (3, 3), (4, 4)]); + let out = tokio::spawn(async move { loader.load(batch).await }) + .await + .expect("task joins") + .expect("load succeeds"); + + assert_eq!(out, HashMap::from([(2, 4), (3, 9), (4, 16)])); + } +} diff --git a/src/config.rs b/src/config.rs new file mode 100644 index 0000000..9fb49d3 --- /dev/null +++ b/src/config.rs @@ -0,0 +1,41 @@ +use std::num::NonZeroUsize; +use std::time::Duration; + +#[derive(Debug, Clone)] +#[non_exhaustive] +pub struct BatchLoaderConfig { + pub window: Duration, + pub max_batch_size: NonZeroUsize, + pub timeout: Duration, + pub concurrency_limit: Option, + // only consulted when concurrency_limit is set + pub max_waiting: Option, +} + +impl Default for BatchLoaderConfig { + fn default() -> Self { + Self { + window: Duration::from_millis(30), + max_batch_size: NonZeroUsize::new(1024).expect("1024 is non-zero"), + timeout: Duration::from_secs(30), + concurrency_limit: None, + max_waiting: None, + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn default_config_matches_documented_values() { + let cfg = BatchLoaderConfig::default(); + + assert_eq!(cfg.window, Duration::from_millis(30)); + assert_eq!(cfg.max_batch_size.get(), 1024); + assert_eq!(cfg.timeout, Duration::from_secs(30)); + assert_eq!(cfg.concurrency_limit, None); + assert_eq!(cfg.max_waiting, None); + } +} diff --git a/src/dispatch.rs b/src/dispatch.rs new file mode 100644 index 0000000..a1acd5a --- /dev/null +++ b/src/dispatch.rs @@ -0,0 +1,259 @@ +use std::collections::HashMap; + +use tokio::sync::oneshot; + +use crate::collector::BatchCollector; +use crate::error::Error; + +// The one-shot reply a waiting caller holds. Factored into an alias because the +// nested generics otherwise trip clippy's type-complexity lint at each use site. +type Responder = + oneshot::Sender::Output, Error<::Error>>>; + +// One queued load: its key, the input to fold into the batch, and the one-shot +// channel that carries the result back to the waiting caller. Stays internal - +// the oneshot never appears in the public `load` signature. +pub(crate) struct Request { + pub(crate) key: C::Key, + pub(crate) input: C::Input, + pub(crate) respond: Responder, +} + +// Collapse a closed window into one downstream call and hand each result back to +// every caller that asked for that key. +pub(crate) async fn dispatch_window(collector: &C, batch: Vec>) { + // Dedup: one representative input per key for downstream, and every waiter + // grouped under its key so the result reaches all of them. Addressing is by + // key, never by position - inputs sharing a key are interchangeable. + let mut inputs: HashMap = HashMap::new(); + let mut waiters: HashMap>> = HashMap::new(); + for Request { + key, + input, + respond, + } in batch + { + waiters.entry(key.clone()).or_default().push(respond); + inputs.entry(key).or_insert(input); + } + + match collector.load(inputs).await { + Err(e) => { + // Downstream failed: the whole batch shares the same error. + for senders in waiters.into_values() { + for respond in senders { + deliver::(respond, Err(Error::Collector(e.clone()))); + } + } + } + Ok(mut response) => { + // An unknown key in the response is an implementor bug that taints the + // whole batch and takes precedence over any missing key. + let unknown = response.keys().filter(|k| !waiters.contains_key(k)).count(); + if unknown > 0 { + for senders in waiters.into_values() { + for respond in senders { + deliver::( + respond, + Err(Error::ContractViolation { + unknown_keys: unknown, + }), + ); + } + } + return; + } + + for (key, senders) in waiters { + match response.remove(&key) { + // Fan-out: each waiter gets its own clone of the shared value. + Some(value) => { + for respond in senders { + deliver::(respond, Ok(value.clone())); + } + } + // Requested key absent: only its waiters get the addressed error. + None => { + for respond in senders { + deliver::(respond, Err(Error::MissingOutput)); + } + } + } + } + } + } +} + +fn deliver(respond: Responder, result: Result>) { + let delivered = respond.send(result).is_ok(); + warn_if_dropped(delivered); +} + +// A dropped waiter is benign - the caller cancelled - but it is never silently +// ignored: when `tracing` is compiled in we surface it at warn level, otherwise +// there is no log facade to write to. +#[cfg(feature = "tracing")] +fn warn_if_dropped(delivered: bool) { + if !delivered { + tracing::warn!("carpool: dropped a load result because its caller went away"); + } +} + +#[cfg(not(feature = "tracing"))] +fn warn_if_dropped(_delivered: bool) {} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + use std::sync::atomic::{AtomicUsize, Ordering}; + + use super::*; + + #[derive(Debug, Clone)] + struct TestError; + + impl std::fmt::Display for TestError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str("test downstream error") + } + } + + impl std::error::Error for TestError {} + + // What the collector does with a batch, so each contract branch can be driven. + #[derive(Clone)] + enum Behavior { + Square, + OmitKey(u64), + InjectUnknown(u64), + Fail, + } + + #[derive(Clone)] + struct TestCollector { + calls: Arc, + batch_len: Arc, + behavior: Behavior, + } + + impl BatchCollector for TestCollector { + type Input = u64; + type Output = u64; + type Key = u64; + type Error = TestError; + + fn key(&self, input: &u64) -> u64 { + *input + } + + async fn load(&self, batch: HashMap) -> Result, TestError> { + self.calls.fetch_add(1, Ordering::SeqCst); + self.batch_len.store(batch.len(), Ordering::SeqCst); + let squared = || batch.iter().map(|(k, v)| (*k, v * v)); + match self.behavior { + Behavior::Square => Ok(squared().collect()), + Behavior::OmitKey(drop) => Ok(squared().filter(|(k, _)| *k != drop).collect()), + Behavior::InjectUnknown(extra) => { + let mut out: HashMap = squared().collect(); + out.insert(extra, 0); + Ok(out) + } + Behavior::Fail => Err(TestError), + } + } + } + + type Reply = oneshot::Receiver>>; + + fn req(key: u64, input: u64) -> (Request, Reply) { + let (tx, rx) = oneshot::channel(); + ( + Request { + key, + input, + respond: tx, + }, + rx, + ) + } + + fn collector(behavior: Behavior) -> (TestCollector, Arc, Arc) { + let calls = Arc::new(AtomicUsize::new(0)); + let batch_len = Arc::new(AtomicUsize::new(0)); + let collector = TestCollector { + calls: calls.clone(), + batch_len: batch_len.clone(), + behavior, + }; + (collector, calls, batch_len) + } + + // Dedup: duplicate keys collapse to one downstream input and one call; both + // waiters of a key receive the shared value. + #[tokio::test] + async fn duplicate_keys_collapse_to_one_input_per_key() { + let (collector, calls, batch_len) = collector(Behavior::Square); + let (a, rx_a) = req(1, 1); + let (b, rx_b) = req(1, 1); + let (c, rx_c) = req(2, 2); + + dispatch_window(&collector, vec![a, b, c]).await; + + assert_eq!(calls.load(Ordering::SeqCst), 1, "one downstream call"); + assert_eq!( + batch_len.load(Ordering::SeqCst), + 2, + "three requests, two unique keys" + ); + assert_eq!(rx_a.await.unwrap().unwrap(), 1); + assert_eq!(rx_b.await.unwrap().unwrap(), 1); + assert_eq!(rx_c.await.unwrap().unwrap(), 4); + } + + // Absent requested key: only its waiter gets the addressed MissingOutput, + // the others still get their values. + #[tokio::test] + async fn missing_key_is_addressed_only_to_its_waiter() { + let (collector, _calls, _len) = collector(Behavior::OmitKey(2)); + let (a, rx_a) = req(1, 1); + let (b, rx_b) = req(2, 2); + + dispatch_window(&collector, vec![a, b]).await; + + assert_eq!(rx_a.await.unwrap().unwrap(), 1); + assert!(matches!(rx_b.await.unwrap(), Err(Error::MissingOutput))); + } + + // Unknown key in the response fails the whole batch, even the waiters whose + // keys were answered correctly. + #[tokio::test] + async fn unknown_key_fails_the_whole_batch() { + let (collector, _calls, _len) = collector(Behavior::InjectUnknown(99)); + let (a, rx_a) = req(1, 1); + let (b, rx_b) = req(2, 2); + + dispatch_window(&collector, vec![a, b]).await; + + assert!(matches!( + rx_a.await.unwrap(), + Err(Error::ContractViolation { unknown_keys: 1 }) + )); + assert!(matches!( + rx_b.await.unwrap(), + Err(Error::ContractViolation { unknown_keys: 1 }) + )); + } + + // Downstream error reaches every waiter in the batch. + #[tokio::test] + async fn collector_error_reaches_every_waiter() { + let (collector, _calls, _len) = collector(Behavior::Fail); + let (a, rx_a) = req(1, 1); + let (b, rx_b) = req(2, 2); + + dispatch_window(&collector, vec![a, b]).await; + + assert!(matches!(rx_a.await.unwrap(), Err(Error::Collector(_)))); + assert!(matches!(rx_b.await.unwrap(), Err(Error::Collector(_)))); + } +} diff --git a/src/error.rs b/src/error.rs new file mode 100644 index 0000000..e6aadb4 --- /dev/null +++ b/src/error.rs @@ -0,0 +1,67 @@ +use thiserror::Error; + +#[derive(Debug, Clone, Error)] +#[non_exhaustive] +pub enum Error { + // #[source], not #[from]: #[from] would generate From and clash for a generic E + #[error("collector failed: {0}")] + Collector(#[source] E), + // implementor bug: response key not in the batch - fails the whole batch + #[error( + "collector broke the key-addressed contract: {unknown_keys} unknown key(s) in the response" + )] + ContractViolation { unknown_keys: usize }, + // implementor bug: no output for a requested key - not a domain "not found" + #[error("collector returned no output for a requested key")] + MissingOutput, + #[error("batch timed out")] + Timeout, + #[error("timed out waiting for a concurrency slot")] + WaitingTimeout, + // background pipeline shut down (for example a downstream panic tore the dispatcher down) + #[error("the batch loader has shut down")] + Closed, +} + +#[cfg(test)] +mod tests { + use super::*; + + #[derive(Debug, Clone)] + struct DownstreamError(&'static str); + + impl std::fmt::Display for DownstreamError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "downstream: {}", self.0) + } + } + + impl std::error::Error for DownstreamError {} + + #[test] + fn collector_variant_exposes_downstream_as_source() { + let err: Error = Error::Collector(DownstreamError("boom")); + let source = std::error::Error::source(&err).expect("Collector has a source"); + + assert_eq!(source.to_string(), "downstream: boom"); + } + + // Contract-bug variants wrap nothing downstream: + // no source, and their Display stands on its own without touching E. + #[test] + fn contract_bug_variants_have_no_source_and_render() { + let missing: Error = Error::MissingOutput; + let violation: Error = Error::ContractViolation { unknown_keys: 2 }; + + assert!(std::error::Error::source(&missing).is_none()); + assert!(std::error::Error::source(&violation).is_none()); + assert_eq!( + missing.to_string(), + "collector returned no output for a requested key" + ); + assert_eq!( + violation.to_string(), + "collector broke the key-addressed contract: 2 unknown key(s) in the response" + ); + } +} diff --git a/src/lib.rs b/src/lib.rs index 1217e7b..fc5a0c6 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,12 +1,13 @@ -//! Deduplicate and batch concurrent async requests. -//! -//! Many concurrent `load(input)` calls are merged within a collection window -//! into a single downstream batch, and duplicate keys share one result. There -//! is no cache; the API is trait-based and built on `tokio`. -//! -//! The public API is not available yet; this version ships the crate skeleton -//! only. - #![forbid(unsafe_code)] -#![warn(missing_docs)] -#![deny(rustdoc::broken_intra_doc_links)] + +mod collector; +mod config; +mod dispatch; +mod error; +mod loader; +mod window; + +pub use collector::BatchCollector; +pub use config::BatchLoaderConfig; +pub use error::Error; +pub use loader::BatchLoader; diff --git a/src/loader.rs b/src/loader.rs new file mode 100644 index 0000000..6b322ff --- /dev/null +++ b/src/loader.rs @@ -0,0 +1,214 @@ +use tokio::sync::{mpsc, oneshot}; + +use crate::collector::BatchCollector; +use crate::config::BatchLoaderConfig; +use crate::dispatch::{Request, dispatch_window}; +use crate::error::Error; +use crate::window; + +// Public runtime: every `load` drops its request into the collection window and +// awaits a one-shot reply. A background dispatcher turns each closed window into +// one deduplicated downstream call and fans the results back out by key. +// Cloning is cheap (a collector clone and a sender), so the loader is shared across tasks. +#[derive(Clone)] +pub struct BatchLoader { + collector: C, + inbound: mpsc::Sender>, +} + +impl BatchLoader { + // spawn, not new: two background tasks start eagerly; a lazy ctor can come later. + pub fn spawn(collector: C, config: BatchLoaderConfig) -> Self { + // inbound buffers a burst up to one window's worth; outbound holds at most + // one closed window awaiting dispatch - deeper buffering would only hide + // backpressure behind memory. + let (inbound, requests) = mpsc::channel::>(config.max_batch_size.get()); + let (windows_tx, windows_rx) = mpsc::channel::>>(1); + + // Both tasks are detached: they end on channel close once every loader + // clone is dropped, so there is nothing to join. + tokio::spawn(window::collect( + requests, + windows_tx, + config.window, + config.max_batch_size, + )); + tokio::spawn(run_dispatcher(collector.clone(), windows_rx)); + + Self { collector, inbound } + } + + // The key is derived through the collector so dedup and dispatch address the + // same key the implementor sees. + pub async fn load(&self, input: C::Input) -> Result> { + let key = self.collector.key(&input); + let (respond, reply) = oneshot::channel(); + if self + .inbound + .send(Request { + key, + input, + respond, + }) + .await + .is_err() + { + return Err(Error::Closed); + } + match reply.await { + Ok(result) => result, + // The dispatcher dropped our responder without answering + // (a downstream panic can tear it down), so the loader cannot serve. + Err(_) => Err(Error::Closed), + } + } +} + +async fn run_dispatcher( + collector: C, + mut windows: mpsc::Receiver>>, +) { + // One closed window at a time, one downstream call each. Ends when the window + // task drops its sender (all loaders gone), winding the runtime down. + while let Some(batch) = windows.recv().await { + dispatch_window(&collector, batch).await; + } +} + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + use std::num::NonZeroUsize; + use std::sync::Arc; + use std::sync::atomic::{AtomicUsize, Ordering}; + use std::time::Duration; + + use super::*; + + #[derive(Clone)] + struct Squares { + calls: Arc, + } + + impl BatchCollector for Squares { + type Input = u64; + type Output = u64; + type Key = u64; + type Error = std::convert::Infallible; + + fn key(&self, input: &u64) -> u64 { + *input + } + + async fn load(&self, batch: HashMap) -> Result, Self::Error> { + self.calls.fetch_add(1, Ordering::SeqCst); + Ok(batch.into_iter().map(|(k, v)| (k, v * v)).collect()) + } + } + + // A loader whose window never closes on the timer (virtual time stays frozen), + // so windows close deterministically at `max` items - no scheduler race. + fn loader(max: usize, calls: Arc) -> BatchLoader { + let config = BatchLoaderConfig { + window: Duration::from_secs(3600), + max_batch_size: NonZeroUsize::new(max).expect("max is non-zero"), + ..BatchLoaderConfig::default() + }; + BatchLoader::spawn(Squares { calls }, config) + } + + fn spawn_load( + loader: &BatchLoader, + input: u64, + ) -> tokio::task::JoinHandle>> { + let loader = loader.clone(); + tokio::spawn(async move { loader.load(input).await }) + } + + async fn collect_results( + handles: Vec>>>, + ) -> Vec { + let mut results = Vec::new(); + for handle in handles { + results.push(handle.await.expect("task joins").expect("load succeeds")); + } + results + } + + // Dedup proven by the call count: five concurrent loads over two keys collapse + // to a single downstream call, and every waiter still gets its key's value. + #[tokio::test(start_paused = true)] + async fn duplicate_keys_collapse_to_one_downstream_call() { + let calls = Arc::new(AtomicUsize::new(0)); + let loader = loader(5, calls.clone()); + + let handles = [1u64, 1, 1, 2, 2] + .into_iter() + .map(|k| spawn_load(&loader, k)) + .collect(); + let results = collect_results(handles).await; + + assert_eq!( + calls.load(Ordering::SeqCst), + 1, + "dedup -> single downstream call" + ); + assert_eq!(results, vec![1, 1, 1, 4, 4]); + } + + // Distinct keys in one window each get their own result from the one call. + #[tokio::test(start_paused = true)] + async fn distinct_keys_each_get_their_own_result() { + let calls = Arc::new(AtomicUsize::new(0)); + let loader = loader(3, calls.clone()); + + let handles = [2u64, 3, 4] + .into_iter() + .map(|k| spawn_load(&loader, k)) + .collect(); + let results = collect_results(handles).await; + + assert_eq!(calls.load(Ordering::SeqCst), 1); + assert_eq!(results, vec![4, 9, 16]); + } + + // A load arriving after the window closed opens a fresh one, so it is a new + // downstream call - the count goes from one to two. + #[tokio::test(start_paused = true)] + async fn a_load_after_the_window_closes_starts_a_new_call() { + let calls = Arc::new(AtomicUsize::new(0)); + let loader = loader(1, calls.clone()); + + let first = spawn_load(&loader, 5) + .await + .expect("task joins") + .expect("load succeeds"); + assert_eq!(first, 25); + assert_eq!(calls.load(Ordering::SeqCst), 1); + + let second = spawn_load(&loader, 6) + .await + .expect("task joins") + .expect("load succeeds"); + assert_eq!(second, 36); + assert_eq!( + calls.load(Ordering::SeqCst), + 2, + "the later load opens a fresh window" + ); + } + + // Regression guard: load must stay Send so a cloned loader can spawn it. + #[tokio::test(start_paused = true)] + async fn load_future_is_send_after_clone() { + let calls = Arc::new(AtomicUsize::new(0)); + let loader = loader(1, calls); + + let value = spawn_load(&loader, 7) + .await + .expect("task joins") + .expect("load succeeds"); + + assert_eq!(value, 49); + } +} diff --git a/src/window.rs b/src/window.rs new file mode 100644 index 0000000..85a3c10 --- /dev/null +++ b/src/window.rs @@ -0,0 +1,135 @@ +use std::num::NonZeroUsize; +use std::pin::pin; +use std::time::Duration; + +use tokio::sync::mpsc; +use tokio::time::{self, Instant}; + +// Collection window: the first item opens a window and arms a `window`-long +// timer; the window closes on whichever comes first - `max_batch_size` items or +// the timer firing - and the next item opens a fresh one. +// +// The boundary rule lives in one place: the select! is `biased` with the timer +// first, so when an item and the timer are ready together the timer wins - the +// window closes and the item stays in the channel for the next one, never lost +// and never duplicated. +// +// `collect` batches opaque `T` by time and size only and spawns no task: both +// channel-close paths are normal teardown, not errors (inbound closed - source +// gone, flush and stop; outbound closed - consumer gone, stop). +pub(crate) async fn collect( + mut inbound: mpsc::Receiver, + outbound: mpsc::Sender>, + window: Duration, + max_batch_size: NonZeroUsize, +) { + let max = max_batch_size.get(); + + // No active window: block until the first item opens one. + while let Some(first) = inbound.recv().await { + let mut buffer = Vec::with_capacity(max); + buffer.push(first); + + let mut timer = pin!(time::sleep_until(Instant::now() + window)); + let mut inbound_open = true; + + // Window open: fill toward max_batch_size or until the timer fires. + // The buffer is owned across select!, so a cancelled branch never drops + // an accepted item - the cancel-safety the later tasks rely on. + while buffer.len() < max { + tokio::select! { + biased; + () = timer.as_mut() => break, + maybe = inbound.recv() => match maybe { + Some(input) => buffer.push(input), + None => { + inbound_open = false; + break; + } + }, + } + } + + // Single hand-off point for the closed window. A closed outbound or a + // closed inbound both end collection - nothing left to do. + if outbound.send(buffer).await.is_err() || !inbound_open { + return; + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + const WINDOW: Duration = Duration::from_millis(30); + + fn spawn_window(max: usize) -> (mpsc::Sender, mpsc::Receiver>) { + let (in_tx, in_rx) = mpsc::channel::(64); + let (out_tx, out_rx) = mpsc::channel::>(64); + let max = NonZeroUsize::new(max).expect("max is non-zero"); + tokio::spawn(collect(in_rx, out_tx, WINDOW, max)); + (in_tx, out_rx) + } + + // Close-by-size at the boundary: max = 3, so item 4 must not ride in the + // first batch - it opens a new window. + #[tokio::test(start_paused = true)] + async fn size_closes_at_max_and_spills_to_next_window() { + let (tx, mut batches) = spawn_window::(3); + for i in 1..=4 { + tx.send(i).await.expect("send"); + } + + assert_eq!(batches.recv().await.expect("first batch"), vec![1, 2, 3]); + + time::advance(WINDOW).await; + assert_eq!(batches.recv().await.expect("second batch"), vec![4]); + } + + // Close-by-timer with a single underfilled item: the window must close on + // `window` and not hang waiting for max_batch_size. + #[tokio::test(start_paused = true)] + async fn single_input_closes_on_timer() { + let (tx, mut batches) = spawn_window::(16); + tx.send(7).await.expect("send"); + + time::advance(WINDOW).await; + assert_eq!(batches.recv().await.expect("batch"), vec![7]); + } + + // Close-by-timer with several underfilled items: all of them ride the one + // batch the timer closes. + #[tokio::test(start_paused = true)] + async fn partial_window_closes_on_timer_with_all_items() { + let (tx, mut batches) = spawn_window::(16); + tx.send(1).await.expect("send 1"); + tx.send(2).await.expect("send 2"); + + time::advance(WINDOW).await; + assert_eq!(batches.recv().await.expect("batch"), vec![1, 2]); + } + + // Boundary: item 2 is waiting in the channel when the timer fires. The + // biased select gives the timer priority, so the window closes with only + // [1] and item 2 falls into the next window - present exactly once across + // both batches, never lost and never duplicated. + #[tokio::test(start_paused = true)] + async fn timer_wins_boundary_against_waiting_input() { + let (tx, mut batches) = spawn_window::(16); + tx.send(1).await.expect("send 1"); + // Let the window open on item 1 and park in the select before the + // deadline. + tokio::task::yield_now().await; + // Reach the deadline first, then deliver item 2 at that instant: the + // timer is already ready, so the biased select closes the window before + // it ever looks at recv. Item 2 lands at exactly now == deadline. + time::advance(WINDOW).await; + tx.try_send(2).expect("channel has capacity"); + + assert_eq!(batches.recv().await.expect("first batch"), vec![1]); + + time::advance(WINDOW).await; + assert_eq!(batches.recv().await.expect("second batch"), vec![2]); + } +}