yo-cli 0.3.15

//! The `yodb` command line tool.
//!
//! Two subcommands: `check`, which is the M1 deliverable, and `serve`, which
//! puts the RESP engine on a socket. The others arrive with the milestones that
//! need them.
//!
//! Argument parsing is done by hand rather than with a library. That is a
//! choice worth defending exactly once, here: `yodb check` is the tool you run
//! when a database will not start, and a tool for that moment that pulls in a
//! dependency tree is a tool that can fail to build on the machine where you
//! need it. It parses six flags across two commands. When this grows a
//! benchmark runner and a config file the calculation changes and so should the
//! code.

mod check;
mod poll;
mod serve;
mod signal;
mod store;

use std::net::SocketAddr;
use std::path::PathBuf;
use std::process::ExitCode;

use check::Severity;

const USAGE: &str = "\
yodb, an embedded knowledge engine

usage:
  yodb check FILE [--quick] [--quiet]
  yodb serve [--bind ADDR] [--port PORT] [--unixsocket PATH] [--no-port]
             [--dir PATH] [--store PATH --maxmemory BYTES]

  check    read a .yo file and report anything wrong with it. Never writes.
             --quick   skip the records and read only the headers
             --quiet   print findings and the summary, nothing else

  serve    speak RESP on a socket, so a Redis client can talk to it.
             --bind        address to listen on, 127.0.0.1 by default
             --port        port to listen on, 6379 by default
             --unixsocket  also listen on a socket file, which skips the
                           TCP stack and is the faster way in for a client
                           on the same machine
             --no-port     no TCP at all, socket file only
             --dir         where the server writes, which is where BACKUP
                           puts its files and what CONFIG GET dir answers.
                           The directory the command was run from by default
             --maxmemory   how much memory to use before something has to
                           go, in the units CONFIG SET takes, so 100mb is
                           a hundred mebibytes and 100m is a hundred
                           million. No limit by default
             --store       a file to put cold values in when memory fills
                           up, instead of throwing keys away. The path has
                           to be a new one, because what a previous run
                           left in a store is reachable only through an
                           index that died with it. Needs --maxmemory,
                           since a server with no limit never fills up

environment:
  YO_ALLOC  what to do when a command path allocates. off by default, which
            is the check turned off. report prints each place it happens once
            and carries on. abort stops the process on the first one.

exit codes:
  0  nothing wrong
  1  something wrong
  2  the arguments did not make sense, or the file could not be read at all
";

/// What a Redis client tries first, so it is what we listen on.
const DEFAULT_PORT: u16 = 6379;

/// Loopback, not every interface.
///
/// Redis shipped bound to every interface for years, and the result was tens of
/// thousands of open databases on the internet. Reaching this server from
/// another machine should be a thing somebody typed on purpose.
const DEFAULT_BIND: &str = "127.0.0.1";

/// The allocator that enforces Y7, no heap on a command path.
///
/// Installed here because picking a global allocator belongs to the program and
/// not to any library it links. It forwards everything to the system allocator
/// and does nothing else until `YO_ALLOC` asks it to, so a release build of
/// `yodb` behaves exactly as it did before this line existed.
#[global_allocator]
static ALLOC: yo_alloc::YoAlloc = yo_alloc::YoAlloc::new();

fn main() -> ExitCode {
    // Before anything else, because it decides what happens for the rest of the
    // process and a value nobody understands has to be an error rather than a
    // quiet off. Somebody typing YO_ALLOC=abrot believes the check is running.
    if yo_alloc::set_mode_from_env().is_none() {
        eprintln!("yodb: YO_ALLOC is off, report or abort");
        return ExitCode::from(2);
    }

    let code = run();

    // The tally, for a run that reaches the end. `serve` normally does not,
    // because Ctrl-C ends the process rather than the loop, so in practice this
    // is for `check` and for a server that was told to stop. Each site already
    // printed itself on the way past.
    if yo_alloc::mode() == yo_alloc::Mode::Report {
        let (sites, total) = yo_alloc::seen();
        eprintln!("yodb: {total} allocation(s) on a command path, at {sites} place(s)");
    }
    code
}

fn run() -> ExitCode {
    let args: Vec<String> = std::env::args().skip(1).collect();
    let mut rest: Vec<&str> = args.iter().map(String::as_str).collect();

    match rest.first().copied() {
        Some("check") => {
            rest.remove(0);
            check_command(&rest)
        }
        Some("serve") => {
            rest.remove(0);
            serve_command(&rest)
        }
        Some("-h" | "--help") | None => {
            print!("{USAGE}");
            ExitCode::SUCCESS
        }
        Some("-V" | "--version") => {
            println!("yo {}", env!("CARGO_PKG_VERSION"));
            ExitCode::SUCCESS
        }
        Some(other) => {
            eprintln!("yo: no such command: {other}\n");
            eprint!("{USAGE}");
            ExitCode::from(2)
        }
    }
}

fn check_command(args: &[&str]) -> ExitCode {
    let mut path: Option<PathBuf> = None;
    let mut quick = false;
    let mut quiet = false;

    for a in args {
        match *a {
            "--quick" => quick = true,
            "--quiet" => quiet = true,
            "-h" | "--help" => {
                print!("{USAGE}");
                return ExitCode::SUCCESS;
            }
            other if other.starts_with('-') => {
                eprintln!("yodb check: no such option: {other}");
                return ExitCode::from(2);
            }
            other if path.is_none() => path = Some(PathBuf::from(other)),
            other => {
                eprintln!("yodb check: takes one file, and was also given {other}");
                return ExitCode::from(2);
            }
        }
    }

    let Some(path) = path else {
        eprintln!("yodb check: which file?\n");
        eprint!("{USAGE}");
        return ExitCode::from(2);
    };

    let report = match check::check(&path, !quick) {
        Ok(r) => r,
        Err(e) => {
            // Getting here means the file could not be opened far enough to say
            // anything at all, which is a different thing from a file with
            // problems in it and gets a different exit code.
            eprintln!("yodb check: {}: {e}", path.display());
            return ExitCode::from(2);
        }
    };

    if !quiet {
        println!("{}", path.display());
    }
    for f in &report.findings {
        println!("{f}");
    }

    let c = report.counts;
    if !quiet {
        // Off `quick` rather than off the counts. A segment that stops parsing
        // leaves the count at zero, and printing "records not walked" there
        // would say the walk did not happen when what happened is that it ran
        // and hit something.
        if quick {
            println!("{} segments, records not walked", c.regions);
        } else {
            println!(
                "{} segments, {} records, {} record bytes, {} dead",
                c.regions, c.records, c.record_bytes, c.dead_bytes
            );
        }
    }

    let errors = report.count(Severity::Error);
    let warns = report.count(Severity::Warn);
    if report.is_sound() {
        println!(
            "OK{}",
            if warns > 0 {
                format!(", with {warns} warning{}", plural(warns))
            } else {
                String::new()
            }
        );
        ExitCode::SUCCESS
    } else {
        println!("FAILED: {errors} problem{}", plural(errors));
        ExitCode::FAILURE
    }
}

fn serve_command(args: &[&str]) -> ExitCode {
    let mut bind = DEFAULT_BIND.to_string();
    let mut port = DEFAULT_PORT;
    let mut unixsocket: Option<std::path::PathBuf> = None;
    let mut tcp = true;
    let mut store: Option<std::path::PathBuf> = None;
    let mut maxmemory: Option<u64> = None;
    let mut dir: Option<std::path::PathBuf> = None;

    let mut at = 0;
    while at < args.len() {
        let arg = args[at];
        at += 1;
        match arg {
            "-h" | "--help" => {
                print!("{USAGE}");
                return ExitCode::SUCCESS;
            }
            "--no-port" => tcp = false,
            "--bind" | "--port" | "--unixsocket" | "--store" | "--maxmemory" | "--dir" => {
                let Some(value) = args.get(at) else {
                    eprintln!("yodb serve: {arg} needs a value");
                    return ExitCode::from(2);
                };
                at += 1;
                if arg == "--bind" {
                    bind = (*value).to_string();
                } else if arg == "--unixsocket" {
                    unixsocket = Some(std::path::PathBuf::from(*value));
                } else if arg == "--store" {
                    store = Some(std::path::PathBuf::from(*value));
                } else if arg == "--dir" {
                    dir = Some(std::path::PathBuf::from(*value));
                } else if arg == "--maxmemory" {
                    // The parser `CONFIG SET maxmemory` uses, so that the two
                    // ways of setting the same limit read it the same way.
                    match yo_resp::dispatch::parse_memory(value.as_bytes()) {
                        Some(n) => maxmemory = Some(n),
                        None => {
                            eprintln!("yodb serve: {value} is not an amount of memory");
                            return ExitCode::from(2);
                        }
                    }
                } else {
                    match value.parse() {
                        Ok(p) => port = p,
                        Err(_) => {
                            eprintln!("yodb serve: {value} is not a port");
                            return ExitCode::from(2);
                        }
                    }
                }
            }
            other => {
                eprintln!("yodb serve: no such option: {other}");
                return ExitCode::from(2);
            }
        }
    }

    let Ok(addr) = format!("{bind}:{port}").parse::<SocketAddr>() else {
        eprintln!("yodb serve: {bind} is not an address to listen on");
        return ExitCode::from(2);
    };
    if !tcp && unixsocket.is_none() {
        eprintln!("yodb serve: --no-port with no --unixsocket leaves nothing to connect to");
        return ExitCode::from(2);
    }
    if store.is_some() && maxmemory.is_none() {
        eprintln!("yodb serve: --store with no --maxmemory is a file nothing would ever be put in");
        return ExitCode::from(2);
    }
    // Checked and made absolute here rather than when a backup is asked for,
    // because the person who mistyped the path is still watching now and will
    // not be then. Joined onto the working directory rather than canonicalised,
    // since canonicalising on Windows produces a path with a prefix on it that
    // nobody expects to read back out of `CONFIG GET dir`.
    let dir = match dir {
        Some(path) if !path.is_dir() => {
            eprintln!(
                "yodb serve: {}: not a directory to write in",
                path.display()
            );
            return ExitCode::from(2);
        }
        Some(path) if path.is_absolute() => Some(path),
        Some(path) => Some(std::env::current_dir().unwrap_or_default().join(path)),
        None => None,
    };

    // Before the listener, because a file that cannot be made should not leave a
    // port bound behind it, and because failing here is a mistyped path and the
    // person who typed it is still watching.
    let opened = match &store {
        Some(path) => match store::Store::create(path) {
            Ok(s) => Some(s),
            Err(e) => {
                eprintln!("yodb serve: {}: {e}", path.display());
                return ExitCode::from(2);
            }
        },
        None => None,
    };

    let want = if tcp { Some(addr) } else { None };
    let mut server = match serve::Server::open(want, unixsocket.clone()) {
        Ok(s) => s,
        Err(e) => {
            eprintln!("yodb serve: {e}");
            return ExitCode::from(2);
        }
    };
    if let Some(limit) = maxmemory {
        server.set_maxmemory(limit);
    }
    if let Some(dir) = dir {
        server.set_dir(dir);
    }
    if let Some(opened) = opened {
        server.use_store(opened);
    }

    // What it actually bound to, which is the only way to find out when the
    // port asked for was zero.
    let version = env!("CARGO_PKG_VERSION");
    match (tcp, &unixsocket) {
        (true, Some(path)) => {
            let bound = server.local_addr().unwrap_or(addr);
            println!(
                "yodb {version} listening on {bound} and on {}",
                path.display()
            );
        }
        (true, None) => {
            let bound = server.local_addr().unwrap_or(addr);
            println!("yodb {version} listening on {bound}");
        }
        (false, Some(path)) => {
            println!("yodb {version} listening on {}", path.display());
        }
        (false, None) => unreachable!("refused above"),
    }

    // Both numbers, on purpose. A server that only printed the quarter would
    // leave the next person to wonder a quarter of what, and the answer to that
    // is the thing they need when the pool comes out the wrong size. See
    // `yo_resp::cap` for why it is a quarter at all.
    let cap = yo_resp::cap::cap();
    match cap.limit() {
        Some(limit) => println!(
            "yodb {version} may use {} and will size pools from {}",
            bytes(limit),
            bytes(cap.budget())
        ),
        None => println!("yodb {version} found no memory limit to size pools from"),
    }

    // Which of the two things a memory limit means here, said out loud at
    // startup, because they are opposites and the difference is a file.
    match (&store, maxmemory) {
        (Some(path), Some(limit)) => println!(
            "yodb {version} keeps {} in memory and moves the rest into {}",
            bytes(limit),
            path.display()
        ),
        (None, Some(limit)) => println!("yodb {version} evicts keys above {}", bytes(limit)),
        (_, None) => {}
    }

    // After the listening line and not before it, so a Ctrl-C that arrives in
    // the moment between the two is a process that was never told to serve
    // rather than one that says it is serving and then stops.
    signal::listen();
    let outcome = match server.run(signal::stop()) {
        Ok(()) => ExitCode::SUCCESS,
        Err(e) => {
            eprintln!("yodb serve: {e}");
            ExitCode::FAILURE
        }
    };
    // Explicitly, and before the line that says so, because dropping the server
    // is what unlinks the socket file and closes the doors. Leaving it to the
    // end of the function would print that it had shut down while the path it
    // was listening on was still there for somebody to connect to.
    drop(server);
    if signal::stopped() {
        println!("yodb {version} shutting down");
    }
    outcome
}

fn plural(n: usize) -> &'static str {
    if n == 1 { "" } else { "s" }
}

/// A byte count a person can read, for the startup lines only.
///
/// Powers of two with the short names, because that is what `maxmemory` takes
/// and a server that prints one unit and accepts another is a server that gets
/// misconfigured.
fn bytes(n: u64) -> String {
    const UNITS: [(u64, &str); 3] = [(1 << 30, "gb"), (1 << 20, "mb"), (1 << 10, "kb")];
    for (size, name) in UNITS {
        if n >= size {
            let whole = n / size;
            let tenths = (n % size) * 10 / size;
            return if tenths == 0 {
                format!("{whole}{name}")
            } else {
                format!("{whole}.{tenths}{name}")
            };
        }
    }
    format!("{n} bytes")
}

#[cfg(test)]
mod tests {
    use super::bytes;

    #[test]
    fn a_byte_count_prints_in_the_units_maxmemory_takes() {
        assert_eq!(bytes(0), "0 bytes");
        assert_eq!(bytes(512), "512 bytes");
        assert_eq!(bytes(1024), "1kb");
        assert_eq!(bytes(2 * 1024 * 1024 * 1024), "2gb");
        // One tenth is enough to tell 7.5gb from 7gb and not so much that the
        // line stops being readable.
        assert_eq!(bytes(7 * (1 << 30) + (1 << 29)), "7.5gb");
        assert_eq!(bytes(1536), "1.5kb");
    }
}