luby/src/main.rs

use clap::Parser;
use mimalloc::MiMalloc;
use owo_colors::OwoColorize;
use std::{backtrace::Backtrace, fmt::Display, net::SocketAddr, num::NonZero, panic, thread};
use sysexits::ExitCode;

#[global_allocator]
static GLOBAL: MiMalloc = MiMalloc;

fn panic_cb(panic: &panic::PanicHookInfo) {
    let trace = Backtrace::force_capture();
    let location = panic.location().unwrap();
    let payload = panic.payload();
    let msg = if let Some(s) = payload.downcast_ref::<&'static str>() {
        s
    } else if let Some(s) = payload.downcast_ref::<String>() {
        s.as_str()
    } else {
        "unknown error"
    };

    eprint!(
        "{}:\n{trace}",
        format_args!(
            "thread '{}' panicked at {location}: {msg}",
            thread::current().name().unwrap_or("<unnamed>")
        )
        .red()
    );

    eprintln!(
        "{}",
        "This is a bug in luby. Please kindly report this at https://git.lua.re/luaneko/luby."
            .yellow()
    );
}

#[derive(Debug, Parser)]
struct Args {
    /// Paths to scripts to execute.
    #[clap(value_name = "SCRIPTS")]
    path: Vec<String>,

    /// Strings to execute.
    #[clap(long, short = 'e', value_name = "CHUNK")]
    eval: Vec<String>,

    /// Libraries to require on startup.
    #[clap(long, short = 'l', value_name = "NAME")]
    lib: Vec<String>,

    /// Console log level.
    #[clap(long, value_name = "LEVEL", default_value = "debug")]
    log: tracing::Level,

    /// LuaJIT control commands.
    #[clap(long, short = 'j', value_name = "CMD=FLAGS")]
    jit: Vec<String>,

    /// Number of tokio worker threads.
    #[clap(long, value_name = "THREADS", default_value_t = Self::threads())]
    threads: NonZero<usize>,

    /// Number of tokio blocking threads.
    #[clap(long, value_name = "THREADS", default_value_t = Self::blocking_threads())]
    blocking_threads: NonZero<usize>,

    /// Enable tokio-console integration.
    #[clap(long)]
    enable_console: bool,

    /// tokio-console publish address.
    #[clap(
        long,
        value_name = "ADDRESS",
        default_value = "127.0.0.1:6669",
        requires = "enable_console"
    )]
    console_addr: SocketAddr,
}

impl Args {
    fn threads() -> NonZero<usize> {
        thread::available_parallelism().unwrap_or(NonZero::new(1).unwrap())
    }

    fn blocking_threads() -> NonZero<usize> {
        NonZero::new(1024).unwrap()
    }
}

fn exit_err<T, E: Display>(code: ExitCode) -> impl FnOnce(E) -> T {
    move |err| {
        eprintln!("{}", err.red());
        code.exit()
    }
}

fn main() -> Result<(), ExitCode> {
    panic::set_hook(Box::new(panic_cb));

    let args = Args::parse();
    init_logger(&args);

    let tokio = init_tokio(&args);
    let lua = init_lua(&args);
    let main = lua.spawn(async |s| main_async(args, s).await);

    tokio.block_on(async {
        lua.await;
        main.await.unwrap()
    })
}

fn init_logger(args: &Args) {
    use tracing::level_filters::LevelFilter;
    use tracing_subscriber::{Layer, util::*};

    let log = tracing_subscriber::fmt()
        .compact()
        .with_env_filter(
            tracing_subscriber::EnvFilter::builder()
                .with_default_directive(LevelFilter::from(args.log).into())
                .from_env_lossy(),
        )
        .with_file(false)
        .with_line_number(false)
        .with_target(false)
        .finish();

    if args.enable_console {
        console_subscriber::ConsoleLayer::builder()
            .with_default_env()
            .server_addr(args.console_addr)
            .spawn()
            .with_subscriber(log)
            .init()
    } else {
        log.init()
    }
}

fn init_tokio(args: &Args) -> tokio::runtime::Runtime {
    let mut rt = match args.threads.get() {
        1 => tokio::runtime::Builder::new_current_thread(),
        n => {
            let mut rt = tokio::runtime::Builder::new_multi_thread();
            rt.worker_threads(n - 1);
            rt
        }
    };

    rt.enable_all()
        .thread_name("luby")
        .max_blocking_threads(args.blocking_threads.get())
        .build()
        .unwrap_or_else(exit_err(ExitCode::OsErr))
}

fn init_lua(args: &Args) -> lb::runtime::Runtime {
    let rt = lb::runtime::Builder::new();
    let mut rt = rt.build().unwrap_or_else(exit_err(ExitCode::Software));

    for arg in args.jit.iter() {
        let mut s = rt.guard();
        if let Some((cmd, flags)) = parse_jitlib_cmd(arg)
            && let Ok(_) = s.require(format!("jit.{cmd}"), 1)
        {
            (s.push("start"), s.get(-2), s.push(flags));
            s.call(1, 0)
        } else {
            s.require("jit", 1).unwrap();
            match arg.as_str() {
                cmd @ ("on" | "off" | "flush") => {
                    (s.push(cmd), s.get(-2));
                    s.call(0, 0)
                }
                arg => {
                    (s.push("opt"), s.get(-2));
                    (s.push("start"), s.get(-2), s.push(arg));
                    s.call(1, 0)
                }
            }
        }
        .unwrap_or_else(exit_err(ExitCode::Usage));
    }

    rt
}

fn parse_jitlib_cmd(s: &str) -> Option<(&str, &str)> {
    match s {
        "p" => Some(("p", "Flspv10")),
        "v" => Some(("v", "-")),
        "dump" => Some(("dump", "tirs")),
        _ => s.split_once('='),
    }
}

async fn main_async(args: Args, state: &mut luajit::State) -> Result<(), ExitCode> {
    for ref path in args.path {
        let mut s = state.guard();
        let chunk = match std::fs::read(path) {
            Ok(chunk) => chunk,
            Err(err) => {
                eprintln!("{}", format_args!("{path}: {err}").red());
                ExitCode::NoInput.exit();
            }
        };

        s.load(&luajit::Chunk::new(chunk).path(path))
            .unwrap_or_else(exit_err(ExitCode::NoInput));

        if let Err(err) = s.call_async(0, 0).await {
            match err.trace() {
                Some(trace) => eprintln!("{}\n{trace}", err.red()), // runtime error
                None => eprintln!("{}", err.red()),
            }

            ExitCode::DataErr.exit();
        }
    }

    Ok(())
}