refactor: replace ctrlc handler with tokio's built-in (#197)

This commit is contained in:
Daniel Gushchyan 2024-09-18 08:26:43 -07:00 committed by GitHub
parent 43a1a49c3f
commit 628d7b4596
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 29 additions and 111 deletions

11
Cargo.lock generated
View file

@ -1135,16 +1135,6 @@ dependencies = [
"cipher",
]
[[package]]
name = "ctrlc"
version = "3.4.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "90eeab0aa92f3f9b4e87f258c72b139c207d251f9cbc1080a0086b86a8870dd3"
dependencies = [
"nix 0.29.0",
"windows-sys 0.59.0",
]
[[package]]
name = "curve25519-dalek"
version = "4.1.3"
@ -4090,7 +4080,6 @@ dependencies = [
"color-eyre",
"cosmrs",
"cosmwasm-std",
"ctrlc",
"dirs 5.0.1",
"displaydoc",
"figment",

View file

@ -38,7 +38,6 @@ watchexec = "4.1.0"
watchexec-events = "3.0.0"
watchexec-signals = "3.0.0"
miette = "7.2.0"
ctrlc = { version = "3.4.5", features=["termination"]}
xxhash-rust = { version = "0.8.12", features=["xxh3"] }
toml = "0.8.19"
figment = { version = "0.10.19", features=["env", "toml"] }

View file

@ -1,14 +1,11 @@
use std::{path::PathBuf, process::exit, time::Duration};
use std::{path::PathBuf, time::Duration};
use async_trait::async_trait;
use color_eyre::owo_colors::OwoColorize;
// todo get rid of this?
use miette::{IntoDiagnostic, Result};
use quartz_common::proto::core_client::CoreClient;
use tokio::{
sync::{mpsc, watch},
time::sleep,
};
use tokio::{sync::mpsc, time::sleep};
use tracing::{debug, info};
use watchexec::Watchexec;
use watchexec_signals::Signal;
@ -63,21 +60,11 @@ async fn dev_driver(
config: Config,
) -> Result<(), Error> {
// State
let mut shutdown_tx: Option<watch::Sender<()>> = None;
let mut first_enclave_message = true;
let mut first_contract_message = true;
let mut contract = String::from("");
// Shutdown enclave upon interruption
let shutdown_tx_cpy = shutdown_tx.clone();
ctrlc::set_handler(move || {
if let Some(tx) = &shutdown_tx_cpy {
let _res = tx.send(());
}
exit(130)
})
.expect("Error setting Ctrl-C handler");
// Drive
while let Some(dev) = rx.recv().await {
@ -98,7 +85,7 @@ async fn dev_driver(
contract_build.handle(&config).await?;
// Start enclave in background
let new_shutdown_tx = spawn_enclave_start(args, &config).await?;
spawn_enclave_start(args, &config)?;
// Deploy new contract and perform handshake
let res = deploy_and_handshake(None, args, &config).await;
@ -108,17 +95,12 @@ async fn dev_driver(
Ok(res_contract) => {
// Set state
contract = res_contract;
shutdown_tx = Some(new_shutdown_tx);
info!("{}", "Enclave is listening for requests...".green().bold());
}
Err(e) => {
eprintln!("Error launching quartz app");
new_shutdown_tx
.send(())
.expect("Could not send signal on channel");
return Err(e);
}
}
@ -133,14 +115,11 @@ async fn dev_driver(
println!("{}", BANNER.yellow().bold());
info!("{}", "Rebuilding Enclave...".green().bold());
if let Some(shutdown_tx) = shutdown_tx.clone() {
let _res = shutdown_tx.send(());
}
info!("Waiting 1 second for the enclave to shut down");
sleep(Duration::from_secs(1)).await;
let new_shutdown_tx = spawn_enclave_start(args, &config).await?;
// Start enclave in background
spawn_enclave_start(args, &config)?;
// todo: should not unconditionally deploy here
let res = deploy_and_handshake(Some(&contract), args, &config).await;
@ -149,17 +128,12 @@ async fn dev_driver(
Ok(res_contract) => {
// Set state
contract = res_contract;
shutdown_tx = Some(new_shutdown_tx);
info!("{}", "Enclave is listening for requests...".green().bold());
}
Err(e) => {
eprintln!("Error restarting enclave and handshake");
new_shutdown_tx
.send(())
.expect("Could not send signal on channel");
return Err(e);
}
}
@ -173,25 +147,15 @@ async fn dev_driver(
println!("{}", BANNER.yellow().bold());
info!("{}", "Rebuilding Contract...".green().bold());
if let Some(shutdown_tx) = shutdown_tx.clone() {
let res = deploy_and_handshake(None, args, &config).await;
let res = deploy_and_handshake(None, args, &config).await;
match res {
Ok(res_contract) => contract = res_contract,
Err(e) => {
eprintln!("Error deploying contract and handshake:");
match res {
Ok(res_contract) => contract = res_contract,
Err(e) => {
eprintln!("Error deploying contract and handshake:");
shutdown_tx
.send(())
.expect("Could not send signal on channel");
return Err(e);
}
return Err(e);
}
} else {
return Err(Error::GenericErr(
"Attempting to redeploy contract, but enclave isn't running".to_string(),
));
}
info!("{}", "Enclave is listening for requests...".green().bold());
@ -203,14 +167,9 @@ async fn dev_driver(
}
// Spawns enclve start in a separate task which runs in the background
async fn spawn_enclave_start(
args: &DevRequest,
config: &Config,
) -> Result<watch::Sender<()>, Error> {
fn spawn_enclave_start(args: &DevRequest, config: &Config) -> Result<(), Error> {
// In separate process, launch the enclave
let (shutdown_tx, shutdown_rx) = watch::channel(());
let enclave_start = EnclaveStartRequest {
shutdown_rx: Some(shutdown_rx),
unsafe_trust_latest: args.unsafe_trust_latest,
};
@ -222,7 +181,7 @@ async fn spawn_enclave_start(
Ok::<Response, Error>(res)
});
Ok(shutdown_tx)
Ok(())
}
// TODO: do not shutdown if cli calls fail, just print

View file

@ -3,10 +3,7 @@ use std::{fs, path::Path};
use async_trait::async_trait;
use cargo_metadata::MetadataCommand;
use color_eyre::owo_colors::OwoColorize;
use tokio::{
process::{Child, Command},
sync::watch,
};
use tokio::process::{Child, Command};
use tracing::{debug, info};
use crate::{
@ -47,7 +44,7 @@ impl Handler for EnclaveStartRequest {
let enclave_child =
create_mock_enclave_child(config.app_dir.as_path(), config.release, enclave_args)
.await?;
handle_process(self.shutdown_rx, enclave_child).await?;
handle_process(enclave_child).await?;
} else {
let enclave_dir = fs::canonicalize(config.app_dir.join("enclave"))?;
@ -69,40 +66,25 @@ impl Handler for EnclaveStartRequest {
// Run quartz enclave and block
let enclave_child = create_gramine_sgx_child(&enclave_dir).await?;
handle_process(self.shutdown_rx, enclave_child).await?;
handle_process(enclave_child).await?;
}
Ok(EnclaveStartResponse.into())
}
}
async fn handle_process(
shutdown_rx: Option<watch::Receiver<()>>,
mut child: Child,
) -> Result<(), Error> {
info!("{}", "Running enclave ...".green().bold());
match shutdown_rx {
Some(mut rx) => {
tokio::select! {
status = child.wait() => {
handle_child_status(status.map_err(|e| Error::GenericErr(e.to_string()))?)?;
}
_ = rx.changed() => {
info!("Enclave shutdown signal received.");
let _ = child.kill().await;
}
}
}
None => {
// If no shutdown receiver is provided, just wait for the child process
let status = child
.wait()
.await
.map_err(|e| Error::GenericErr(e.to_string()))?;
handle_child_status(status)?;
}
}
async fn handle_process(mut child: Child) -> Result<(), Error> {
let status = child
.wait()
.await
.map_err(|e| Error::GenericErr(e.to_string()))?;
if !status.success() {
return Err(Error::GenericErr(format!(
"Couldn't build enclave. {:?}",
status
)));
}
Ok(())
}
@ -139,22 +121,13 @@ async fn create_mock_enclave_child(
info!("{}", "🚧 Spawning enclave process ...".green().bold());
let child = command
.kill_on_drop(true)
.spawn()
.map_err(|e| Error::GenericErr(e.to_string()))?;
Ok(child)
}
fn handle_child_status(status: std::process::ExitStatus) -> Result<(), Error> {
if !status.success() {
return Err(Error::GenericErr(format!(
"Couldn't build enclave. {:?}",
status
)));
}
Ok(())
}
async fn gramine_sgx_gen_private_key(enclave_dir: &Path) -> Result<(), Error> {
// Launch the gramine-sgx-gen-private-key command
Command::new("gramine-sgx-gen-private-key")
@ -244,6 +217,7 @@ async fn create_gramine_sgx_child(enclave_dir: &Path) -> Result<Child, Error> {
let child = Command::new("gramine-sgx")
.arg("./quartz")
.kill_on_drop(true)
.current_dir(enclave_dir)
.spawn()?;

View file

@ -97,7 +97,6 @@ impl TryFrom<EnclaveCommand> for Request {
match cmd {
EnclaveCommand::Build(_) => Ok(EnclaveBuildRequest {}.into()),
EnclaveCommand::Start(args) => Ok(EnclaveStartRequest {
shutdown_rx: None,
unsafe_trust_latest: args.unsafe_trust_latest,
}
.into()),

View file

@ -1,5 +1,4 @@
use tendermint::{block::Height, Hash};
use tokio::sync::watch;
use tracing::debug;
use crate::{
@ -9,7 +8,6 @@ use crate::{
#[derive(Clone, Debug)]
pub struct EnclaveStartRequest {
pub shutdown_rx: Option<watch::Receiver<()>>,
pub unsafe_trust_latest: bool,
}