Autobatch

Autobatch #

Autobatch is the small runtime in libs/rust/bonsai-utils/src/autobatch.rs that the batch extraction commands run on. This page explains why it exists and how it works inside. It is about 460 lines of code plus tests; reading this page first should make the source easy to follow.

Why #

A batch command receives up to 300 extractions. Each extraction goes through the same steps (load, validate, check with the provider, plan learning, write, notify), but not every extraction takes every step: one is missing, one is deleted, one skips validation, one fails the provider check.

There are two obvious ways to write that, and both are bad:

  • Per-item code in a loop. It reads well, because the logic for one extraction is a straight line with early returns. It costs one round trip per step per extraction, which is where “too many sequential calls” came from.
  • Hand-batched stages. Load everything, split the survivors from the failures, validate the survivors, split again, write, then merge every branch back into one result per id. It is fast, but the business logic is spread across stages, every branch needs its own split-and-merge plumbing, and it is easy to drop an item or report it twice.

Autobatch keeps the first shape and gets the cost of the second. The command is written as code for one extraction. Wherever that code needs I/O it calls a batch. The runtime runs all items side by side, collects their calls, and performs each batch once for everyone who asked.

The two properties the design is built around:

  1. One fetch per step, however the items branch.
  2. A region cannot hang. A runtime that parks items until “the right moment” has an obvious failure mode: an item parked on a batch nobody will ever flush, and a request that never returns. Autobatch rules that out with types, not with discipline.

The model # 

autobatch::scope(async |mut batcher| {
    let load = batcher.define_batch("load", fetch_many);      // Vec<I> -> one result per input
    let write = batcher.define_batch("write", write_many);
    batcher.run_per_item(ids, async |ctx, id| {
        let row = load.call(ctx, id).await?;                  // queues `id`, parks this item
        write.call(ctx, plan(row)?).await
    }).await
}).await
  • A region is one scope call.
  • A batch is a queue of inputs plus a fetch function that turns Vec<Input> into one result per input.
  • An item is one invocation of the per_item closure: an ordinary Rust future.
  • call puts an input in the batch’s queue and suspends the item until its result exists.

Everything in a region runs on one task. The items are not spawned; they are futures that the region polls itself. That is the whole trick: because the region does the polling, it can tell when every item is stuck, and that is the moment to flush.

Inside a batch #

Each batch owns this state behind an Rc<RefCell<..>>:

struct State<I, O, E> {
    queued: Vec<(usize, usize, I)>,     // (item index, slot, input), in call order
    next_slot: usize,
    done: HashMap<usize, Result<O, E>>, // slot -> result, filled by a flush
    parked: Vec<Waker>,                 // items waiting on this batch
    flushes: usize,
}

call is the only thing an item can do with a batch:

pub async fn call(&self, ctx: &ItemContext<'region>, input: I) -> Result<O, E> {
    // 1. Take a slot and queue the input. The borrow ends before the first await.
    let slot = {
        let mut state = self.state.borrow_mut();
        let slot = state.next_slot;
        state.next_slot += 1;
        state.queued.push((ctx.index, slot, input));
        slot
    };
    // 2. Each time this item is polled: is my slot answered yet?
    poll_fn(|cx| {
        let mut state = self.state.borrow_mut();
        match state.done.remove(&slot) {
            Some(result) => Poll::Ready(result),      // a flush filled it in
            None => {
                state.parked.push(cx.waker().clone()); // park until the next flush
                Poll::Pending
            }
        }
    })
    .await
}

The slot routes a result back to the call that asked for it. It is per call, not per item, so one item may call the same batch several times.

The batcher keeps its own handle to every batch, type-erased, in creation order:

pub struct Batcher<'region, 'a> {
    batches: Vec<Box<dyn Flush + 'a>>,
    brand: Brand<'region>,
}

trait Flush {
    fn has_queued(&self) -> bool;
    fn flush(&self, flush: usize) -> LocalFuture<'_, ()>;
}

define_batch creates the State, pushes a Flusher { name, state, fetch } into batches and returns a Batch holding the same Rc. Batches with different input and output types live in the one list because the driver only needs those two methods.

The driver loop #

run_per_item turns the items into futures and hands them to the driver:

pub async fn run_per_item<X, T, E>(
    self,                                             // consumes the batcher
    items: impl IntoIterator<Item = X>,
    per_item: impl AsyncFn(&ItemContext<'region>, X) -> Result<T, E>,
) -> Result<Vec<T>, E> {
    let items = items.into_iter().collect::<Vec<_>>();
    let contexts = (0..items.len())
        .map(|index| ItemContext { index, brand: PhantomData })
        .collect::<Vec<_>>();
    let futures = contexts.iter().zip(items).map(|(ctx, item)| per_item(ctx, item));
    let items = try_join_in_order(futures);           // one future for "all items"
    drive(&self.batches, items).await
}

Calling per_item(ctx, item) runs nothing yet. It only creates a future. Nothing happens until that future is polled.

Step 1: polling all items is one future #

try_join_in_order is a hand-written join. Each time it is polled, it polls every unfinished item once, in index order:

async fn try_join_in_order<T, E, F>(rows: impl IntoIterator<Item = F>) -> Result<Vec<T>, E>
where
    F: Future<Output = Result<T, E>>,
{
    let mut rows: Vec<_> = rows.into_iter().map(|row| Some(Box::pin(row))).collect();
    let mut done: Vec<Option<T>> = rows.iter().map(|_| None).collect();
    poll_fn(|cx| {
        for (index, slot) in rows.iter_mut().enumerate() {
            let Some(row) = slot else { continue };              // already finished
            match row.as_mut().poll(cx) {
                Poll::Pending => {}                              // parked on a call, or on I/O
                Poll::Ready(Err(error)) => return Poll::Ready(Err(error)), // stop everything
                Poll::Ready(Ok(value)) => {
                    done[index] = Some(value);
                    *slot = None;
                }
            }
        }
        match rows.iter().all(Option::is_none) {
            true => Poll::Ready(Ok(done.iter_mut().filter_map(Option::take).collect())),
            false => Poll::Pending,
        }
    })
    .await
}

One poll of this future means: every item runs as far as it can, until it finishes or hits a call whose slot is not answered. After one poll, every unfinished item is parked somewhere.

Step 2: the driver decides what “stuck” means #

drive polls that future itself, with a waker of its own, so it can see whether anything is still able to move:

#[derive(Default)]
struct Signal {
    woken: AtomicBool,   // "something woke up since I last looked"
    outer: AtomicWaker,  // the real task running the region
}

impl Wake for Signal {
    fn wake_by_ref(self: &Arc<Self>) {
        self.woken.store(true, Ordering::SeqCst);
        self.outer.wake();
    }
}

async fn drive<T>(batches: &[Box<dyn Flush + '_>], rows: impl Future<Output = T>) -> T {
    let mut rows = pin!(rows);
    let signal = Arc::new(Signal::default());
    let waker = Waker::from(signal.clone());
    let mut flushes = 0;
    loop {
        // A. Poll the items until they finish or settle.
        let settled = poll_fn(|cx| {
            signal.outer.register(cx.waker());
            loop {
                signal.woken.store(false, Ordering::SeqCst);
                if let Poll::Ready(output) = rows.as_mut().poll(&mut Context::from_waker(&waker)) {
                    return Poll::Ready(Some(output));        // all items done
                }
                if signal.woken.load(Ordering::SeqCst) {
                    continue;                                // an item woke itself: poll again
                }
                return match batches.iter().any(|batch| batch.has_queued()) {
                    true => Poll::Ready(None),               // settled: someone needs a flush
                    false => Poll::Pending,                  // waiting on I/O outside any batch
                };
            }
        })
        .await;
        if let Some(output) = settled {
            return output;
        }
        // B. Flush the earliest batch that has queued inputs, then poll again.
        if let Some(batch) = batches.iter().find(|batch| batch.has_queued()) {
            flushes += 1;
            batch.flush(flushes).await;
        }
    }
}

Reading part A as a decision table, after one poll of the items:

Items woken A batch has queued inputs The driver
finished – – returns their output
pending set – polls again; an item is runnable (it yielded, or its I/O just completed)
pending clear yes has settled: go to B and flush
pending clear no returns Pending; the task sleeps until outside I/O calls the waker

The last row is why Signal forwards to outer: when that outside I/O completes it calls the Signal waker, which sets woken and wakes the real task, and the runtime polls drive again.

Step 3: a flush # 

fn flush(&self, flush: usize) -> LocalFuture<'_, ()> {
    Box::pin(async move {
        // Take the whole queue, in item order.
        let (slots, inputs): (Vec<_>, Vec<_>) = {
            let mut state = self.state.borrow_mut();
            state.flushes += 1;
            let mut queued = mem::take(&mut state.queued);
            queued.sort_by_key(|(row, _, _)| *row);          // stable: an item's own calls keep order
            queued.into_iter().map(|(_, slot, input)| (slot, input)).unzip()
        };                                                   // borrow released before the fetch
        let expected = inputs.len();
        let span = info_span!("autobatch.flush", batch = %self.name, rows = expected, flush);
        let results = match (self.fetch)(inputs).instrument(span).await {
            Ok(results) if results.len() == expected => results,
            Ok(results) => {                                 // wrong count: fail every caller
                let error = E::from(WrongResultCount { expected, actual: results.len() });
                (0..expected).map(|_| Err(error.clone())).collect()
            }
            Err(error) => (0..expected).map(|_| Err(error.clone())).collect(), // fetch failed
        };
        // Hand each result to its slot and wake the callers.
        let mut state = self.state.borrow_mut();
        state.done.extend(slots.into_iter().zip(results));
        state.parked.drain(..).for_each(Waker::wake);
    })
}

The two Err arms are why the error type must be Clone: one error is copied into every slot of that flush.

Only one flush runs at a time, and the items are not polled while it runs. A batch’s RefCell is therefore only borrowed inside one call poll or on either side of the fetch await, never across an await, so a double borrow cannot happen.

One round, end to end #

Two items, batches load then write:

drive: poll items
  item 0: load.call(a)  -> queued [(0, slot 0, a)], parked
  item 1: load.call(b)  -> queued [(0, slot 0, a), (1, slot 1, b)], parked
  woken clear, load has queued inputs -> settled
drive: flush load   -> fetch([a, b]) -> done {0: A, 1: B}, wake both
drive: poll items
  item 0: slot 0 ready -> plan -> write.call(x) -> queued, parked
  item 1: slot 1 ready -> plan -> write.call(y) -> queued, parked
  settled
drive: flush write  -> fetch([x, y]) -> done, wake both
drive: poll items
  item 0: Ready(Ok)   item 1: Ready(Ok)   -> all done, return

Why the earliest batch flushes first #

When the items settle, several batches may have queued inputs, because items take different paths. Part B of drive uses find, so it always picks the batch that was created first. This test from autobatch.rs shows the effect:

let validate = doubling(&mut batcher, "validate", &log);   // created first
let write = doubling(&mut batcher, "write", &log);
batcher
    .run_per_item(0..6, async |ctx, row| {
        let value = match row % 3 {
            0 => validate.call(ctx, row).await? / 2,        // items 0 and 3 validate
            _ => row,                                       // the others skip it
        };
        write.call(ctx, value).await
    })
    .await

// What the fetches were called with:
// [("validate", [0, 3]), ("write", [0, 1, 2, 3, 4, 5])]
Round Items do Queues after polling Flushed
1 0 and 3 call validate; 1, 2, 4, 5 call write validate: 0, 3 — write: 1, 2, 4, 5 validate with [0, 3]
2 0 and 3 resume and call write write: 1, 2, 4, 5, 0, 3 write with [0, 1, 2, 3, 4, 5]

Had write flushed in round 1, items 0 and 3 would have needed a second write. Flushing the earliest batch lets the items that are behind catch up before a later batch runs. As long as items call batches in creation order, each batch flushes once per region, whatever the branching.

The sort_by_key in the flush is why write sees [0, 1, 2, 3, 4, 5] and not the queue order [1, 2, 4, 5, 0, 3]: a fetch sees its inputs in the order the items were given to the region, whichever round each was queued in.

The opposite case costs an extra flush. An item that goes back to an earlier batch after a later one makes that earlier batch flush again:

let a = first.call(ctx, x).await?;
let b = second.call(ctx, a).await?;
let c = first.call(ctx, b).await?;   // `first` flushes a second time for these calls

I/O outside a batch #

An item may await something that is not a batch:

batcher.run_per_item(0..3, async |ctx, row| {
    elsewhere().await;                 // not a batch: a cache lookup, a timer, a spawned task
    batch.call(ctx, row).await
})

While every item is inside elsewhere(), no batch has queued inputs, so drive returns Pending (last row of the decision table). When the I/O completes it calls the Signal waker, the task is polled, the items move on to batch.call, and the batch flushes once with all three.

It only costs something when items disagree:

batcher.run_per_item(0..3, async |ctx, row| {
    if row == 0 {
        slow_lookup().await;           // item 0 is still here ...
    }
    batch.call(ctx, row).await         // ... while items 1 and 2 are parked here
})
// batch flushes with [1, 2]; later, when the lookup completes, again with [0]

The driver sees pending items, woken clear and a batch with queued inputs, so it flushes without item 0. That is a deliberate trade: waiting for item 0 would hang the region if its I/O never completed, so the cost is an extra flush. It is also the reason to put I/O inside batches.

How hanging is ruled out #

A region hangs if an item is parked on a batch that will never flush. There are three ways that could happen, and each is a compile error. The examples below are the module’s own compile_fail doctests, so these guarantees are tested.

1. A call from something the driver is not polling. call needs an &ItemContext, and contexts exist only inside run_per_item:

autobatch::scope(async |mut batcher| {
    let load = batcher.define_batch("load", fetch);
    let _ = load.call(7).await;          // error[E0061]: missing the item context
}).await;

A context cannot be smuggled out of its item either, because it is lent by reference for the duration of that item’s closure:

let kept: RefCell<Vec<&ItemContext<'_>>> = RefCell::new(vec![]);
batcher.run_per_item([1], async |ctx, id: u32| {
    kept.borrow_mut().push(ctx);         // error[E0521]: borrowed data escapes the closure
    Ok::<_, ()>(id)
}).await

2. A fetch calling a batch of its own region. While a fetch runs, the driver is awaiting it and flushes nothing else, so a fetch parked on a sibling batch would wait forever. A fetch is not given a context, and it cannot make one, because run_per_item takes the batcher by value:

let load = batcher.define_batch("load", |ids: Vec<u32>| async {
    batcher.run_per_item([1], async |_, id: u32| Ok::<_, E>(id)).await?;  // uses `batcher`
    Ok::<_, E>(ids.into_iter().map(Ok).collect())
});
batcher.run_per_item([1], async |ctx, id: u32| load.call(ctx, id).await).await
// error[E0505]: cannot move out of `batcher` because it is borrowed

The same move is what makes the driver’s list complete: once run_per_item owns the batcher, no batch can be defined behind its back.

3. Mixing regions. A nested region’s items are polled by the inner driver, so a call from them to an outer batch would park where the outer driver cannot see it:

autobatch::scope(async |mut batcher| {
    let load = batcher.define_batch("load", fetch);
    batcher.run_per_item([1], async |_, id: u32| {
        autobatch::scope(async |inner| {
            inner.run_per_item([id], async |ctx, id: u32| load.call(ctx, id).await).await
            //                                            ^^^^ outer batch, inner context: error
        }).await
    }).await
}).await;

What makes that an error is a brand. scope has this signature:

pub async fn scope<'a, T>(region: impl for<'region> AsyncFnOnce(Batcher<'region, 'a>) -> T) -> T

type Brand<'region> = PhantomData<fn(&'region ()) -> &'region ()>;

pub struct Batch<'region, I, O, E> { state: Rc<RefCell<State<I, O, E>>>, brand: Brand<'region> }
pub struct ItemContext<'region>    { index: usize,                       brand: Brand<'region> }

impl<'region, I, O, E> Batch<'region, I, O, E> {
    pub async fn call(&self, ctx: &ItemContext<'region>, input: I) -> Result<O, E>
    //                                         ^^^^^^^ must be this batch's own region
}
  • for<'region> means the closure must compile for any lifetime, so inside it 'region is a fresh lifetime that cannot be named or equated with anything else.
  • fn(&'region ()) -> &'region () uses the lifetime in both argument and return position, which makes the type invariant: the compiler will not shrink or stretch 'region to make two regions agree.
  • call demands a context with the batch’s own 'region. The inner region’s contexts carry the inner lifetime, so the types do not match.

For the same reason a batch cannot be returned out of scope: its type would have to mention 'region, which does not exist outside the closure.

A fetch may open a region of its own, with its own batcher and items. That is safe because the inner region drives itself to completion inside the fetch.

Single-threaded by design #

Nothing in a region is Send or Sync: batch state is Rc<RefCell<..>>, and items and fetches borrow freely from their surroundings (the request, the services, the transaction). That removes every lock, and it lets fetch closures borrow instead of cloning into Arcs.

The cost is that a region’s future is not Send, so it cannot go to tokio::spawn. Actix handlers already run on a single-threaded worker, so bonsapi needs nothing extra. Doc hub runs its jobs through a LocalPoolHandle for the same reason:

let pool = LocalPoolHandle::new(config.concurrency.max(1).into());
let job = pool.spawn_pinned(move || async move { handler.handle(job).await });

Sharing a transaction: Exclusive #

Several fetches often need the same &mut transaction, and Rust will not let several closures capture one &mut. Exclusive lends it to one fetch at a time:

pub struct Exclusive<T>(RefCell<Option<T>>);

impl<T> Exclusive<T> {
    pub fn take(&self) -> Taken<'_, T> {
        let value = self.0.borrow_mut().take();
        Taken { exclusive: self, value: Some(value.expect("another fetch is using this value")) }
    }
}

impl<T> Drop for Taken<'_, T> {
    fn drop(&mut self) {
        *self.exclusive.0.borrow_mut() = self.value.take();   // put it back
    }
}

In use (from status_update.rs, trimmed). Both fetches are handed the same &Exclusive, and each takes the transaction for the duration of its own queries:

let tx = &Exclusive::new(tx);                                   // tx: &mut T
autobatch::scope(async |mut batcher| {
    let plan_learning = batcher.define_batch("status_update.plan_learning", move |rows| {
        plan_learning(db, mq, tx, resource, entity_id, organization_id, rows)
    });
    let write = batcher.define_batch("status_update.write", move |prepared| {
        write_prepared(db, mq, tx, resource, entity_id, organization_id, prepared)
    });
    // ...
})

async fn plan_learning(db: &D, mq: &M, tx: &Exclusive<&mut T>, /* ... */ rows: Vec<_>)
    -> BulkOutput<(LoadedExtraction<R>, LearningPlan), Fail>
{
    let plans = ExtractionLearningService::new(db, mq)
        .plan_many(&mut **tx.take(), resource, entity_id, organization_id, &planning)
        .await
        .map_err(fail_whole_chunk)?;
    // ...
}

take() moves the value out and the guard puts it back when the fetch ends. Flushes never overlap, so only one fetch holds it at a time. Taking it while it is already taken is a bug and panics instead of deadlocking.

Serial batches #

define_serial_batch is define_batch with a fetch that loops:

pub fn define_serial_batch(&mut self, name, each: impl Fn(I) -> Fut) -> Batch<'region, I, O, E> {
    let each = Rc::new(each);
    self.define_batch(name, move |inputs| {
        let each = each.clone();
        async move {
            let mut results = Vec::with_capacity(inputs.len());
            for input in inputs {
                results.push(each(input).await);    // one at a time, in item order
            }
            Ok(results)
        }
    })
}

Items still meet at that step and it still runs in item order, but the work inside is one call per input. It is for side effects that have no bulk form or must keep their order, such as notifications and some queue messages.

Tracing #

run_per_item opens an autobatch.region span (rows). Every flush opens an autobatch.flush span, and the fetch runs inside it, so the queries a fetch issues nest under the flush that caused them:

autobatch.region rows=50
  autobatch.flush batch=transition.invoices.load   rows=50 flush=1
    db.invoices_by_ids
    db.latest_invoice_data_by_invoice_ids
  autobatch.flush batch=transition.invoices.write  rows=48 flush=2
    db.begin
    db.update_extraction_rows
    db.commit

The span names under each flush are illustrative. A batch name appearing twice in one region means items reached it in different rounds.

What it does not do #

  • No deduplication or caching. Two items asking for the same input produce two inputs in the fetch. A fetch that needs distinct keys dedupes them itself.
  • No batching across requests or chunks. A region is one chunk of one request.
  • No concurrency between fetches. Flushes run one after another. The saving comes from fewer round trips, not from parallelism.
  • No size limit. A flush carries whatever was queued. The commands bound it by running a region per chunk of 50 (COMMAND_CHUNK_SIZE), up to 300 per request.

How bonsapi uses it #

apps/bonsapi/src/service/extraction/utils.rs adds a thin layer. Errors are split into “this extraction failed” and “stop the chunk”:

pub enum Fail {
    Item(ApiErrorResponse),    // only this extraction's outcome is an error
    Chunk(ApiErrorResponse),   // database or infrastructure error: fail the whole chunk
}

pub async fn run_per_item<'region, X, T>(
    batcher: Batcher<'region, '_>,
    items: impl IntoIterator<Item = X>,
    per_item: impl AsyncFn(&ItemContext<'region>, X) -> std::result::Result<T, Fail>,
) -> Result<Vec<Result<T>>> {
    batcher
        .run_per_item(items, async |ctx, item| match per_item(ctx, item).await {
            Ok(value) => Ok(Ok(value)),
            Err(Fail::Item(error)) => Ok(Err(error)),   // the item "succeeds" with an Err inside
            Err(Fail::Chunk(error)) => Err(error),      // try_join_in_order stops here
        })
        .await
}

An item failure is turned into a successful item whose value is an Err, so the driver keeps going; only a chunk failure reaches try_join_in_order as an Err and stops the region.

A by-ids read can return fewer rows than it was given ids, but a fetch must return one result per input. align_to_ids puts rows back in input order:

async fn refresh(req: &EntityRequestContext<'_, X>, ids: Vec<Id<Extraction>>)
    -> BulkOutput<Option<Extraction>, Fail>
{
    let extractions = ExtractionReadService::new(req.infra.db)
        .extractions_by_ids(ids.iter().copied(), req.entity_id, req.organization_id)
        .await?;                                                  // `?` fails the whole flush
    let found = align_to_ids(&ids, extractions, Extraction::id);  // Vec<Result<Option<_>>>
    Ok(found.into_iter().map(|row| row.map_err(Fail::Item)).collect())
}

A real command, reextract (trimmed from reextract.rs):

let outcomes = autobatch::scope(async |mut batcher| {
    let execute = batcher.define_batch("reextract.execute", |requests| {
        execute(infra, execution, api_user, entity_id, organization_id, requests)
    });
    let notify = notifies(&mut batcher, "reextract.notify", req);
    run_per_item(batcher, &request.extractions, async |ctx, item| {
        let original = originals.get(&item.extraction_id).cloned().ok_or(ApiErrorResponse::NotFound)?;
        ExtractionCommandValidation::ensure_resource_type(&original, ApiErrorResponse::UnprocessableEntity)?;
        let request = ExtractionExecutionRequest::new(original, trigger(item));
        let prepared = execute.call(ctx, request).await?;   // one execute for the whole chunk
        notify.call(ctx, prepared.extraction.id()).await?;  // then the notifications
        Ok(prepared.extraction)
    })
    .await
})
.await?;

A missing or untyped extraction returns early with ? and never reaches execute; the rest meet there.

run_in_chunks and ExtractionCommandResponse::in_chunks run one region per chunk of 50 and collect one outcome per id. Each resource type gets its own set of batches (transition.invoices.*, transition.direct_expenses.*, transition.credit_notes.*, transition.bank_statements.*), so a mixed request does one write per type, each in its own transaction.

When adding a step: write it as a function from Vec<Input> to one result per input, define its batch in the position where it should run, and do its I/O for all inputs at once.