narju/

gate.rs

use tokio::net::{TcpListener, TcpStream};
use tokio::sync::mpsc::{channel, Receiver, Sender};

use crate::floor::host::Str;
use crate::floor::{list, list_elements, pair, Val};
use crate::link;
use crate::sched::Addr;
use crate::world::{reply, request, ConnId, Event};

/// The task a node is reached at before it has said anything. A convention,
/// because nothing else could be: an address is only learned by being told, so
/// a node just dialled is unreachable unless one address on it is known already.
pub const GREETER: usize = 0;

fn sym(s: &str) -> Val {
    Val::Sym(s.into())
}

fn throw(why: Val) -> Val {
    pair(sym("throw"), why)
}

fn io(e: &std::io::Error) -> Val {
    throw(pair(sym("io"), Str::val(e.to_string())))
}

/// `('tcp "host:port")`, or nothing this gate fronts.
fn where_to(spec: &Val) -> Option<String> {
    match list_elements(spec).as_deref() {
        Some([kind, at]) if *kind == sym("tcp") => Some(Str::of(at)?.to_string()),
        _ => None,
    }
}

/// `('host gate)`. The only thing that opens a link, so the only thing that
/// names one: the counter needs neither sharing nor an atomic, which is why a
/// listener's accepts come back here instead of becoming links where they
/// arrive. Numbering starts past the ids a launcher may have registered before
/// the loop began ([`crate::world::World::link`]).
pub async fn gate(mut inbox: Receiver<Val>, back: Sender<Event>) {
    let (accepts, mut accepted) = channel::<(TcpStream, Val)>(8);
    let mut next: ConnId = 16;
    loop {
        tokio::select! {
            msg = inbox.recv() => {
                let Some(msg) = msg else { return };
                let Some((from, id, ask)) = request(&msg) else { continue };
                let v = match list_elements(&ask).as_deref() {
                    Some([verb, spec]) if *verb == sym("dial") => {
                        match where_to(spec) {
                            Some(at) => dial(&at, &mut next, &back).await,
                            None => throw(pair(sym("no-transport"), spec.clone())),
                        }
                    }
                    Some([verb, spec]) if *verb == sym("listen") => {
                        match where_to(spec) {
                            Some(at) => listen(&at, &accepts, from.clone()).await,
                            None => throw(pair(sym("no-transport"), spec.clone())),
                        }
                    }
                    _ => throw(pair(sym("bad-request"), ask)),
                };
                if back.send(Event::Post(from, reply(id, v))).await.is_err() {
                    return;
                }
            }
            // The gate holds a sender, so this never closes.
            Some((sock, to)) = accepted.recv() => {
                let peer = attach(sock, &mut next, &back).await;
                let told = back.send(Event::Post(to, list(&[sym("peer"), peer]))).await;
                if told.is_err() {
                    return;
                }
            }
        }
    }
}

async fn dial(at: &str, next: &mut ConnId, back: &Sender<Event>) -> Val {
    match TcpStream::connect(at).await {
        Ok(sock) => attach(sock, next, back).await,
        Err(e) => io(&e),
    }
}

/// Bring a socket up as a link and answer the address of the greeter on it. The
/// registration and the reply go through the one queue so the first cannot be
/// overtaken by the second, which would answer `('no-link conn)` for a link that
/// exists. Nothing waits for the registration to be *drained*: that would be
/// waiting on the node this reply is for.
async fn attach(sock: TcpStream, next: &mut ConnId, back: &Sender<Event>) -> Val {
    let conn = *next;
    *next += 1;
    let out = link::bridge(conn, sock, back.clone());
    if back.send(Event::Link(conn, out)).await.is_err() {
        return throw(sym("no-world"));
    }
    Addr::Remote(conn, GREETER).val()
}

/// Answers the address bound rather than the one asked for: a port of 0 is how
/// a caller says it does not care which, and it has to be told the answer.
async fn listen(at: &str, accepts: &Sender<(TcpStream, Val)>, to: Val) -> Val {
    let bound = match TcpListener::bind(at).await {
        Ok(l) => l,
        Err(e) => return io(&e),
    };
    let here = match bound.local_addr() {
        Ok(a) => a.to_string(),
        Err(e) => return io(&e),
    };
    tokio::spawn(accepting(bound, accepts.clone(), to));
    Str::val(here)
}

/// The second arm is why this is not a bare loop: a listener otherwise outlives
/// the world it was opened for, parked on an `accept` for a node with nothing
/// left to hand the connection to.
async fn accepting(bound: TcpListener, accepts: Sender<(TcpStream, Val)>, to: Val) {
    loop {
        tokio::select! {
            got = bound.accept() => match got {
                Ok((sock, _)) => {
                    if accepts.send((sock, to.clone())).await.is_err() {
                        return;
                    }
                }
                // The listening socket itself failed, which no message can
                // repair - a per-connection error would have been an `Ok` here.
                Err(_) => return,
            },
            () = accepts.closed() => return,
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::floor::machine::{Machine, NoMail, Store};
    use crate::floor::Env;
    use crate::sched::addr;
    use crate::world::World;

    fn val(src: &str) -> Val {
        let Val::Code(e) = crate::surface::compile(src).expect("compiles") else {
            panic!("not code")
        };
        Machine::new(Env::default(), e)
            .run(&mut Store::<NoMail>::default())
            .expect("no fault")
    }

    fn ok(v: Val) -> Val {
        pair(sym("ok"), v)
    }

    #[tokio::test]
    async fn a_node_dials_itself_and_is_answered() {
        let mut w = World::new(100_000);
        w.heap()
            .spawn(val("(lambda g (_) (send (car (receive)) 'pong))"));
        w.heap().spawn(val("(lambda t (me)
               (begin
                 (send '(host gate) (list 'req me 1 '(listen (tcp \"127.0.0.1:0\"))))
                 (let ((at (car (cdr (cdr (receive))))))
                   (begin
                     (spawn (lambda d (dme)
                              (begin
                                (send '(host gate)
                                      (list 'req dme 2 (list 'dial (list 'tcp at))))
                                (let ((there (car (cdr (cdr (receive))))))
                                  (begin (send there (cons dme 'hi)) (receive))))))
                     (receive)))))"));

        let (tx, rx) = channel(8);
        w.serve("gate", tx);
        tokio::spawn(gate(rx, w.port()));

        let mut ended = w.run().await.expect("no bug");
        ended.sort_by_key(|(id, _)| *id);
        assert_eq!(
            ended,
            vec![
                (0, ok(Val::Nil)),
                (1, ok(list(&[sym("peer"), Addr::Remote(17, GREETER).val()]))),
                (2, ok(sym("pong"))),
            ]
        );
    }

    #[tokio::test]
    async fn a_name_is_looked_up_across_a_link_and_used() {
        let mut t = crate::tower::Tower::load().expect("loads");
        let mut entry = |src: &str| {
            let forms = crate::floor::read::read(src).expect("reads");
            t.task(&crate::surface::desugar_body(&forms).expect("desugars"))
                .expect("no fault")
        };

        let greeter = entry(
            "(greet self
               (list (cons 'echo
                           (spawn (lambda (me)
                                    (task-loop me
                                               (lambda h (st msg)
                                                 (begin (reply msg 'pong)
                                                        (cons 'stop 'ok)))
                                               '()))))))",
        );
        let driver = entry(
            "(task-loop self
               (lambda h (st msg)
                 (if (nil? st)
                     (let ((at (call '(host gate) '(listen (tcp \"127.0.0.1:0\")))))
                       (let ((d (spawn (lambda (me)
                                         (task-loop me
                                           (lambda h2 (st2 msg2)
                                             (let ((there (call '(host gate)
                                                                (list 'dial (list 'tcp at)))))
                                               (let ((echo (call there '(lookup echo))))
                                                 (let ((got (call echo 'ping)))
                                                   (begin (call there '(stop))
                                                          (cons 'stop got))))))
                                           '())))))
                         (begin (send d 'go) (cons 'listening at))))
                     (cons 'stop msg)))
               '())",
        );

        let mut w = World::new(50_000_000);
        w.heap().spawn(greeter);
        let drive = w.heap().spawn(driver);
        w.heap().post(&addr(drive), sym("go"));

        let (tx, rx) = channel(8);
        w.serve("gate", tx);
        tokio::spawn(gate(rx, w.port()));

        let mut ended = w.run().await.expect("no bug");
        ended.sort_by_key(|(id, _)| *id);
        assert_eq!(
            ended,
            vec![
                (0, ok(sym("ok"))),
                (1, ok(list(&[sym("peer"), Addr::Remote(17, GREETER).val()]))),
                (2, ok(sym("ok"))),
                (3, ok(sym("pong"))),
            ]
        );
    }

    #[tokio::test]
    async fn a_transport_that_is_not_fronted_is_refused() {
        let mut w = World::new(1000);
        let id = w.heap().spawn(val("(lambda t (me)
               (begin (send '(host gate) (list 'req me 1 '(dial (carrier \"pigeon\"))))
                      (receive)))"));

        let (tx, rx) = channel(8);
        w.serve("gate", tx);
        tokio::spawn(gate(rx, w.port()));

        assert_eq!(
            w.run().await.expect("no bug"),
            vec![(
                id,
                ok(list(&[
                    sym("reply"),
                    Val::Num(1),
                    throw(pair(
                        sym("no-transport"),
                        list(&[sym("carrier"), Str::val("pigeon")])
                    ))
                ]))
            )]
        );
    }
}