From 06f45f338c60aed3e0c73fa29ce28e7ca2cc1a55 Mon Sep 17 00:00:00 2001
From: nharper285 <74685766+nharper285@users.noreply.github.com>
Date: Fri, 26 Feb 2021 10:36:10 -0800
Subject: [PATCH] Update Task Heartbeat to include Job_id (#594)

---
 src/agent/onefuzz-agent/src/tasks/config.rs              | 2 +-
 src/agent/onefuzz-agent/src/tasks/heartbeat.rs           | 8 +++++++-
 src/api-service/__app__/queue_task_heartbeat/__init__.py | 2 +-
 src/pytypes/onefuzztypes/models.py                       | 1 +
 4 files changed, 10 insertions(+), 3 deletions(-)

diff --git a/src/agent/onefuzz-agent/src/tasks/config.rs b/src/agent/onefuzz-agent/src/tasks/config.rs
index d2d0a0e17..730f21c5a 100644
--- a/src/agent/onefuzz-agent/src/tasks/config.rs
+++ b/src/agent/onefuzz-agent/src/tasks/config.rs
@@ -47,7 +47,7 @@ impl CommonConfig {
     pub async fn init_heartbeat(&self) -> Result<Option<TaskHeartbeatClient>> {
         match &self.heartbeat_queue {
             Some(url) => {
-                let hb = init_task_heartbeat(url.clone(), self.task_id).await?;
+                let hb = init_task_heartbeat(url.clone(), self.task_id, self.job_id).await?;
                 Ok(Some(hb))
             }
             None => Ok(None),
diff --git a/src/agent/onefuzz-agent/src/tasks/heartbeat.rs b/src/agent/onefuzz-agent/src/tasks/heartbeat.rs
index 10ac41f1c..634ecada6 100644
--- a/src/agent/onefuzz-agent/src/tasks/heartbeat.rs
+++ b/src/agent/onefuzz-agent/src/tasks/heartbeat.rs
@@ -26,18 +26,24 @@ struct Heartbeat {
 #[derive(Clone)]
 pub struct TaskContext {
     task_id: Uuid,
+    job_id: Uuid,
     machine_id: Uuid,
     machine_name: String,
 }
 
 pub type TaskHeartbeatClient = HeartbeatClient<TaskContext, HeartbeatData>;
 
-pub async fn init_task_heartbeat(queue_url: Url, task_id: Uuid) -> Result<TaskHeartbeatClient> {
+pub async fn init_task_heartbeat(
+    queue_url: Url,
+    task_id: Uuid,
+    job_id: Uuid,
+) -> Result<TaskHeartbeatClient> {
     let machine_id = get_machine_id().await?;
     let machine_name = get_machine_name().await?;
     let hb = HeartbeatClient::init_heartbeat(
         TaskContext {
             task_id,
+            job_id,
             machine_id,
             machine_name,
         },
diff --git a/src/api-service/__app__/queue_task_heartbeat/__init__.py b/src/api-service/__app__/queue_task_heartbeat/__init__.py
index b81b2880a..e1bc69aa2 100644
--- a/src/api-service/__app__/queue_task_heartbeat/__init__.py
+++ b/src/api-service/__app__/queue_task_heartbeat/__init__.py
@@ -30,7 +30,7 @@ def main(msg: func.QueueMessage, dashboard: func.Out[str]) -> None:
             task.heartbeat = datetime.utcnow()
             task.save()
     except ValidationError:
-        logging.error("invalid task heartbat: %s", raw)
+        logging.error("invalid task heartbeat: %s", raw)
 
     events = get_events()
     if events:
diff --git a/src/pytypes/onefuzztypes/models.py b/src/pytypes/onefuzztypes/models.py
index 02d5888b8..f6a4996c0 100644
--- a/src/pytypes/onefuzztypes/models.py
+++ b/src/pytypes/onefuzztypes/models.py
@@ -516,6 +516,7 @@ class Job(BaseModel):
 
 class TaskHeartbeatEntry(BaseModel):
     task_id: UUID
+    job_id: Optional[UUID]
     machine_id: UUID
     data: List[Dict[str, HeartbeatType]]