[aptos-workspace-server] add timeout + refactor cli command

aptos-move/aptos-workspace-server/Cargo.toml

@@ -23,6 +23,7 @@ aptos-types = { workspace = true }
 # third party deps
 anyhow = { workspace = true }
 bollard = { workspace = true }
+clap = { workspace = true }
 diesel = { workspace = true, features = [
     "postgres_backend",
 ] }

aptos-move/aptos-workspace-server/src/lib.rs

@@ -28,16 +28,18 @@ mod services;
 
 use anyhow::{anyhow, Context, Result};
 use aptos_localnet::docker::get_docker;
+use clap::Parser;
 use common::make_shared;
 use futures::TryFutureExt;
 use services::{
     docker_common::create_docker_network_permanent, indexer_api::start_indexer_api,
     processors::start_all_processors,
 };
+use std::time::Duration;
 use tokio_util::sync::CancellationToken;
 use uuid::Uuid;
 
-pub async fn run_all_services() -> Result<()> {
+async fn run_all_services(timeout: u64) -> Result<()> {
     let test_dir = tempfile::tempdir()?;
     let test_dir = test_dir.path();
     println!("Created test directory: {}", test_dir.display());
@@ -51,9 +53,15 @@ pub async fn run_all_services() -> Result<()> {
         //       waiting for it to trigger.
         let shutdown = shutdown.clone();
         tokio::spawn(async move {
-            tokio::signal::ctrl_c().await.unwrap();
-
-            println!("\nCtrl-C received. Shutting down services. This may take a while.\n");
+            tokio::select! {
+                res = tokio::signal::ctrl_c() => {
+                    res.unwrap();
+                    println!("\nCtrl-C received. Shutting down services. This may take a while.\n");
+                }
+                _ = tokio::time::sleep(Duration::from_secs(timeout)) => {
+                    println!("\nTimeout reached. Shutting down services. This may take a while.\n");
+                }
+            }
 
             shutdown.cancel();
         });
@@ -189,3 +197,19 @@ pub async fn run_all_services() -> Result<()> {
 
     Ok(())
 }
+
+#[derive(Parser)]
+pub enum WorkspaceCommand {
+    Run {
+        #[arg(long, default_value_t = 1800)]
+        timeout: u64,
+    },
+}
+
+impl WorkspaceCommand {
+    pub async fn run(self) -> Result<()> {
+        match self {
+            WorkspaceCommand::Run { timeout } => run_all_services(timeout).await,
+        }
+    }
+}

aptos-move/aptos-workspace-server/src/main.rs

@@ -2,10 +2,11 @@
 // SPDX-License-Identifier: Apache-2.0
 
 use anyhow::Result;
+use clap::Parser;
 
 #[tokio::main]
 async fn main() -> Result<()> {
-    aptos_workspace_server::run_all_services().await?;
-
-    Ok(())
+    aptos_workspace_server::WorkspaceCommand::parse()
+        .run()
+        .await
 }

crates/aptos/src/lib.rs

@@ -22,6 +22,7 @@ use crate::common::{
     types::{CliCommand, CliResult, CliTypedResult},
     utils::cli_build_information,
 };
+use aptos_workspace_server::WorkspaceCommand;
 use async_trait::async_trait;
 use clap::Parser;
 use std::collections::BTreeMap;
@@ -53,7 +54,7 @@ pub enum Tool {
     #[clap(subcommand)]
     Update(update::UpdateTool),
     #[clap(subcommand, hide(true))]
-    Workspace(workspace::WorkspaceTool),
+    Workspace(WorkspaceCommand),
 }
 
 impl Tool {
@@ -73,7 +74,7 @@ impl Tool {
             Node(tool) => tool.execute().await,
             Stake(tool) => tool.execute().await,
             Update(tool) => tool.execute().await,
-            Workspace(tool) => tool.execute().await,
+            Workspace(workspace) => workspace.execute_serialized_without_logger().await,
         }
     }
 }

crates/aptos/src/workspace/mod.rs

@@ -1,37 +1,18 @@
 // Copyright (c) Aptos Foundation
 // SPDX-License-Identifier: Apache-2.0
 
-use crate::common::types::{CliCommand, CliResult, CliTypedResult};
+use crate::common::types::{CliCommand, CliTypedResult};
+use aptos_workspace_server::WorkspaceCommand;
 use async_trait::async_trait;
-use clap::Parser;
-
-/// Tool for operations related to Aptos Workspace
-#[derive(Parser)]
-pub enum WorkspaceTool {
-    Run(RunWorkspace),
-}
-
-impl WorkspaceTool {
-    pub async fn execute(self) -> CliResult {
-        use WorkspaceTool::*;
-
-        match self {
-            Run(cmd) => cmd.execute_serialized_without_logger().await,
-        }
-    }
-}
-
-#[derive(Parser)]
-pub struct RunWorkspace;
 
 #[async_trait]
-impl CliCommand<()> for RunWorkspace {
+impl CliCommand<()> for WorkspaceCommand {
     fn command_name(&self) -> &'static str {
-        "RunWorkspace"
+        "Workspace"
     }
 
     async fn execute(self) -> CliTypedResult<()> {
-        aptos_workspace_server::run_all_services().await?;
+        self.run().await?;
 
         Ok(())
     }