flow-mumu 0.2.0-rc.9

// src/drain_async.rs
//
// flow:drain_async(source [, callback])
//   • Spawns a background poller that drains `source` in **micro-bursts per tick**.
//     Each tick processes up to a configurable **item burst cap** OR until a tiny
//     **time budget** is reached, whichever happens first.
//   • If `callback` (Function) is provided, it is called once at the end
//     with the total count: callback(count).
//   • Returns immediately (Int(0)).
//
// IMPORTANT — cooperative throttling:
//   Some stages (e.g. flow:throttle) may deliberately return a **temporary no-data**
//   signal instead of EOF. We standardise this as the error string: "AGAIN".
//   This poller treats "AGAIN" as **not done**. To prevent the interpreter’s
//   idle cut-off from exiting too early while we are simply waiting for the
//   throttle window, we return a **keep-alive** (processed=1) when a tick sees
//   only "AGAIN" responses.
//
// Tuning (via environment variables; optional):
//   • LAVA_FLOW_BURST       — max items processed per poll tick (default: 64)
//   • LAVA_FLOW_BUDGET_US   — per-tick time budget in microseconds (default: 500)

use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};

use mumu::parser::interpreter::{DynamicFnInfo, Interpreter, PollerFn};
use mumu::parser::interpreter::apply::apply_n_ary_function_value;
use mumu::parser::types::{FunctionValue, Value};

use crate::common::to_transform;

pub fn register_flow_drain_async(interp: &mut Interpreter) {
    let f = Arc::new(Mutex::new(flow_drain_async_bridge));
    let info = DynamicFnInfo::new(f, true);
    interp.register_dynamic_function_ex("flow:drain_async", info);

    interp.set_variable(
        "flow:drain_async",
        Value::Function(Box::new(FunctionValue::Named("flow:drain_async".to_string()))),
    );
}

fn flow_drain_async_bridge(interp: &mut Interpreter, args: Vec<Value>) -> Result<Value, String> {
    match args.len() {
        1 => spawn_drain(interp, args[0].clone(), None),
        2 => {
            let src = args[0].clone();
            let cb  = match &args[1] {
                Value::Function(fb) => Some(fb.clone()),
                other => return Err(format!("flow:drain_async => callback must be Function, got {:?}", other)),
            };
            spawn_drain(interp, src, cb)
        }
        n => Err(format!("flow:drain_async => expected 1 or 2 args (source [, callback]), got {}", n)),
    }
}

fn spawn_drain(interp: &mut Interpreter, source: Value, cb: Option<Box<FunctionValue>>) -> Result<Value, String> {
    // Normalize Iterator or zero-arg Function into a pullable transform
    let tf = to_transform(source)?;

    // State shared with the poller
    struct DrainState {
        tf: Box<FunctionValue>,
        count: i32,
        done: bool,
        callback: Option<Box<FunctionValue>>,
        callback_fired: bool,
    }

    let state = Arc::new(Mutex::new(DrainState {
        tf,
        count: 0,
        done: false,
        callback: cb,
        callback_fired: false,
    }));

    // Micro-burst tuning (read once here; captured by the poller)
    let burst_per_tick: usize = std::env::var("LAVA_FLOW_BURST")
        .ok()
        .and_then(|s| s.parse::<usize>().ok())
        .unwrap_or(64);
    let budget_us: u64 = std::env::var("LAVA_FLOW_BUDGET_US")
        .ok()
        .and_then(|s| s.parse::<u64>().ok())
        .unwrap_or(500);
    let time_budget = Duration::from_micros(budget_us);

    // Poller: drain up to `burst_per_tick` items or until `time_budget` elapses
    let st_for_poller = Arc::clone(&state);
    let poller: PollerFn = Arc::new(Mutex::new(move |interp: &mut Interpreter| {
        let tick_start = Instant::now();
        let mut processed = 0usize;
        let mut saw_again = false;
        let mut cb_to_fire: Option<(Box<FunctionValue>, i32)> = None;

        // Early exit if already done; fire callback (0) if not fired yet.
        {
            let mut guard = st_for_poller.lock().unwrap();
            if guard.done {
                if !guard.callback_fired {
                    let cb = guard.callback.take();
                    guard.callback_fired = true;
                    if let Some(cb_fn) = cb {
                        cb_to_fire = Some((cb_fn, 0));
                    }
                }
            }
        }
        if let Some((cb_fn, c)) = cb_to_fire.take() {
            let _ = apply_n_ary_function_value(interp, cb_fn, vec![Value::Int(c)]);
            return 0;
        }

        // Main micro-burst loop
        loop {
            // Pull one item without holding the state lock during the call.
            let tf_clone = {
                let guard = st_for_poller.lock().unwrap();
                if guard.done {
                    break;
                }
                guard.tf.clone()
            };

            match apply_n_ary_function_value(interp, tf_clone, vec![]) {
                Ok(_v) => {
                    // Successfully drained an item
                    let mut guard = st_for_poller.lock().unwrap();
                    guard.count = guard.count.saturating_add(1);
                    drop(guard);
                    processed += 1;
                }
                Err(e) if e == "AGAIN" => {
                    // Temporary no-data; not EOF. Note for keep-alive and stop this tick.
                    saw_again = true;
                    break;
                }
                Err(e) if e == "NO_MORE_DATA" => {
                    // True EOF for this pipeline run. Mark done and fire callback.
                    let mut guard = st_for_poller.lock().unwrap();
                    guard.done = true;
                    let total = guard.count;
                    let cb = guard.callback.take();
                    guard.callback_fired = true;
                    drop(guard);
                    if let Some(cb_fn) = cb {
                        let _ = apply_n_ary_function_value(interp, cb_fn, vec![Value::Int(total)]);
                    }
                    break;
                }
                Err(_e) => {
                    // Treat any other error as terminal for safety
                    let mut guard = st_for_poller.lock().unwrap();
                    guard.done = true;
                    drop(guard);
                    break;
                }
            }

            // Stop this tick if we hit the burst cap or exceeded the tiny time budget
            if processed >= burst_per_tick || tick_start.elapsed() >= time_budget {
                break;
            }
        }

        // Keep-alive: if this tick saw only "AGAIN", report 1 so the outer
        // idle-round logic does not cut us off just before the throttle window.
        if processed == 0 && saw_again {
            return 1;
        }

        processed
    }));

    interp.add_poller(poller);
    // Non-blocking return; for flow:compose we need a Value; 0 is harmless.
    Ok(Value::Int(0))
}