improvise/crates/improvise-incremental/tests/sched_dag.rs

//! A 7-unit DAG with fan-out and fan-in on real threads and on virtual workers, against the same
//! graph in one circuit (from the review of #364). Evidence: `context/test-evidence/improvise-incremental.md`.

mod support;

use improvise_incremental::sched::{Claimed, Executor, Pool, Turn, Unit};
use std::sync::{Arc, Mutex};
use support::watchdog::within;
use zset::{
    ChannelPipe, Circuit, Direct, EdgeSink, EdgeSource, Frame, Output, Stream, ZSet, channel_pipes,
};

type P = ChannelPipe<Frame<i64>>;
type Z = ZSet<i64>;

struct Rng(u64);
impl Rng {
    fn next(&mut self) -> u64 {
        self.0 ^= self.0 << 13;
        self.0 ^= self.0 >> 7;
        self.0 ^= self.0 << 17;
        self.0
    }
    fn below(&mut self, n: u64) -> u64 {
        self.next() % n
    }
}

fn st_a(c: &mut Circuit, s: &Stream<Z>) -> Stream<Z> {
    c.map(s, |x| x * 2)
}
fn st_b(c: &mut Circuit, s: &Stream<Z>) -> Stream<Z> {
    c.filter(s, |x| x % 3 != 0)
}
fn st_c(c: &mut Circuit, a: &Stream<Z>, b: &Stream<Z>) -> Stream<Z> {
    let u = c.plus(a, b);
    let k = c.map(&u, |x| (x.rem_euclid(7), *x));
    let g = c.aggregate(&k, |_, vs| Some(vs.iter().map(|(v, w)| v * w).sum::<i64>()));
    c.map(&g, |(k, s)| k * 1_000_000 + s)
}
fn st_d(c: &mut Circuit, s: &Stream<Z>) -> Stream<Z> {
    let m = c.map(s, |x| x.rem_euclid(1000));
    c.distinct(&m)
}
fn st_e(c: &mut Circuit, s: &Stream<Z>) -> Stream<Z> {
    c.filter(s, |x| x % 2 == 0)
}
fn st_f(c: &mut Circuit, s: &Stream<Z>) -> Stream<Z> {
    c.map(s, |x| x / 2)
}

fn reference(cmds: &[Z]) -> (Z, Z) {
    let mut c = Circuit::new();
    let (i, s) = c.input::<i64>();
    let a = st_a(&mut c, &s);
    let b = st_b(&mut c, &s);
    let cc = st_c(&mut c, &a, &b);
    let d = st_d(&mut c, &cc);
    let e = st_e(&mut c, &cc);
    let f = st_f(&mut c, &e);
    let (od, of) = (c.output(&d), c.output(&f));
    for d in cmds {
        c.push(&i, d.clone());
        c.step();
    }
    (c.take(&od), c.take(&of))
}

struct Head {
    sinks: Vec<EdgeSink<i64, Direct, P>>,
    inbox: Arc<Mutex<Vec<Z>>>,
    around: Vec<usize>,
}
impl Unit for Head {
    fn turn(&mut self, _: usize) -> Turn {
        let before: u64 = self.sinks.iter().map(|s| s.traffic()).sum();
        let cmds = std::mem::take(&mut *self.inbox.lock().unwrap());
        for d in cmds {
            for s in &mut self.sinks {
                s.committed(d.clone());
            }
        }
        for s in &mut self.sinks {
            s.pump();
        }
        let after: u64 = self.sinks.iter().map(|s| s.traffic()).sum();
        Turn {
            progress: after != before,
            more: false,
        }
    }
    fn neighbours(&self) -> &[usize] {
        &self.around
    }
}

struct Node {
    c: Circuit,
    sources: Vec<EdgeSource<i64, Direct, P>>,
    sinks: Vec<(EdgeSink<i64, Direct, P>, Output<i64>)>,
    collect: Option<(Output<i64>, Arc<Mutex<Z>>)>,
    around: Vec<usize>,
}
impl Node {
    fn traffic(&self) -> u64 {
        self.sources.iter().map(|s| s.traffic()).sum::<u64>()
            + self.sinks.iter().map(|(s, _)| s.traffic()).sum::<u64>()
    }
}
impl Unit for Node {
    fn turn(&mut self, budget: usize) -> Turn {
        let before = self.traffic();
        let mut read = 0;
        let mut more = false;
        for s in &mut self.sources {
            let r = s.poll(&mut self.c, budget);
            // Only a leaf (nothing downstream to kick it back) asks to run again at its budget.
            more |= self.sinks.is_empty() && r == budget;
            read += r;
        }
        if read > 0 {
            self.c.begin();
            self.c.step();
            self.c.commit();
            for (s, o) in &mut self.sinks {
                s.committed_from(&mut self.c, o);
            }
            if let Some((o, z)) = &self.collect {
                *z.lock().unwrap() += self.c.take(o);
            }
        }
        for (s, _) in &mut self.sinks {
            s.pump();
        }
        Turn {
            progress: self.traffic() != before,
            more,
        }
    }
    fn neighbours(&self) -> &[usize] {
        &self.around
    }
}

struct Dag {
    exec: Arc<Executor>,
    inbox: Arc<Mutex<Vec<Z>>>,
    d: Arc<Mutex<Z>>,
    f: Arc<Mutex<Z>>,
}

fn dag(rng: &mut Rng, budget: usize) -> Dag {
    let mut cap = || 1 + rng.below(3) as usize;
    let (h_a, a_in) = channel_pipes::<Frame<i64>>(cap());
    let (h_b, b_in) = channel_pipes::<Frame<i64>>(cap());
    let (a_c, c_in_a) = channel_pipes::<Frame<i64>>(cap());
    let (b_c, c_in_b) = channel_pipes::<Frame<i64>>(cap());
    let (c_d, d_in) = channel_pipes::<Frame<i64>>(cap());
    let (c_e, e_in) = channel_pipes::<Frame<i64>>(cap());
    let (e_f, f_in) = channel_pipes::<Frame<i64>>(cap());
    let inbox = Arc::new(Mutex::new(Vec::new()));
    let (dz, fz) = (
        Arc::new(Mutex::new(Z::new())),
        Arc::new(Mutex::new(Z::new())),
    );
    let head = Head {
        sinks: vec![EdgeSink::new(Direct, h_a, 1), EdgeSink::new(Direct, h_b, 1)],
        inbox: inbox.clone(),
        around: vec![1, 2],
    };
    let unary = |pipe_in: P,
                 f: fn(&mut Circuit, &Stream<Z>) -> Stream<Z>,
                 out: Option<P>,
                 col: Option<Arc<Mutex<Z>>>,
                 around: Vec<usize>| {
        let mut c = Circuit::new();
        let (i, s) = c.input::<i64>();
        let o = f(&mut c, &s);
        let o = c.output(&o);
        let source = EdgeSource::new(Direct, pipe_in, i);
        let (sinks, collect) = match (out, col) {
            (Some(p), None) => (vec![(EdgeSink::new(Direct, p, 1), o)], None),
            (None, Some(z)) => (vec![], Some((o, z))),
            _ => unreachable!(),
        };
        Node {
            c,
            sources: vec![source],
            sinks,
            collect,
            around,
        }
    };
    let a = unary(a_in, st_a, Some(a_c), None, vec![0, 3]);
    let b = unary(b_in, st_b, Some(b_c), None, vec![0, 3]);
    let cnode = {
        let mut c = Circuit::new();
        let (ia, sa) = c.input::<i64>();
        let (ib, sb) = c.input::<i64>();
        let o = st_c(&mut c, &sa, &sb);
        let o1 = c.output(&o);
        let o2 = c.output(&o);
        Node {
            c,
            sources: vec![
                EdgeSource::new(Direct, c_in_a, ia),
                EdgeSource::new(Direct, c_in_b, ib),
            ],
            sinks: vec![
                (EdgeSink::new(Direct, c_d, 1), o1),
                (EdgeSink::new(Direct, c_e, 1), o2),
            ],
            collect: None,
            around: vec![1, 2, 4, 5],
        }
    };
    let d = unary(d_in, st_d, None, Some(dz.clone()), vec![3]);
    let e = unary(e_in, st_e, Some(e_f), None, vec![3, 6]);
    let f = unary(f_in, st_f, None, Some(fz.clone()), vec![5]);
    let units: Vec<Box<dyn Unit>> = vec![
        Box::new(head),
        Box::new(a),
        Box::new(b),
        Box::new(cnode),
        Box::new(d),
        Box::new(e),
        Box::new(f),
    ];
    let exec = Arc::new(Executor::new(units, budget));
    exec.kick_all();
    Dag {
        exec,
        inbox,
        d: dz,
        f: fz,
    }
}

fn commands(rng: &mut Rng, n: usize) -> Vec<Z> {
    (0..n)
        .map(|_| {
            let k = rng.below(4);
            (0..k)
                .map(|_| (rng.below(40) as i64 - 20, rng.below(5) as i64 - 2))
                .collect()
        })
        .collect()
}

fn submit(g: &Dag, d: Z) {
    g.inbox.lock().unwrap().push(d);
    g.exec.kick(0);
}

fn out(g: &Dag) -> (Z, Z) {
    (g.d.lock().unwrap().clone(), g.f.lock().unwrap().clone())
}

#[test]
fn a_dag_on_real_threads_equals_one_circuit_and_idle_means_done() {
    within(30, "the DAG on a pool", || {
        let mut rng = Rng(0x9e3779b97f4a7c15);
        for round in 0..8 {
            let workers = 1 + round % 8;
            let budget = 1 + rng.below(4) as usize;
            let g = dag(&mut rng, budget);
            let cmds = commands(&mut rng, 300);
            let pool = Pool::start(g.exec.clone(), workers);
            // Concurrent submitter thread, with idle checkpoints from the host.
            let mut at = 0;
            while at < cmds.len() {
                let chunk = 1 + rng.below(80) as usize;
                let end = (at + chunk).min(cmds.len());
                for d in &cmds[at..end] {
                    submit(&g, d.clone());
                    if rng.below(7) == 0 {
                        std::thread::yield_now();
                    }
                }
                at = end;
                if rng.below(3) == 0 {
                    pool.executor().wait_idle().unwrap();
                    assert_eq!(
                        out(&g),
                        reference(&cmds[..at]),
                        "[idle] round {round} w{workers} at {at}"
                    );
                }
            }
            pool.finish().unwrap();
            assert_eq!(
                out(&g),
                reference(&cmds),
                "[final] round {round} w{workers} b{budget}"
            );
        }
    });
}

#[test]
fn submitter_threads_race_the_pool() {
    within(30, "racing submitters", || {
        let mut rng = Rng(0x1234567);
        for workers in 1..=4 {
            let g = Arc::new(dag(&mut rng, 2));
            let cmds = commands(&mut rng, 600);
            let pool = Pool::start(g.exec.clone(), workers);
            // Four submitters; order between them is free, the result is a sum so order-free.
            let parts: Vec<Vec<Z>> = (0..4)
                .map(|k| cmds.iter().skip(k).step_by(4).cloned().collect())
                .collect();
            let hs: Vec<_> = parts
                .into_iter()
                .map(|p| {
                    let g = g.clone();
                    std::thread::spawn(move || {
                        for d in p {
                            submit(&g, d);
                        }
                    })
                })
                .collect();
            for h in hs {
                h.join().unwrap();
            }
            pool.finish().unwrap();
            // `distinct` and aggregates are functions of the integral, so order-free.
            assert_eq!(out(&g), reference(&cmds), "[race] w{workers}");
        }
    });
}

fn run_virtual(exec: &Executor, workers: usize, rng: &mut Rng, steps: usize, drain: bool) {
    let mut held: Vec<(Claimed, Turn)> = Vec::new();
    let mut n = 0;
    loop {
        if !drain && n >= steps {
            break;
        }
        n += 1;
        assert!(n < 2_000_000, "never quiet");
        let c = rng.next() as usize;
        if held.len() < workers
            && (held.is_empty() || !c.is_multiple_of(3))
            && let Some(mut cl) = exec.claim(|k| c / 3 % k)
        {
            let t = cl.run(exec.budget());
            held.push((cl, t));
            continue;
        }
        if held.is_empty() {
            if drain {
                break;
            }
            continue;
        }
        let (cl, t) = held.remove(c / 3 % held.len());
        exec.finish(cl, t);
    }
    for (cl, t) in held {
        exec.finish(cl, t);
    }
}

#[test]
fn a_dag_on_virtual_workers_in_any_order_equals_one_circuit() {
    let mut rng = Rng(0xdeadbeef);
    for round in 0..60 {
        let workers = 1 + round % 8;
        let budget = 1 + rng.below(3) as usize;
        let g = dag(&mut rng, budget);
        let n = 1 + rng.below(40) as usize;
        let cmds = commands(&mut rng, n);
        for d in &cmds {
            submit(&g, d.clone());
            let s = rng.below(10) as usize;
            run_virtual(&g.exec, workers, &mut rng, s, false);
        }
        run_virtual(&g.exec, workers, &mut rng, 0, true);
        assert!(g.exec.idle());
        assert_eq!(out(&g), reference(&cmds), "[virt] round {round}");
    }
}