flow-mumu 0.2.0-rc.9

// src/throttle.rs
//
// Non-blocking throttle for Flow pipelines.
// Each **instance** has its own timing state, so multiple pipelines can run
// at different speeds concurrently.
//
// IMPORTANT: Do **not** sleep inside the transform. Sleeping blocks the single-threaded
// poll loop and starves other pipelines. Instead, when the period hasn't elapsed yet,
// return the **temporary no-data** signal "AGAIN". The async drains (e.g. flow:drain_async)
// are aware of "AGAIN" and will keep the interpreter alive until the next window.

use mumu::parser::interpreter::{Interpreter, DynamicFnInfo};
use mumu::parser::types::{
    Value,
    FunctionValue,
    IteratorHandle,
    IteratorKind,
};
use mumu::parser::interpreter::apply::apply_n_ary_function_value;
use std::sync::{Arc, Mutex};
use std::time::{Instant, Duration};

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

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

fn flow_throttle_bridge_fn(_interp: &mut Interpreter, args: Vec<Value>) -> Result<Value, String> {
    match args.len() {
        0 => Ok(make_partial(None, None)),
        1 => {
            let one_val = &args[0];
            if is_placeholder(one_val) {
                Ok(make_partial(None, None))
            } else if let Value::Int(ms_i) = one_val {
                Ok(make_partial(Some(*ms_i), None))
            } else {
                Err(format!("flow:throttle => first param must be an int or '_', got {:?}", one_val))
            }
        }
        2 => {
            let a_ms = &args[0];
            let a_src = &args[1];

            if is_placeholder(a_ms) || is_placeholder(a_src) {
                let ms_opt = match a_ms {
                    Value::Int(n) => Some(*n),
                    _ if is_placeholder(a_ms) => None,
                    other => return Err(format!("flow:throttle => first param must be int or '_', got {:?}", other)),
                };
                let src_opt = if is_placeholder(a_src) { None } else { Some(a_src.clone()) };

                if ms_opt.is_some() && src_opt.is_some() {
                    let ms_val = ms_opt.unwrap();
                    let src_val = src_opt.unwrap();
                    finalize_throttle(ms_val, src_val)
                } else {
                    Ok(make_partial(ms_opt, src_opt))
                }
            } else {
                if let Value::Int(m) = a_ms {
                    finalize_throttle(*m, a_src.clone())
                } else {
                    Err(format!("flow:throttle => first param must be int, got {:?}", a_ms))
                }
            }
        }
        n => Err(format!("flow:throttle => up to 2 arguments, got {}", n)),
    }
}

fn finalize_throttle(ms: i32, source_val: Value) -> Result<Value, String> {
    if ms < 0 {
        return Err(format!("flow:throttle => negative ms not allowed: {}", ms));
    }
    match source_val {
        Value::Iterator(h)   => Ok(build_throttle_iter(ms, h)),
        Value::Function(tf)  => Ok(build_throttle_chain(ms, tf)),
        other => Err(format!("flow:throttle => second param must be Iterator or transform Function, got {:?}", other)),
    }
}

/// Re-entrant, stateless partial:
/// Captures the initial (ms, src) as *defaults*, folds any new args,
/// and either finalizes or returns a brand-new partial with the new defaults.
fn make_partial(init_ms: Option<i32>, init_src: Option<Value>) -> Value {
    use mumu::parser::types::FunctionValue::RustClosure;

    Value::Function(Box::new(RustClosure(
        "flow:throttle-partial".to_string(),
        Arc::new(Mutex::new(move |_interp: &mut Interpreter, new_args: Vec<Value>| {
            // Start from captured defaults each invocation
            let mut ms = init_ms;
            let mut src = init_src.clone();

            for val in new_args {
                if ms.is_none() {
                    if is_placeholder(&val) {
                        // keep None
                    } else if let Value::Int(i) = val {
                        ms = Some(i);
                    } else {
                        return Err(format!("flow:throttle => first param must be int or '_', got {:?}", val));
                    }
                } else if src.is_none() {
                    if is_placeholder(&val) {
                        // keep None
                    } else {
                        src = Some(val);
                    }
                } else {
                    return Err("flow:throttle => partial => too many arguments".to_string());
                }
            }

            if let (Some(m), Some(s)) = (ms, src.clone()) {
                finalize_throttle(m, s)
            } else {
                Ok(make_partial(ms, src))
            }
        })),
        0
    )))
}

/* ──────────────────────────────────────────────────────────────────────────
   NON-BLOCKING throttle implementations
   If the period hasn't elapsed yet, we immediately return **"AGAIN"** (not EOF).
   This yields control back to the poll loop so other pipelines can progress.
   ───────────────────────────────────────────────────────────────────────── */

fn build_throttle_iter(ms: i32, handle: IteratorHandle) -> Value {
    use mumu::parser::types::FunctionValue::RustClosure;

    // Set last_item_time so the first item can pass immediately.
    let period = Duration::from_millis(ms as u64);
    let env = ThrottleIterEnv {
        ms_delay: ms,
        upstream: handle,
        last_item_time: Instant::now() - period,
    };
    let shared = Arc::new(Mutex::new(env));

    let clos = RustClosure(
        "flow:throttle-iterator".to_string(),
        Arc::new(Mutex::new(move |_itp: &mut Interpreter, _a: Vec<Value>| {
            // Quick gate WITHOUT sleeping.
            let can_emit = {
                let guard = shared
                    .lock()
                    .map_err(|_| "flow:throttle => lock error (iterator)".to_string())?;
                let period = Duration::from_millis(guard.ms_delay as u64);
                Instant::now().duration_since(guard.last_item_time) >= period
            };

            if !can_emit {
                return Err("AGAIN".to_string());
            }

            // Pull from upstream now.
            let res = {
                let guard = shared
                    .lock()
                    .map_err(|_| "flow:throttle => lock error (iterator)".to_string())?;
                match &guard.upstream.kind {
                    IteratorKind::Core(state_arc) => {
                        let mut i_guard = state_arc
                            .lock()
                            .map_err(|_| "flow:throttle => lock error on Iterator".to_string())?;
                        if i_guard.done || i_guard.current >= i_guard.end {
                            i_guard.done = true;
                            Err("NO_MORE_DATA".to_string())
                        } else {
                            let val = Value::Int(i_guard.current);
                            i_guard.current += 1;
                            if i_guard.current >= i_guard.end {
                                i_guard.done = true;
                            }
                            Ok(val)
                        }
                    }
                    IteratorKind::Plugin(plugin_arc) => {
                        let mut plugin = plugin_arc
                            .lock()
                            .map_err(|_| "flow:throttle => plugin lock error".to_string())?;
                        match plugin.next_value() {
                            Ok(val) => Ok(val),
                            Err(e) => Err(e),
                        }
                    }
                }
            };

            // On success, update the gate time.
            if res.is_ok() {
                let mut guard = shared
                    .lock()
                    .map_err(|_| "flow:throttle => lock error (iterator set)".to_string())?;
                guard.last_item_time = Instant::now();
            }

            res
        })),
        0
    );

    Value::Function(Box::new(clos))
}

fn build_throttle_chain(ms: i32, upstream_tf: Box<FunctionValue>) -> Value {
    use mumu::parser::types::FunctionValue::RustClosure;

    let period = Duration::from_millis(ms as u64);
    let env = ThrottleChainEnv {
        ms_delay: ms,
        upstream: upstream_tf,
        last_item_time: Instant::now() - period,
    };
    let shared = Arc::new(Mutex::new(env));

    let clos = RustClosure(
        "flow:throttle-transform".to_string(),
        Arc::new(Mutex::new(move |interp: &mut Interpreter, _args: Vec<Value>| {
            // Quick gate WITHOUT sleeping.
            let can_emit = {
                let guard = shared
                    .lock()
                    .map_err(|_| "flow:throttle => lock error (chain)".to_string())?;
                let period = Duration::from_millis(guard.ms_delay as u64);
                Instant::now().duration_since(guard.last_item_time) >= period
            };

            if !can_emit {
                return Err("AGAIN".to_string());
            }

            // Pull one from upstream
            let upstream_clone = {
                let guard = shared
                    .lock()
                    .map_err(|_| "flow:throttle => lock error (chain)".to_string())?;
                guard.upstream.clone()
            };
            let res = apply_n_ary_function_value(interp, upstream_clone, vec![]);

            // On success, update the gate time.
            if res.is_ok() {
                let mut guard = shared
                    .lock()
                    .map_err(|_| "flow:throttle => lock error (chain set)".to_string())?;
                guard.last_item_time = Instant::now();
            }

            res
        })),
        0
    );

    Value::Function(Box::new(clos))
}

struct ThrottleIterEnv {
    ms_delay: i32,
    upstream: IteratorHandle,
    last_item_time: std::time::Instant,
}

struct ThrottleChainEnv {
    ms_delay: i32,
    upstream: Box<FunctionValue>,
    last_item_time: std::time::Instant,
}

fn is_placeholder(v: &Value) -> bool {
    match v {
        Value::Placeholder => true,
        Value::SingleString(s) if s == "_" => true,
        Value::StrArray(ss) if ss.len() == 1 && ss[0] == "_" => true,
        _ => false,
    }
}