This is an archive of the discontinued Mercurial Phabricator instance.

Differential D8447

rust-chg: reimplement locator by using async/await and tokio-0.2
ClosedPublic

Authored by yuja on Apr 16 2020, 8:30 AM.

Details

Reviewers
Alphare
Group Reviewers
hg-reviewers
Commits
rHGa347a329e48d: rust-chg: reimplement locator by using async/await and tokio-0.2
Summary

connect_spawned() is rewritten from scratch by using std::process. Before,
it would select completion of either connection or server process. New code
could be implemented as such, but it's much simpler to occasionally run
try_wait() to detect server death.

Diff Detail

Repository
rHG Mercurial
Branch
default
Lint
No Linters Available
Unit
No Unit Test Coverage

Event Timeline

yuja created this revision.Apr 16 2020, 8:30 AM
Alphare accepted this revision.Apr 17 2020, 8:30 AM
yuja added a commit: rHGa347a329e48d: rust-chg: reimplement locator by using async/await and tokio-0.2.Apr 23 2020, 2:03 PM
This revision was not accepted when it landed; it landed in state Needs Review.
Closed by commit rHGa347a329e48d: rust-chg: reimplement locator by using async/await and tokio-0.2 (authored by yuja). · Explain Why
This revision was automatically updated to reflect the committed changes.

Revision Contents
Changeset List

PathPackages
Mrust/chg/src/lib.rs (2 lines)
Mrust/chg/src/locator.rs (202 lines)

Diff 21124

rust/chg/src/lib.rs

// Copyright 2018 Yuya Nishihara <yuya@tcha.org> // Copyright 2018 Yuya Nishihara <yuya@tcha.org>
// //
// This software may be used and distributed according to the terms of the // This software may be used and distributed according to the terms of the
// GNU General Public License version 2 or any later version. // GNU General Public License version 2 or any later version.
mod attachio; mod attachio;
mod clientext; mod clientext;
//pub mod locator; pub mod locator;
pub mod message; pub mod message;
pub mod procutil; pub mod procutil;
mod runcommand; mod runcommand;
mod uihandler; mod uihandler;
pub use clientext::ChgClient; pub use clientext::ChgClient;
pub use uihandler::{ChgUiHandler, SystemHandler}; pub use uihandler::{ChgUiHandler, SystemHandler};

rust/chg/src/locator.rs

// Copyright 2011, 2018 Yuya Nishihara <yuya@tcha.org> // Copyright 2011, 2018 Yuya Nishihara <yuya@tcha.org>
// //
// This software may be used and distributed according to the terms of the // This software may be used and distributed according to the terms of the
// GNU General Public License version 2 or any later version. // GNU General Public License version 2 or any later version.
//! Utility for locating command-server process. //! Utility for locating command-server process.
use futures::future::{self, Either, Loop};
use log::debug; use log::debug;
use std::env; use std::env;
use std::ffi::{OsStr, OsString}; use std::ffi::{OsStr, OsString};
use std::fs::{self, DirBuilder}; use std::fs::{self, DirBuilder};
use std::io; use std::io;
use std::os::unix::ffi::{OsStrExt, OsStringExt}; use std::os::unix::ffi::{OsStrExt, OsStringExt};
use std::os::unix::fs::{DirBuilderExt, MetadataExt}; use std::os::unix::fs::{DirBuilderExt, MetadataExt};
use std::path::{Path, PathBuf}; use std::path::{Path, PathBuf};
use std::process; use std::process::{self, Child, Command};
use std::time::Duration; use std::time::{Duration, Instant};
use tokio::prelude::*;
use tokio::process::{Child, Command};
use tokio::time; use tokio::time;
use tokio_hglib::UnixClient;
use crate::clientext::ChgClientExt; use crate::clientext::ChgClient;
use crate::message::{Instruction, ServerSpec}; use crate::message::{Instruction, ServerSpec};
use crate::procutil; use crate::procutil;
const REQUIRED_SERVER_CAPABILITIES: &[&str] = &[ const REQUIRED_SERVER_CAPABILITIES: &[&str] = &[
"attachio", "attachio",
"chdir", "chdir",
"runcommand", "runcommand",
"setenv", "setenv",
/// Specifies the arguments to be passed to the server at start. /// Specifies the arguments to be passed to the server at start.
pub fn set_early_args(&mut self, args: impl IntoIterator<Item = impl AsRef<OsStr>>) { pub fn set_early_args(&mut self, args: impl IntoIterator<Item = impl AsRef<OsStr>>) {
self.hg_early_args = args.into_iter().map(|a| a.as_ref().to_owned()).collect(); self.hg_early_args = args.into_iter().map(|a| a.as_ref().to_owned()).collect();
} }
/// Connects to the server. /// Connects to the server.
/// ///
/// The server process will be spawned if not running. /// The server process will be spawned if not running.
pub fn connect(self) -> impl Future<Item = (Self, UnixClient), Error = io::Error> { pub async fn connect(&mut self) -> io::Result<ChgClient> {
future::loop_fn((self, 0), |(loc, cnt)| { for _cnt in 0..10 {
if cnt < 10 { let mut client = self.try_connect().await?;
let fut = loc let instructions = client.validate(&self.hg_early_args).await?;
.try_connect() let reconnect = self.run_instructions(&instructions)?;
.and_then(|(loc, client)| { if !reconnect {
client return Ok(client);
.validate(&loc.hg_early_args) }
.map(|(client, instructions)| (loc, client, instructions)) }
})
.and_then(move |(loc, client, instructions)| { // TODO: unindent
loc.run_instructions(client, instructions, cnt) {
}); {
Either::A(fut)
} else {
let msg = format!( let msg = format!(
concat!( concat!(
"too many redirections.\n", "too many redirections.\n",
"Please make sure {:?} is not a wrapper which ", "Please make sure {:?} is not a wrapper which ",
"changes sensitive environment variables ", "changes sensitive environment variables ",
"before executing hg. If you have to use a ", "before executing hg. If you have to use a ",
"wrapper, wrap chg instead of hg.", "wrapper, wrap chg instead of hg.",
), ),
loc.hg_command self.hg_command
); );
Either::B(future::err(io::Error::new(io::ErrorKind::Other, msg))) Err(io::Error::new(io::ErrorKind::Other, msg))
}
} }
})
} }
/// Runs instructions received from the server. /// Runs instructions received from the server.
fn run_instructions( ///
mut self, /// Returns true if the client should try connecting to the other server.
client: UnixClient, fn run_instructions(&mut self, instructions: &[Instruction]) -> io::Result<bool> {
instructions: Vec<Instruction>,
cnt: usize,
) -> io::Result<Loop<(Self, UnixClient), (Self, usize)>> {
let mut reconnect = false; let mut reconnect = false;
for inst in instructions { for inst in instructions {
debug!("instruction: {:?}", inst); debug!("instruction: {:?}", inst);
match inst { match inst {
Instruction::Exit(_) => { Instruction::Exit(_) => {
// Just returns the current connection to run the // Just returns the current connection to run the
// unparsable command and report the error // unparsable command and report the error
return Ok(Loop::Break((self, client))); return Ok(false);
} }
Instruction::Reconnect => { Instruction::Reconnect => {
reconnect = true; reconnect = true;
} }
Instruction::Redirect(path) => { Instruction::Redirect(path) => {
if path.parent() != self.base_sock_path.parent() { if path.parent() != self.base_sock_path.parent() {
let msg = format!( let msg = format!(
"insecure redirect instruction from server: {}", "insecure redirect instruction from server: {}",
path.display() path.display()
); );
return Err(io::Error::new(io::ErrorKind::InvalidData, msg)); return Err(io::Error::new(io::ErrorKind::InvalidData, msg));
} }
self.redirect_sock_path = Some(path); self.redirect_sock_path = Some(path.to_owned());
reconnect = true; reconnect = true;
} }
Instruction::Unlink(path) => { Instruction::Unlink(path) => {
if path.parent() != self.base_sock_path.parent() { if path.parent() != self.base_sock_path.parent() {
let msg = format!( let msg = format!(
"insecure unlink instruction from server: {}", "insecure unlink instruction from server: {}",
path.display() path.display()
); );
return Err(io::Error::new(io::ErrorKind::InvalidData, msg)); return Err(io::Error::new(io::ErrorKind::InvalidData, msg));
} }
fs::remove_file(path).unwrap_or(()); // may race fs::remove_file(path).unwrap_or(()); // may race
} }
} }
} }
if reconnect { Ok(reconnect)
Ok(Loop::Continue((self, cnt + 1)))
} else {
Ok(Loop::Break((self, client)))
}
} }
/// Tries to connect to the existing server, or spawns new if not running. /// Tries to connect to the existing server, or spawns new if not running.
fn try_connect(self) -> impl Future<Item = (Self, UnixClient), Error = io::Error> { async fn try_connect(&mut self) -> io::Result<ChgClient> {
let sock_path = self let sock_path = self
.redirect_sock_path .redirect_sock_path
.as_ref() .as_ref()
.unwrap_or(&self.base_sock_path) .unwrap_or(&self.base_sock_path)
.clone(); .clone();
debug!("try connect to {}", sock_path.display()); debug!("try connect to {}", sock_path.display());
UnixClient::connect(sock_path) // TODO: unindent
.then(|res| { {
match res { {
Ok(client) => Either::A(future::ok((self, client))), let mut client = match ChgClient::connect(sock_path).await {
Ok(client) => client,
Err(_) => { Err(_) => {
// Prevent us from being re-connected to the outdated // Prevent us from being re-connected to the outdated
// master server: We were told by the server to redirect // master server: We were told by the server to redirect
// to redirect_sock_path, which didn't work. We do not // to redirect_sock_path, which didn't work. We do not
// want to connect to the same master server again // want to connect to the same master server again
// because it would probably tell us the same thing. // because it would probably tell us the same thing.
if self.redirect_sock_path.is_some() { if self.redirect_sock_path.is_some() {
fs::remove_file(&self.base_sock_path).unwrap_or(()); fs::remove_file(&self.base_sock_path).unwrap_or(());
// may race // may race
} }
Either::B(self.spawn_connect()) self.spawn_connect().await?
} }
} };
})
.and_then(|(loc, client)| {
check_server_capabilities(client.server_spec())?; check_server_capabilities(client.server_spec())?;
Ok((loc, client))
})
.and_then(|(loc, client)| {
// It's purely optional, and the server might not support this command. // It's purely optional, and the server might not support this command.
if client.server_spec().capabilities.contains("setprocname") { if client.server_spec().capabilities.contains("setprocname") {
let fut = client
.set_process_name(format!("chg[worker/{}]", loc.process_id))
.map(|client| (loc, client));
Either::A(fut)
} else {
Either::B(future::ok((loc, client)))
}
})
.and_then(|(loc, client)| {
client client
.set_current_dir(&loc.current_dir) .set_process_name(format!("chg[worker/{}]", self.process_id))
.map(|client| (loc, client)) .await?;
}) }
.and_then(|(loc, client)| { client.set_current_dir(&self.current_dir).await?;
client client
.set_env_vars_os(loc.env_vars.iter().cloned()) .set_env_vars_os(self.env_vars.iter().cloned())
.map(|client| (loc, client)) .await?;
}) Ok(client)
}
}
} }
/// Spawns new server process and connects to it. /// Spawns new server process and connects to it.
/// ///
/// The server will be spawned at the current working directory, then /// The server will be spawned at the current working directory, then
/// chdir to "/", so that the server will load configs from the target /// chdir to "/", so that the server will load configs from the target
/// repository. /// repository.
fn spawn_connect(self) -> impl Future<Item = (Self, UnixClient), Error = io::Error> { async fn spawn_connect(&mut self) -> io::Result<ChgClient> {
let sock_path = self.temp_sock_path(); let sock_path = self.temp_sock_path();
debug!("start cmdserver at {}", sock_path.display()); debug!("start cmdserver at {}", sock_path.display());
Command::new(&self.hg_command) let server = Command::new(&self.hg_command)
.arg("serve") .arg("serve")
.arg("--cmdserver") .arg("--cmdserver")
.arg("chgunix") .arg("chgunix")
.arg("--address") .arg("--address")
.arg(&sock_path) .arg(&sock_path)
.arg("--daemon-postexec") .arg("--daemon-postexec")
.arg("chdir:/") .arg("chdir:/")
.args(&self.hg_early_args) .args(&self.hg_early_args)
.current_dir(&self.current_dir) .current_dir(&self.current_dir)
.env_clear() .env_clear()
.envs(self.env_vars.iter().cloned()) .envs(self.env_vars.iter().cloned())
.env("CHGINTERNALMARK", "") .env("CHGINTERNALMARK", "")
.spawn() .spawn()?;
.into_future() let client = self.connect_spawned(server, &sock_path).await?;
.and_then(|server| self.connect_spawned(server, sock_path)) // TODO: unindent
.and_then(|(loc, client, sock_path)| { {
{
debug!( debug!(
"rename {} to {}", "rename {} to {}",
sock_path.display(), sock_path.display(),
loc.base_sock_path.display() self.base_sock_path.display()
); );
fs::rename(&sock_path, &loc.base_sock_path)?; fs::rename(&sock_path, &self.base_sock_path)?;
Ok((loc, client)) Ok(client)
}) }
}
} }
/// Tries to connect to the just spawned server repeatedly until timeout /// Tries to connect to the just spawned server repeatedly until timeout
/// exceeded. /// exceeded.
fn connect_spawned( async fn connect_spawned(
self, &mut self,
server: Child, mut server: Child,
sock_path: PathBuf, sock_path: &Path,
) -> impl Future<Item = (Self, UnixClient, PathBuf), Error = io::Error> { ) -> io::Result<ChgClient> {
debug!("try connect to {} repeatedly", sock_path.display()); debug!("try connect to {} repeatedly", sock_path.display());
let connect = future::loop_fn(sock_path, |sock_path| {
UnixClient::connect(sock_path.clone()).then(|res| {
match res {
Ok(client) => Either::A(future::ok(Loop::Break((client, sock_path)))),
Err(_) => {
// try again with slight delay
let fut = time::delay_for(Duration::from_millis(10))
.map(|()| Loop::Continue(sock_path))
.map_err(|err| io::Error::new(io::ErrorKind::Other, err));
Either::B(fut)
}
}
})
});
// waits for either connection established or server failed to start // waits for either connection established or server failed to start
connect let start_time = Instant::now();
.select2(server) while start_time.elapsed() < self.timeout {
.map_err(|res| res.split().0) if let Ok(client) = ChgClient::connect(&sock_path).await {
.timeout(self.timeout) // server handle is dropped here, but the detached process
.map_err(|err| { // will continue running in background
err.into_inner().unwrap_or_else(|| { return Ok(client);
io::Error::new(
io::ErrorKind::TimedOut,
"timed out while connecting to server",
)
})
})
.and_then(|res| {
match res {
Either::A(((client, sock_path), server)) => {
server.forget(); // continue to run in background
Ok((self, client, sock_path))
} }
Either::B((st, _)) => Err(io::Error::new(
if let Some(st) = server.try_wait()? {
return Err(io::Error::new(
io::ErrorKind::Other, io::ErrorKind::Other,
format!("server exited too early: {}", st), format!("server exited too early: {}", st),
)), ));
} }
})
// try again with slight delay
time::delay_for(Duration::from_millis(10)).await;
}
Err(io::Error::new(
io::ErrorKind::TimedOut,
"timed out while connecting to server",
))
} }
} }
/// Determines the server socket to connect to. /// Determines the server socket to connect to.
/// ///
/// If no `$CHGSOCKNAME` is specified, the socket directory will be created /// If no `$CHGSOCKNAME` is specified, the socket directory will be created
/// as necessary. /// as necessary.
fn prepare_server_socket_path() -> io::Result<PathBuf> { fn prepare_server_socket_path() -> io::Result<PathBuf> {