save

Author: shinnku-nikaidou

src/components/controlled_mode/server.rs

@@ -58,9 +58,16 @@ pub async fn connect_to_server(url: Url) -> anyhow::Result<()> {
     let self_ip = get_self_ip().await?;
     let mut grpc_client = HeartbeatServiceClient::connect(url.to_string()).await?;
     let mut executor = ServerCommandExecutor::new();
+    let mut stop_time = chrono::Utc::now().naive_utc() + chrono::Duration::days(30);
 
     loop {
         let now = chrono::Utc::now().naive_utc();
+        if now > stop_time {
+            log::info!("Stopping execution as the time limit has been reached.");
+            executor.shutdown_workers().await?;
+            stop_time = chrono::Utc::now().naive_utc() + chrono::Duration::days(30);
+        }
+
         let ServerResponse {
             server_timestamp,
             next_operation,
@@ -91,17 +98,38 @@ pub async fn connect_to_server(url: Url) -> anyhow::Result<()> {
                 NextOperation::ContinueCurrent(Empty {}) => {}
                 NextOperation::StopCurrent(Empty {}) => {
                     executor.shutdown_workers().await?;
+                    stop_time = chrono::Utc::now().naive_utc() + chrono::Duration::days(30);
                 }
                 NextOperation::StopAndExecute(commands) => {
                     executor.shutdown_workers().await?;
-                    executor
+                    stop_time = chrono::Utc::now().naive_utc() + chrono::Duration::days(30);
+                    let max_time = executor
                         .execute(commands.try_into()?, command_id.unwrap_or(0))
                         .await;
+                    match max_time {
+                        Ok(max_time) => {
+                            stop_time = now + chrono::Duration::seconds(max_time as i64);
+                        }
+                        Err(e) => {
+                            log::error!("Failed to execute commands: {}", e);
+                            continue;
+                        }
+                    }
                 }
                 NextOperation::Execute(commands) => {
-                    executor
+                    stop_time = chrono::Utc::now().naive_utc() + chrono::Duration::days(30);
+                    let max_time = executor
                         .execute(commands.try_into()?, command_id.unwrap_or(0))
                         .await;
+                    match max_time {
+                        Ok(max_time) => {
+                            stop_time = now + chrono::Duration::seconds(max_time as i64);
+                        }
+                        Err(e) => {
+                            log::error!("Failed to execute commands: {}", e);
+                            continue;
+                        }
+                    }
                 }
             }
         }

src/components/controlled_mode/server_command_executor.rs

@@ -5,7 +5,9 @@ use crate::components::controlled_mode::server::{commands, heartbeat, send_heart
 use crate::components::controlled_mode::server_command::{ParallelCommands, RemoteCommand};
 use chrono::NaiveDateTime;
 use log::{info, warn};
+use serde::de;
 use std::sync::Arc;
+use std::time::Duration;
 use tokio::sync::{Mutex, mpsc};
 use tokio::task::JoinSet;
 use tonic::transport::Channel;
@@ -104,7 +106,7 @@ impl ServerCommandExecutor {
         Ok(())
     }
 
-    pub async fn execute(&mut self, commands: ParallelCommands, id: u64) {
+    pub async fn execute(&mut self, commands: ParallelCommands, id: u64) -> anyhow::Result<u64> {
         log::debug!("Executing commands: {:?}", commands);
 
         let now = chrono::Utc::now().naive_utc();
@@ -123,6 +125,16 @@ impl ServerCommandExecutor {
                 None => true,
             })
             .collect::<Vec<_>>();
+        let max_time = commands
+            .iter()
+            .filter_map(|c| match c {
+                RemoteCommand::Request(c) => c.time,
+                RemoteCommand::Shell(_) => Some(0),
+            })
+            .collect::<Vec<u64>>()
+            .iter()
+            .fold(0, |arg0, arg1: &u64| u64::max(arg0, *arg1));
+        log::debug!("max_time: {}", max_time);
         for command in commands {
             match command {
                 RemoteCommand::Request(request) => {
@@ -144,7 +156,7 @@ impl ServerCommandExecutor {
                 }
             }
         }
-        
+        Ok(max_time)
     }
 
     pub async fn check_idle(&mut self) {