Merge pull request #83 from Thaumy/main

Supports notifying remote server of task completion

src/main.rs

@@ -68,6 +68,7 @@ fn main() -> Result<()> {
 
     if let Some(args) = component_args {
         let task = Task {
+            id: None,
             wasm_component: fs::read(&args[0])?,
             wasm_component_args: args,
             end_time: Utc.with_ymd_and_hms(3000, 1, 1, 1, 1, 1).unwrap(),
@@ -106,12 +107,17 @@ async fn async_tasks(
         task_rt.spawn(Some(client.clone()))?;
         client.send_info().await?;
         loop {
-            if let Some(task) = client.heartbeat(task_rt.is_idle()).await? {
+            let finished_task_id = task_rt.finished_task_id();
+            if let Some(task) = client
+                .heartbeat(task_rt.is_idle(), finished_task_id)
+                .await?
+            {
                 let end_time = match Utc.timestamp_millis_opt(task.end_time as _) {
                     LocalResult::Single(t) => t,
                     _ => bail!("Invalid task end time"),
                 };
                 let task = Task {
+                    id: Some(task.id),
                     wasm_component: task.wasm,
                     wasm_component_args: task.wasm_args,
                     end_time,

src/runtime/mod.rs

@@ -26,6 +26,7 @@ use std::sync::mpsc::channel;
 use std::sync::mpsc::Receiver;
 use std::sync::mpsc::Sender;
 use std::sync::Arc;
+use std::sync::Mutex;
 use std::thread;
 use std::thread::JoinHandle;
 
@@ -41,6 +42,7 @@ pub use state::PshState;
 use crate::services::rpc::RpcClient;
 
 pub struct Task {
+    pub id: Option<String>,
     pub wasm_component: Vec<u8>,
     pub wasm_component_args: Vec<String>,
     pub end_time: DateTime<Utc>,
@@ -50,6 +52,7 @@ pub struct TaskRuntime {
     tx: Sender<Task>,
     rx: Option<Receiver<Task>>,
     len: Arc<AtomicUsize>,
+    finished_task_id: Arc<Mutex<Vec<String>>>,
 }
 
 impl TaskRuntime {
@@ -60,6 +63,7 @@ impl TaskRuntime {
             tx,
             rx: Some(rx),
             len: Arc::new(AtomicUsize::new(0)),
+            finished_task_id: Arc::new(Mutex::new(vec![])),
         })
     }
 
@@ -74,6 +78,10 @@ impl TaskRuntime {
         len == 0
     }
 
+    pub fn finished_task_id(&mut self) -> Option<String> {
+        self.finished_task_id.lock().unwrap().pop()
+    }
+
     pub fn spawn(&mut self, rpc_client: Option<RpcClient>) -> Result<JoinHandle<()>> {
         let rx = match self.rx.take() {
             Some(rx) => rx,
@@ -84,6 +92,7 @@ impl TaskRuntime {
         let data_export_ctx = DataExportCtx { rpc_client };
 
         let len = self.len.clone();
+        let finished_task_id = self.finished_task_id.clone();
         let handle = thread::spawn(move || {
             while let Ok(task) = rx.recv() {
                 let mut envs = envs.clone();
@@ -108,7 +117,9 @@ impl TaskRuntime {
                     }
                     Err(e) => eprintln!("{}", e),
                 };
-
+                if let Some(id) = task.id {
+                    finished_task_id.lock().unwrap().push(id);
+                }
                 len.fetch_sub(1, Ordering::Release);
             }
         });

src/services/host_info.rs

@@ -58,6 +58,7 @@ impl From<RawInfo> for HostInfoRequest {
             local_ipv6_addr: value.ipv6.map(|v| v.into()),
             instance_id: value.instance_id,
             idle: value.idle,
+            task_id: value.task_id,
         }
     }
 }
@@ -73,6 +74,7 @@ impl From<&RawInfo> for HostInfoRequest {
             local_ipv6_addr: value.ipv6.map(|v| v.into()),
             instance_id: value.instance_id.clone(),
             idle: value.idle,
+            task_id: value.task_id.clone(),
         }
     }
 }
@@ -87,6 +89,7 @@ pub struct RawInfo {
     hostname: Option<String>,
     instance_id: Option<String>,
     idle: bool,
+    task_id: Option<String>,
 }
 
 impl RawInfo {
@@ -117,6 +120,7 @@ impl RawInfo {
             kernel_version: None,
             instance_id,
             idle: false,
+            task_id: None,
         };
 
         let cpu_hd = CpuHandle::new();
@@ -162,6 +166,7 @@ impl RawInfo {
             hostname: None,
             instance_id: self.instance_id.clone(),
             idle: false,
+            task_id: None,
         }
     }
 

src/services/rpc.rs

@@ -81,11 +81,16 @@ impl RpcClient {
         Ok(())
     }
 
-    pub async fn heartbeat(&mut self, idle: bool) -> Result<Option<pb::Task>> {
+    pub async fn heartbeat(
+        &mut self,
+        idle: bool,
+        finished_task_id: Option<String>,
+    ) -> Result<Option<pb::Task>> {
         let req: Request<HostInfoRequest> = {
             let raw_info = self.raw_info.to_heartbeat();
             let mut req: HostInfoRequest = raw_info.into();
             req.idle = idle;
+            req.task_id = finished_task_id;
             let mut req = Request::new(req);
             req.metadata_mut()
                 .insert("authorization", format!("Bearer {}", self.token).parse()?);
@@ -215,6 +220,7 @@ mod rpc_tests {
             local_ipv6_addr: None,
             instance_id: None,
             idle: false,
+            task_id: None,
         };
 
         let heartbeat = test_send_info(client, info_req);