From ad3baf259b20a984d16a6149864f8b2f338c87b2 Mon Sep 17 00:00:00 2001
From: Jun Wu <quark@meta.com>
Date: Mon, 3 Jul 2023 22:36:01 -0700
Subject: [PATCH] hgcommands: integrate with commandserver server

Summary:
Make `start-commandserver` actually starts a server, similar to chgserver.
Unlike chgserver, this server is implemented in Rust and only serves one
client (for now).

Reviewed By: zzl0

Differential Revision: D47071254

fbshipit-source-id: 259ef08f2f72f4c22311e0b3bc9dbbee43fd5298
---
 eden/scm/lib/hgcommands/Cargo.toml      |  1 +
 eden/scm/lib/hgcommands/src/hgpython.rs | 11 ++++++
 eden/scm/lib/hgcommands/src/run.rs      | 49 +++++++++++++++++++++++--
 3 files changed, 58 insertions(+), 3 deletions(-)

diff --git a/eden/scm/lib/hgcommands/Cargo.toml b/eden/scm/lib/hgcommands/Cargo.toml
index b570519d1d..366f49e47e 100644
--- a/eden/scm/lib/hgcommands/Cargo.toml
+++ b/eden/scm/lib/hgcommands/Cargo.toml
@@ -17,6 +17,7 @@ clidispatch = { version = "0.1.0", path = "../clidispatch" }
 cliparser = { version = "0.1.0", path = "../cliparser", features = ["python"] }
 clone = { version = "0.1.0", path = "../clone" }
 comfy-table = "6.1.4"
+commandserver = { version = "0.1.0", path = "../commandserver" }
 configloader = { version = "0.1.0", path = "../config/loader" }
 configmodel = { version = "0.1.0", path = "../config/model" }
 cpython = { version = "0.7.1", default-features = false }
diff --git a/eden/scm/lib/hgcommands/src/hgpython.rs b/eden/scm/lib/hgcommands/src/hgpython.rs
index 37caebbc0d..9a53d241b0 100644
--- a/eden/scm/lib/hgcommands/src/hgpython.rs
+++ b/eden/scm/lib/hgcommands/src/hgpython.rs
@@ -229,6 +229,17 @@ impl HgPython {
             1
         }
     }
+
+    /// Pre-import Python modules.
+    /// Returns after importing the modules.
+    pub fn pre_import_modules(&self) -> Result<(), cpython_ext::PyErr> {
+        // cpython_ext::PyErr can render traceback when RUST_BACKTRACE=1.
+        let gil = Python::acquire_gil();
+        let py = gil.python();
+        let dispatch = py.import("edenscm.dispatch")?;
+        dispatch.call(py, "_preimportmodules", NoArgs, None)?;
+        Ok(())
+    }
 }
 
 impl Drop for HgPython {
diff --git a/eden/scm/lib/hgcommands/src/run.rs b/eden/scm/lib/hgcommands/src/run.rs
index 1932b88dc7..9b5548a82b 100644
--- a/eden/scm/lib/hgcommands/src/run.rs
+++ b/eden/scm/lib/hgcommands/src/run.rs
@@ -15,6 +15,7 @@ use std::ops::DerefMut;
 use std::path::PathBuf;
 use std::str::FromStr;
 use std::sync::atomic::AtomicBool;
+use std::sync::atomic::Ordering;
 use std::sync::atomic::Ordering::SeqCst;
 use std::sync::Arc;
 use std::sync::Weak;
@@ -59,10 +60,17 @@ use crate::HgPython;
 pub fn run_command(args: Vec<String>, io: &IO) -> i32 {
     let start_time = SystemTime::now();
 
-    // The pfcserver does not want tracing or blackbox or ctrlc setup,
+    // The pfcserver or commandserver do not want tracing or blackbox or ctrlc setup,
     // or going through the Rust command table. Bypass them.
-    if args.get(1).map(|s| s.as_ref()) == Some("start-pfc-server") {
-        return HgPython::new(&args).run_hg(args, io, &ConfigSet::new());
+    if let Some(arg1) = args.get(1).map(|s| s.as_ref()) {
+        match arg1 {
+            "start-pfc-server" => return HgPython::new(&args).run_hg(args, io, &ConfigSet::new()),
+            "start-commandserver" => {
+                commandserver_serve(&args, io);
+                return 0;
+            }
+            _ => {}
+        }
     }
 
     // Initialize NodeIpc:
@@ -887,3 +895,38 @@ fn setup_nodeipc() {
     // Trigger `Lazy` initialization.
     let _ = nodeipc::get_singleton();
 }
+
+// Useful to prevent a commandserver connecting to another commandserver.
+static IS_COMMANDSERVER: AtomicBool = AtomicBool::new(false);
+
+fn commandserver_serve(args: &[String], io: &IO) -> i32 {
+    IS_COMMANDSERVER.store(true, Ordering::Release);
+
+    #[cfg(unix)]
+    unsafe {
+        libc::setsid();
+    }
+
+    let _ = setup_tracing_io(io, None);
+    tracing::debug!("preparing commandserver");
+
+    let python = HgPython::new(args);
+    if let Err(e) = python.pre_import_modules() {
+        tracing::warn!("cannot pre-import modules:\n{:?}", &e);
+        return 1;
+    }
+
+    let run_func = |args: Vec<String>| -> i32 {
+        tracing::debug!("commandserver is about to run command: {:?}", &args);
+        run_command(args, io)
+    };
+
+    // TODO(quark): We need to route `ui.system` through IPC.
+    tracing::debug!("commandserver is about to serve");
+    if let Err(e) = commandserver::server::serve_one_client(&run_func) {
+        tracing::warn!("cannot serve:\n{:?}", &e);
+        return 1;
+    }
+    tracing::debug!("commandserver is about to exit cleanly");
+    0
+}