rhg: Make configuration available as early as possible in main()

Differential Revision: https://phab.mercurial-scm.org/D10090

rust/rhg/src/blackbox.rs

@@ -1,161 +1,163 b''
 //! Logging for repository events, including commands run in the repository.
 
 use crate::CliInvocation;
 use format_bytes::format_bytes;
 use hg::errors::HgError;
 use hg::repo::Repo;
 use hg::utils::{files::get_bytes_from_os_str, shell_quote};
 
 const ONE_MEBIBYTE: u64 = 1 << 20;
 
 // TODO: somehow keep defaults in sync with `configitem` in `hgext/blackbox.py`
 const DEFAULT_MAX_SIZE: u64 = ONE_MEBIBYTE;
 const DEFAULT_MAX_FILES: u32 = 7;
 
 // Python does not support %.3f, only %f
 const DEFAULT_DATE_FORMAT: &str = "%Y/%m/%d %H:%M:%S%.3f";
 
 type DateTime = chrono::DateTime<chrono::Local>;
 
 pub struct ProcessStartTime {
     /// For measuring duration
     monotonic_clock: std::time::Instant,
     /// For formatting with year, month, day, etc.
     calendar_based: DateTime,
 }
 
 impl ProcessStartTime {
     pub fn now() -> Self {
         Self {
             monotonic_clock: std::time::Instant::now(),
             calendar_based: chrono::Local::now(),
         }
     }
 }
 
 pub struct Blackbox<'a> {
     process_start_time: &'a ProcessStartTime,
     /// Do nothing if this is `None`
     configured: Option<ConfiguredBlackbox<'a>>,
 }
 
 struct ConfiguredBlackbox<'a> {
     repo: &'a Repo,
     max_size: u64,
     max_files: u32,
     date_format: &'a str,
 }
 
 impl<'a> Blackbox<'a> {
     pub fn new(
         invocation: &'a CliInvocation<'a>,
         process_start_time: &'a ProcessStartTime,
     ) -> Result<Self, HgError> {
         let configured = if let Ok(repo) = invocation.repo {
-            let config = invocation.config();
-            if config.get(b"extensions", b"blackbox").is_none() {
+            if invocation.config.get(b"extensions", b"blackbox").is_none() {
                 // The extension is not enabled
                 None
             } else {
                 Some(ConfiguredBlackbox {
                     repo,
-                    max_size: config
+                    max_size: invocation
+                        .config
                         .get_byte_size(b"blackbox", b"maxsize")?
                         .unwrap_or(DEFAULT_MAX_SIZE),
-                    max_files: config
+                    max_files: invocation
+                        .config
                         .get_u32(b"blackbox", b"maxfiles")?
                         .unwrap_or(DEFAULT_MAX_FILES),
-                    date_format: config
+                    date_format: invocation
+                        .config
                         .get_str(b"blackbox", b"date-format")?
                         .unwrap_or(DEFAULT_DATE_FORMAT),
                 })
             }
         } else {
             // Without a local repository there’s no `.hg/blackbox.log` to
             // write to.
             None
         };
         Ok(Self {
             process_start_time,
             configured,
         })
     }
 
     pub fn log_command_start(&self) {
         if let Some(configured) = &self.configured {
             let message = format_bytes!(b"(rust) {}", format_cli_args());
             configured.log(&self.process_start_time.calendar_based, &message);
         }
     }
 
     pub fn log_command_end(&self, exit_code: i32) {
         if let Some(configured) = &self.configured {
             let now = chrono::Local::now();
             let duration = self
                 .process_start_time
                 .monotonic_clock
                 .elapsed()
                 .as_secs_f64();
             let message = format_bytes!(
                 b"(rust) {} exited {} after {} seconds",
                 format_cli_args(),
                 exit_code,
                 format_bytes::Utf8(format_args!("{:.03}", duration))
             );
             configured.log(&now, &message);
         }
     }
 }
 
 impl ConfiguredBlackbox<'_> {
     fn log(&self, date_time: &DateTime, message: &[u8]) {
         let date = format_bytes::Utf8(date_time.format(self.date_format));
         let user = users::get_current_username().map(get_bytes_from_os_str);
         let user = user.as_deref().unwrap_or(b"???");
         let rev = format_bytes::Utf8(match self.repo.dirstate_parents() {
             Ok(parents) if parents.p2 == hg::revlog::node::NULL_NODE => {
                 format!("{:x}", parents.p1)
             }
             Ok(parents) => format!("{:x}+{:x}", parents.p1, parents.p2),
             Err(_dirstate_corruption_error) => {
                 // TODO: log a non-fatal warning to stderr
                 "???".to_owned()
             }
         });
         let pid = std::process::id();
         let line = format_bytes!(
             b"{} {} @{} ({})> {}\n",
             date,
             user,
             rev,
             pid,
             message
         );
         let result =
             hg::logging::LogFile::new(self.repo.hg_vfs(), "blackbox.log")
                 .max_size(Some(self.max_size))
                 .max_files(self.max_files)
                 .write(&line);
         match result {
             Ok(()) => {}
             Err(_io_error) => {
                 // TODO: log a non-fatal warning to stderr
             }
         }
     }
 }
 
 fn format_cli_args() -> Vec<u8> {
     let mut args = std::env::args_os();
     let _ = args.next(); // Skip the first (or zeroth) arg, the name of the `rhg` executable
     let mut args = args.map(|arg| shell_quote(&get_bytes_from_os_str(arg)));
     let mut formatted = Vec::new();
     if let Some(arg) = args.next() {
         formatted.extend(arg)
     }
     for arg in args {
         formatted.push(b' ');
         formatted.extend(arg)
     }
     formatted
 }

rust/rhg/src/commands/config.rs

@@ -1,36 +1,36 b''
 use crate::error::CommandError;
 use clap::Arg;
 use format_bytes::format_bytes;
 use hg::errors::HgError;
 use hg::utils::SliceExt;
 
 pub const HELP_TEXT: &str = "
 With one argument of the form section.name, print just the value of that config item.
 ";
 
 pub fn args() -> clap::App<'static, 'static> {
     clap::SubCommand::with_name("config")
         .arg(
             Arg::with_name("name")
                 .help("the section.name to print")
                 .value_name("NAME")
                 .required(true)
                 .takes_value(true),
         )
         .about(HELP_TEXT)
 }
 
 pub fn run(invocation: &crate::CliInvocation) -> Result<(), CommandError> {
     let (section, name) = invocation
         .subcommand_args
         .value_of("name")
         .expect("missing required CLI argument")
         .as_bytes()
         .split_2(b'.')
         .ok_or_else(|| HgError::abort(""))?;
 
-    let value = invocation.config().get(section, name).unwrap_or(b"");
+    let value = invocation.config.get(section, name).unwrap_or(b"");
 
     invocation.ui.write_stdout(&format_bytes!(b"{}\n", value))?;
     Ok(())
 }

rust/rhg/src/main.rs

@@ -1,179 +1,228 b''
 extern crate log;
 use crate::ui::Ui;
 use clap::App;
 use clap::AppSettings;
 use clap::Arg;
 use clap::ArgMatches;
 use format_bytes::format_bytes;
 use hg::config::Config;
 use hg::repo::{Repo, RepoError};
-use std::path::{Path, PathBuf};
+use hg::utils::files::{get_bytes_from_os_str, get_path_from_bytes};
+use hg::utils::SliceExt;
+use std::ffi::OsString;
+use std::path::PathBuf;
 
 mod blackbox;
 mod error;
 mod exitcode;
 mod ui;
 use error::CommandError;
 
 fn main_with_result(
+    process_start_time: &blackbox::ProcessStartTime,
     ui: &ui::Ui,
-    process_start_time: &blackbox::ProcessStartTime,
+    repo: Result<&Repo, &NoRepoInCwdError>,
+    config: &Config,
 ) -> Result<(), CommandError> {
-    env_logger::init();
     let app = App::new("rhg")
         .global_setting(AppSettings::AllowInvalidUtf8)
         .setting(AppSettings::SubcommandRequired)
         .setting(AppSettings::VersionlessSubcommands)
         .arg(
             Arg::with_name("repository")
                 .help("repository root directory")
                 .short("-R")
                 .long("--repository")
                 .value_name("REPO")
                 .takes_value(true)
                 // Both ok: `hg -R ./foo log` or `hg log -R ./foo`
                 .global(true),
         )
         .arg(
             Arg::with_name("config")
                 .help("set/override config option (use 'section.name=value')")
                 .long("--config")
                 .value_name("CONFIG")
                 .takes_value(true)
                 .global(true)
                 // Ok: `--config section.key1=val --config section.key2=val2`
                 .multiple(true)
                 // Not ok: `--config section.key1=val section.key2=val2`
                 .number_of_values(1),
         )
         .version("0.0.1");
     let app = add_subcommand_args(app);
 
     let matches = app.clone().get_matches_safe()?;
 
     let (subcommand_name, subcommand_matches) = matches.subcommand();
     let run = subcommand_run_fn(subcommand_name)
         .expect("unknown subcommand name from clap despite AppSettings::SubcommandRequired");
     let subcommand_args = subcommand_matches
         .expect("no subcommand arguments from clap despite AppSettings::SubcommandRequired");
 
-    let config_args = matches
-        .values_of_os("config")
-        // Turn `Option::None` into an empty iterator:
-        .into_iter()
-        .flatten()
-        .map(hg::utils::files::get_bytes_from_os_str);
-    let non_repo_config = &hg::config::Config::load(config_args)?;
-
-    let repo_path = matches.value_of_os("repository").map(Path::new);
-    let repo = match Repo::find(non_repo_config, repo_path) {
-        Ok(repo) => Ok(repo),
-        Err(RepoError::NotFound { at }) if repo_path.is_none() => {
-            // Not finding a repo is not fatal yet, if `-R` was not given
-            Err(NoRepoInCwdError { cwd: at })
-        }
-        Err(error) => return Err(error.into()),
-    };
-
     let invocation = CliInvocation {
         ui,
         subcommand_args,
-        non_repo_config,
-        repo: repo.as_ref(),
+        config,
+        repo,
     };
     let blackbox = blackbox::Blackbox::new(&invocation, process_start_time)?;
     blackbox.log_command_start();
     let result = run(&invocation);
     blackbox.log_command_end(exit_code(&result));
     result
 }
 
 fn main() {
     // Run this first, before we find out if the blackbox extension is even
     // enabled, in order to include everything in-between in the duration
     // measurements. Reading config files can be slow if they’re on NFS.
     let process_start_time = blackbox::ProcessStartTime::now();
 
+    env_logger::init();
     let ui = ui::Ui::new();
 
-    let result = main_with_result(&ui, &process_start_time);
-    if let Err(CommandError::Abort { message }) = &result {
-        if !message.is_empty() {
-            // Ignore errors when writing to stderr, we’re already exiting
-            // with failure code so there’s not much more we can do.
-            let _ = ui.write_stderr(&format_bytes!(b"abort: {}\n", message));
+    let early_args = EarlyArgs::parse(std::env::args_os());
+    let non_repo_config = Config::load(early_args.config)
+        .unwrap_or_else(|error| exit(&ui, Err(error.into())));
+
+    let repo_path = early_args.repo.as_deref().map(get_path_from_bytes);
+    let repo_result = match Repo::find(&non_repo_config, repo_path) {
+        Ok(repo) => Ok(repo),
+        Err(RepoError::NotFound { at }) if repo_path.is_none() => {
+            // Not finding a repo is not fatal yet, if `-R` was not given
+            Err(NoRepoInCwdError { cwd: at })
         }
-    }
-    std::process::exit(exit_code(&result))
+        Err(error) => exit(&ui, Err(error.into())),
+    };
+
+    let config = if let Ok(repo) = &repo_result {
+        repo.config()
+    } else {
+        &non_repo_config
+    };
+
+    let result = main_with_result(
+        &process_start_time,
+        &ui,
+        repo_result.as_ref(),
+        config,
+    );
+    exit(&ui, result)
 }
 
 fn exit_code(result: &Result<(), CommandError>) -> i32 {
     match result {
         Ok(()) => exitcode::OK,
         Err(CommandError::Abort { .. }) => exitcode::ABORT,
 
         // Exit with a specific code and no error message to let a potential
         // wrapper script fallback to Python-based Mercurial.
         Err(CommandError::Unimplemented) => exitcode::UNIMPLEMENTED,
     }
 }
 
+fn exit(ui: &Ui, result: Result<(), CommandError>) -> ! {
+    if let Err(CommandError::Abort { message }) = &result {
+        if !message.is_empty() {
+            // Ignore errors when writing to stderr, we’re already exiting
+            // with failure code so there’s not much more we can do.
+            let _ = ui.write_stderr(&format_bytes!(b"abort: {}\n", message));
+        }
+    }
+    std::process::exit(exit_code(&result))
+}
+
 macro_rules! subcommands {
     ($( $command: ident )+) => {
         mod commands {
             $(
                 pub mod $command;
             )+
         }
 
         fn add_subcommand_args<'a, 'b>(app: App<'a, 'b>) -> App<'a, 'b> {
             app
             $(
                 .subcommand(commands::$command::args())
             )+
         }
 
         pub type RunFn = fn(&CliInvocation) -> Result<(), CommandError>;
 
         fn subcommand_run_fn(name: &str) -> Option<RunFn> {
             match name {
                 $(
                     stringify!($command) => Some(commands::$command::run),
                 )+
                 _ => None,
             }
         }
     };
 }
 
 subcommands! {
     cat
     debugdata
     debugrequirements
     files
     root
     config
 }
 pub struct CliInvocation<'a> {
     ui: &'a Ui,
     subcommand_args: &'a ArgMatches<'a>,
-    non_repo_config: &'a Config,
+    config: &'a Config,
     /// References inside `Result` is a bit peculiar but allow
     /// `invocation.repo?` to work out with `&CliInvocation` since this
     /// `Result` type is `Copy`.
     repo: Result<&'a Repo, &'a NoRepoInCwdError>,
 }
 
 struct NoRepoInCwdError {
     cwd: PathBuf,
 }
 
-impl CliInvocation<'_> {
-    fn config(&self) -> &Config {
-        if let Ok(repo) = self.repo {
-            repo.config()
-        } else {
-            self.non_repo_config
+/// CLI arguments to be parsed "early" in order to be able to read
+/// configuration before using Clap. Ideally we would also use Clap for this,
+/// see <https://github.com/clap-rs/clap/discussions/2366>.
+///
+/// These arguments are still declared when we do use Clap later, so that Clap
+/// does not return an error for their presence.
+struct EarlyArgs {
+    /// Values of all `--config` arguments. (Possibly none)
+    config: Vec<Vec<u8>>,
+    /// Value of the `-R` or `--repository` argument, if any.
+    repo: Option<Vec<u8>>,
+}
+
+impl EarlyArgs {
+    fn parse(args: impl IntoIterator<Item = OsString>) -> Self {
+        let mut args = args.into_iter().map(get_bytes_from_os_str);
+        let mut config = Vec::new();
+        let mut repo = None;
+        // Use `while let` instead of `for` so that we can also call
+        // `args.next()` inside the loop.
+        while let Some(arg) = args.next() {
+            if arg == b"--config" {
+                if let Some(value) = args.next() {
+                    config.push(value)
+                }
+            } else if let Some(value) = arg.drop_prefix(b"--config=") {
+                config.push(value.to_owned())
+            }
+
+            if arg == b"--repository" || arg == b"-R" {
+                if let Some(value) = args.next() {
+                    repo = Some(value)
+                }
+            } else if let Some(value) = arg.drop_prefix(b"--repository=") {
+                repo = Some(value.to_owned())
+            } else if let Some(value) = arg.drop_prefix(b"-R") {
+                repo = Some(value.to_owned())
+            }
         }
+        Self { config, repo }
     }
 }