luby/src/main.rs

use clap::Parser;
use luajit::Chunk;
use mimalloc::MiMalloc;
use owo_colors::OwoColorize;
use std::{backtrace::Backtrace, fmt::Display, num::NonZero, panic, process, 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()
        .bold()
    );

    eprintln!(
        "{}",
        format_args!(
            "luby should never panic. Please kindly report this bug at {}.",
            env!("CARGO_PKG_REPOSITORY")
        )
        .yellow()
        .bold()
    );
}

fn error_cb(err: &luajit::Error) {
    match err.trace() {
        Some(trace) => eprintln!("{}\n{trace}", err.red().bold()),
        None => eprintln!("{}", err.red().bold()),
    }

    process::exit(1);
}

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

#[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.
    #[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', help_heading = "Runtime", value_name = "CMD=FLAGS")]
    jit: Vec<String>,

    /// Allow global variables.
    #[clap(
        long,
        help_heading = "Runtime",
        value_name = "ENABLED",
        default_value_t = true
    )]
    allow_globals: bool,

    /// Number of worker threads.
    #[clap(
        long,
        short = 'T',
        help_heading = "Runtime",
        value_name = "COUNT",
        default_value_t = Self::threads()
    )]
    threads: NonZero<usize>,

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

    /// Enable tokio-console integration.
    #[cfg(feature = "tokio-console")]
    #[clap(long, help_heading = "Debugging", value_name = "ENABLED")]
    enable_console: bool,

    /// tokio-console publish address.
    #[cfg(feature = "tokio-console")]
    #[clap(
        long,
        help_heading = "Debugging",
        value_name = "ADDRESS",
        default_value = "127.0.0.1:6669",
        requires = "enable_console"
    )]
    console_addr: std::net::SocketAddr,

    /// Dump internal data.
    #[clap(
        long,
        help_heading = "Debugging",
        value_name = "DATA",
        value_parser = ["cdef"]
    )]
    dump: Vec<String>,

    /// Print version.
    #[clap(long, short = 'V')]
    version: bool,
}

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 main() -> ExitCode {
    panic::set_hook(Box::new(panic_cb));

    let args = Args::parse();
    if args.version {
        print_version();
        return ExitCode::Ok;
    }

    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;
        match main.await {
            Ok(res) => res,
            Err(err) => panic::resume_unwind(err.into_panic()),
        }
    })
}

fn print_version() {
    println!("luby {}", env!("VERGEN_GIT_DESCRIBE"));
    println!("{}\n", env!("CARGO_PKG_HOMEPAGE"));
    println!("Compiled with {} -- {}", luajit::version(), luajit::url());
    println!(
        "Compiled with rustc {} on {} for {}",
        env!("VERGEN_RUSTC_SEMVER"),
        env!("VERGEN_RUSTC_HOST_TRIPLE"),
        env!("VERGEN_CARGO_TARGET_TRIPLE"),
    );
}

fn init_logger(args: &Args) {
    use tracing::level_filters::LevelFilter;
    use tracing_subscriber::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();

    #[cfg(feature = "tokio-console")]
    {
        use tracing_subscriber::Layer;
        console_subscriber::ConsoleLayer::builder()
            .with_default_env()
            .server_addr(args.console_addr)
            .spawn()
            .with_subscriber(log)
            .init();
    }

    #[cfg(not(feature = "tokio-console"))]
    log.init();
}

fn init_tokio(args: &Args) -> tokio::runtime::Runtime {
    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
        }
    }
    .enable_all()
    .thread_name("luby")
    .max_blocking_threads(args.blocking_threads.get())
    .build()
    .unwrap_or_else(unwrap_exit(ExitCode::OsErr))
}

fn init_lua(args: &Args) -> lb::runtime::Runtime {
    let mut rt = lb::runtime::Builder::new();
    luby::open(&mut rt);

    if args.dump.iter().find(|s| *s == "cdef").is_some() {
        print!("{}", rt.registry()); // for cdef debugging
    }

    rt.unhandled_error(error_cb)
        .prohibit_globals(!args.allow_globals)
        .jit_opts(args.jit.iter())
        .build()
        .unwrap()
}

async fn main_async(args: Args, cx: &mut lb::runtime::Context) -> ExitCode {
    for ref path in args.path {
        let chunk = match std::fs::read(path) {
            Ok(chunk) => chunk,
            Err(err) => {
                eprintln!("{}", format_args!("{path}: {err}").red().bold());
                return ExitCode::NoInput;
            }
        };

        if let Err(ref err) = cx.load(&Chunk::new(chunk).with_path(path)) {
            cx.report_error(err);
        } else if let Err(ref err) = cx.call_async(0, Some(0)).await {
            cx.report_error(err);
        }
    }

    ExitCode::Ok
}