Merge pull request #89 from movementlabsxyz/l-monninger/88

chore: type safe API for embedded Aptos runner

Author: l-monninger

Cargo.lock

@@ -123,12 +123,12 @@ secure-signer-loader = { git = "https://github.com/movementlabsxyz/secure-signin
 secure-signer-aws-kms = { git = "https://github.com/movementlabsxyz/secure-signing.git", rev = "f37eebd6d494d21b26b2faa44e1188de42ff1007" }
 
 # kestrel 
-kestrel = { git = "https://github.com/movementlabsxyz/kestrel.git", rev = "8395e2dce55192269a41f5bacacedaf49555f875" }
-jsonlvar = { git = "https://github.com/movementlabsxyz/kestrel.git", rev = "8395e2dce55192269a41f5bacacedaf49555f875" }
-commander = { git = "https://github.com/movementlabsxyz/kestrel.git", rev = "8395e2dce55192269a41f5bacacedaf49555f875" }
-include-dir = { git = "https://github.com/movementlabsxyz/kestrel.git", rev = "8395e2dce55192269a41f5bacacedaf49555f875" }
-include-vendor = { git = "https://github.com/movementlabsxyz/kestrel.git", rev = "8395e2dce55192269a41f5bacacedaf49555f875" }
-ready-docker = { git = "https://github.com/movementlabsxyz/kestrel.git", rev = "8395e2dce55192269a41f5bacacedaf49555f875" }
+kestrel = { git = "https://github.com/movementlabsxyz/kestrel.git", rev = "96fa1a31088af992748c8dbe6b8fe4ccfba771fa" }
+jsonlvar = { git = "https://github.com/movementlabsxyz/kestrel.git", rev = "96fa1a31088af992748c8dbe6b8fe4ccfba771fa" }
+commander = { git = "https://github.com/movementlabsxyz/kestrel.git", rev = "96fa1a31088af992748c8dbe6b8fe4ccfba771fa" }
+include-dir = { git = "https://github.com/movementlabsxyz/kestrel.git", rev = "96fa1a31088af992748c8dbe6b8fe4ccfba771fa" }
+include-vendor = { git = "https://github.com/movementlabsxyz/kestrel.git", rev = "96fa1a31088af992748c8dbe6b8fe4ccfba771fa" }
+ready-docker = { git = "https://github.com/movementlabsxyz/kestrel.git", rev = "96fa1a31088af992748c8dbe6b8fe4ccfba771fa" }
 
 # orfile
 orfile = { git = "https://github.com/movementlabsxyz/orfile.git", rev = "b49cb42495816fa00a1107be5bbf41aff3a8255d" }

Cargo.toml

@@ -4,7 +4,7 @@ pub mod test {
 	use mtma_migrator_test_types::check::checked_migration;
 	use mtma_migrator_types::migrator::MovementMigrator;
 	use mtma_node_null_core::config::Config as MtmaNullConfig;
-	use mtma_node_test_types::prelude::{Prelude, PreludeError};
+	use mtma_node_test_types::prelude::Prelude;
 
 	#[tokio::test]
 	#[tracing_test::traced_test]

checks/migrator/checks/sketchpad/src/example.rs

@@ -2,7 +2,7 @@ use anyhow::Context;
 use aptos_config::config::NodeConfig;
 use aptos_rest_client::Client as MovementAptosRestClient;
 use kestrel::WaitCondition;
-use movement_aptos_core::MovementAptos;
+use movement_aptos_core::{runtime, MovementAptos};
 use mtma_node_types::executor::MovementAptosNode;
 
 /// An enum supporting different types of runners.
@@ -12,7 +12,7 @@ use mtma_node_types::executor::MovementAptosNode;
 #[derive(Clone)]
 pub enum Runner {
 	/// [MovementAptos] runner.
-	MovementAptos(MovementAptos),
+	MovementAptos(MovementAptos<runtime::TokioTest>),
 }
 
 /// The [MovementAptos] migration struct as would be presented in the criterion.

migration/util/migrator-types/src/migrator/movement_aptos_migrator.rs

@@ -6,7 +6,7 @@ use serde::{Deserialize, Serialize};
 use std::path::PathBuf;
 use std::str::FromStr;
 
-use crate::movement_aptos::MovementAptos;
+use crate::movement_aptos::{runtime, MovementAptos};
 use aptos_node::create_single_node_test_config;
 
 #[derive(Debug, Serialize, Deserialize, Clone)]
@@ -66,11 +66,6 @@ pub struct Config {
 	#[orfile(config)]
 	#[clap(long)]
 	pub log_file: Option<PathBuf>,
-
-	/// Whether to create a global rayon pool.
-	#[orfile(config)]
-	#[clap(long)]
-	pub create_global_rayon_pool: bool,
 }
 
 impl Config {
@@ -90,19 +85,15 @@ impl Config {
 		)
 		.map_err(|e| ConfigError::Internal(e.into()))?;
 
-		Ok(Config {
-			node_config: NodeConfigWrapper(node_config),
-			log_file: None,
-			create_global_rayon_pool: false,
-		})
+		Ok(Config { node_config: NodeConfigWrapper(node_config), log_file: None })
 	}
 
 	/// Builds the config into a [MovementAptos] runner.
-	pub fn build(&self) -> Result<MovementAptos, ConfigError> {
-		Ok(MovementAptos::new(
+	pub fn build(&self) -> Result<MovementAptos<runtime::TokioTest>, ConfigError> {
+		Ok(MovementAptos::<runtime::TokioTest>::try_new(
 			self.node_config.0.clone(),
 			self.log_file.clone(),
-			self.create_global_rayon_pool,
-		))
+		)
+		.map_err(|e| ConfigError::Internal(e.into()))?)
 	}
 }

util/movement-aptos/movement-aptos-core/src/config.rs

@@ -2,8 +2,11 @@ use aptos_config::config::NodeConfig;
 use kestrel::State;
 use std::path::PathBuf;
 pub mod rest_api;
+pub mod runtime;
 pub use rest_api::RestApi;
 
+use runtime::Runtime;
+
 /// Errors thrown when running [MovementAptos].
 #[derive(Debug, thiserror::Error)]
 pub enum MovementAptosError {
@@ -12,28 +15,41 @@ pub enum MovementAptosError {
 }
 
 #[derive(Clone)]
-pub struct MovementAptos {
+pub struct MovementAptos<R>
+where
+	R: Runtime,
+{
 	/// The [NodeConfig] for the Aptos node.
 	pub node_config: NodeConfig,
 	/// The path to the log file.
 	pub log_file: Option<PathBuf>,
-	/// Whether to create a global rayon pool.
-	pub create_global_rayon_pool: bool,
+	/// The runtime for the Aptos node.
+	pub runtime: std::marker::PhantomData<R>,
 	/// The [MovementAptosRestApi] for the Aptos node.
 	pub rest_api: State<RestApi>,
 }
 
-impl MovementAptos {
-	pub fn new(
+impl<R> MovementAptos<R>
+where
+	R: Runtime,
+{
+	/// If you have something that marks your ability to get a runtime, you can use this.
+	pub fn new(node_config: NodeConfig, log_file: Option<PathBuf>, _runtime: R) -> Self {
+		Self { node_config, log_file, runtime: std::marker::PhantomData, rest_api: State::new() }
+	}
+
+	/// Checks runtime availability and creates a new [MovementAptos].
+	pub fn try_new(
 		node_config: NodeConfig,
 		log_file: Option<PathBuf>,
-		create_global_rayon_pool: bool,
-	) -> Self {
-		Self { node_config, log_file, create_global_rayon_pool, rest_api: State::new() }
+	) -> Result<Self, anyhow::Error> {
+		let runtime = R::try_new()?;
+		let movement_aptos = MovementAptos::new(node_config, log_file, runtime);
+		Ok(movement_aptos)
 	}
 
 	pub fn from_config(config: NodeConfig) -> Result<Self, anyhow::Error> {
-		let movement_aptos = MovementAptos::new(config, None, false);
+		let movement_aptos = MovementAptos::new(config, None, R::try_new()?);
 		Ok(movement_aptos)
 	}
 
@@ -43,15 +59,37 @@ impl MovementAptos {
 	}
 
 	/// Runs the internal node logic
-	pub(crate) fn run_node(&self) -> Result<(), MovementAptosError> {
-		aptos_node::start(
-			self.node_config.clone(),
-			self.log_file.clone(),
-			self.create_global_rayon_pool,
-		)
-		.map_err(|e| MovementAptosError::Internal(e.into()))?;
-
-		Ok(())
+	pub(crate) async fn run_node(&self) -> Result<(), MovementAptosError> {
+		// Clone necessary data for the closure
+		let node_config = self.node_config.clone();
+		let log_file = self.log_file.clone();
+
+		// Spawn the blocking task
+		let blocking_task_result = tokio::task::spawn_blocking(move || {
+			// This closure runs on a blocking thread
+			aptos_node::start(
+				node_config,
+				log_file,
+				R::create_global_rayon_pool(), // Assuming R is in scope and its result is Send
+			)
+			// The closure should return the direct result from aptos_node::start.
+			// The error type from aptos_node::start (let's call it AptosNodeError)
+			// needs to be Send + 'static for the closure.
+		})
+		.await;
+
+		match blocking_task_result {
+			Ok(Ok(())) => Ok(()), // aptos_node::start succeeded
+			Ok(Err(aptos_node_err)) => {
+				// aptos_node::start failed. We need aptos_node_err to be convertible
+				// into the Box<dyn Error> for MovementAptosError::Internal.
+				Err(MovementAptosError::Internal(aptos_node_err.into()))
+			}
+			Err(join_err) => {
+				// spawn_blocking task failed (e.g., panicked or was cancelled by Tokio)
+				Err(MovementAptosError::Internal(Box::new(join_err)))
+			}
+		}
 	}
 
 	/// Runs the node and fills state.
@@ -66,18 +104,22 @@ impl MovementAptos {
 
 		let runner = self.clone();
 		let runner_task = kestrel::task(async move {
-			runner.run_node()?;
+			runner.run_node().await?;
 			Ok::<_, MovementAptosError>(())
 		});
 
 		// rest api state
 		let rest_api_state = self.rest_api.clone();
 		let rest_api_polling = kestrel::task(async move {
 			loop {
+				tokio::time::sleep(std::time::Duration::from_secs(1)).await;
+				println!("POLLING REST API: {:?}", rest_api);
 				// wait for the rest api to be ready
 				let response = reqwest::get(rest_api.rest_api_url.clone())
 					.await
 					.map_err(|e| MovementAptosError::Internal(e.into()))?;
+
+				println!("REST API RESPONSE: {:?}", response);
 				if response.status().is_success() {
 					rest_api_state.write().set(rest_api).await;
 					break;
@@ -102,7 +144,7 @@ mod tests {
 	use rand::thread_rng;
 	use std::path::Path;
 
-	#[tokio::test]
+	#[tokio::test(flavor = "multi_thread")]
 	async fn test_movement_aptos() -> Result<(), anyhow::Error> {
 		// open in a new db
 		let unique_id = uuid::Uuid::new_v4();
@@ -128,11 +170,20 @@ mod tests {
 			rng,
 		)?;
 
-		let movement_aptos = MovementAptos::new(node_config, None, false);
+		let movement_aptos = MovementAptos::<runtime::TokioTest>::try_new(node_config, None)?;
 		let rest_api_state = movement_aptos.rest_api().read().clone();
-		movement_aptos.run().await?;
+
+		let movement_aptos_task = kestrel::task(async move {
+			movement_aptos.run().await?;
+			Ok::<_, MovementAptosError>(())
+		});
+
 		rest_api_state.wait_for(tokio::time::Duration::from_secs(30)).await?;
 
+		println!("ENDING MOVEMENT APTOS");
+
+		kestrel::end!(movement_aptos_task)?;
+
 		Ok(())
 	}
 }