Merge pull request #285 from monadicus/feat/misc-task-finishing

Feat: misc task finishing
monadicus · Sep 17, 2024 · 4db274e · 4db274e
2 parents ae44e41 + 8914dc1
commit 4db274e
Show file tree

Hide file tree

Showing 7 changed files with 50 additions and 3 deletions.
diff --git a/crates/aot/src/runner/rpc/node.rs b/crates/aot/src/runner/rpc/node.rs
@@ -34,9 +34,11 @@ pub struct NodeRpcServer<N: Network> {
 }
 
 impl<N: Network> NodeService for NodeRpcServer<N> {
-    async fn set_log_level(self, _: context::Context, verbosity: u8) -> Result<(), AgentError> {
-        tracing::debug!("NodeService Setting log verbosity {verbosity:?}");
+    async fn status(self, _: context::Context) -> Result<(), AgentError> {
+        Ok(())
+    }
 
+    async fn set_log_level(self, _: context::Context, verbosity: u8) -> Result<(), AgentError> {
         self.log_level_handler
             .modify(|filter| *filter = make_env_filter(verbosity))
             .map_err(|_| AgentError::FailedToChangeLogLevel)?;

diff --git a/crates/snops-agent/src/main.rs b/crates/snops-agent/src/main.rs
@@ -189,6 +189,7 @@ async fn main() {
             .execute(
                 AgentRpcServer {
                     state: state.to_owned(),
+                    version: env!("CARGO_PKG_VERSION"),
                 }
                 .serve(),
             )

diff --git a/crates/snops-agent/src/rpc/control.rs b/crates/snops-agent/src/rpc/control.rs
@@ -15,7 +15,8 @@ use snops_common::{
     rpc::{
         control::{
             agent::{
-                AgentMetric, AgentService, AgentServiceRequest, AgentServiceResponse, Handshake,
+                AgentMetric, AgentService, AgentServiceRequest, AgentServiceResponse, AgentStatus,
+                Handshake,
             },
             ControlServiceRequest, ControlServiceResponse,
         },
@@ -42,6 +43,7 @@ define_rpc_mux!(child;
 #[derive(Clone)]
 pub struct AgentRpcServer {
     pub state: AppState,
+    pub version: &'static str,
 }
 
 impl AgentService for AgentRpcServer {
@@ -624,4 +626,14 @@ impl AgentService for AgentRpcServer {
             .await
             .map_err(|_| AgentError::FailedToMakeRequest)?
     }
+
+    async fn get_status(self, ctx: context::Context) -> Result<AgentStatus, AgentError> {
+        let lock = self.state.node_client.lock().await;
+        let node_client = lock.as_ref().ok_or(AgentError::NodeClientNotSet)?;
+
+        Ok(AgentStatus {
+            aot_online: node_client.status(ctx).await.is_ok(),
+            version: self.version.to_string(),
+        })
+    }
 }
diff --git a/crates/snops-cli/src/commands/agent.rs b/crates/snops-cli/src/commands/agent.rs
@@ -70,6 +70,9 @@ enum AgentCommands {
     /// Get the specific agent's TPS.
     Tps,
 
+    /// Get the specific agent's status.
+    Status,
+
     SetLogLevel {
         /// The log level to set.
         level: String,
@@ -138,6 +141,11 @@ impl Agent {
 
                 client.post(ep).send()?
             }
+            Status => {
+                let ep = format!("{url}/api/v1/agents/{}/status", self.id);
+
+                client.get(ep).send()?
+            }
             Tps => {
                 let ep = format!("{url}/api/v1/agents/{}/tps", self.id);
 

diff --git a/crates/snops-common/src/rpc/agent/node.rs b/crates/snops-common/src/rpc/agent/node.rs
@@ -3,6 +3,7 @@ use crate::{rpc::error::AgentError, state::snarkos_status::SnarkOSLiteBlock};
 #[tarpc::service]
 pub trait NodeService {
     // todo @gluax this should return an A different kind of error.
+    async fn status() -> Result<(), AgentError>;
     async fn set_log_level(verbosity: u8) -> Result<(), AgentError>;
     async fn get_block_lite(block_hash: String) -> Result<Option<SnarkOSLiteBlock>, AgentError>;
     async fn find_transaction(tx_id: String) -> Result<Option<String>, AgentError>;

diff --git a/crates/snops-common/src/rpc/control/agent.rs b/crates/snops-common/src/rpc/control/agent.rs
@@ -62,6 +62,14 @@ pub trait AgentService {
     ) -> Result<Option<SnarkOSLiteBlock>, AgentError>;
 
     async fn set_aot_log_level(verbosity: u8) -> Result<(), AgentError>;
+
+    async fn get_status() -> Result<AgentStatus, AgentError>;
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct AgentStatus {
+    pub aot_online: bool,
+    pub version: String,
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize)]

diff --git a/crates/snops/src/server/api.rs b/crates/snops/src/server/api.rs
@@ -44,6 +44,7 @@ pub(super) fn routes() -> Router<AppState> {
         .route("/log/:level", post(set_log_level))
         .route("/agents", get(get_agents))
         .route("/agents/:id", get(get_agent))
+        .route("/agents/:id/status", get(get_agent_status))
         .route("/agents/:id/kill", post(kill_agent))
         .route("/agents/:id/tps", get(get_agent_tps))
         .route("/agents/:id/log/:level", post(set_agent_log_level))
@@ -333,6 +334,20 @@ async fn get_agent(state: State<AppState>, Path(id): Path<String>) -> Response {
     Json(AgentStatusResponse::from(agent.value())).into_response()
 }
 
+async fn get_agent_status(state: State<AppState>, Path(id): Path<String>) -> Response {
+    let id = unwrap_or_not_found!(id_or_none(&id));
+    let agent = unwrap_or_not_found!(state.pool.get(&id));
+
+    let Some(rpc) = agent.rpc() else {
+        return StatusCode::SERVICE_UNAVAILABLE.into_response();
+    };
+
+    match rpc.get_status(tarpc::context::current()).await {
+        Ok(status) => Json(status).into_response(),
+        Err(_e) => StatusCode::INTERNAL_SERVER_ERROR.into_response(),
+    }
+}
+
 async fn kill_agent(state: State<AppState>, Path(id): Path<String>) -> Response {
     let id = unwrap_or_not_found!(id_or_none(&id));
     let client = unwrap_or_not_found!(state.pool.get(&id).and_then(|a| a.client_owned()));