split node and task heartbeats in two nodes (#163)

src/api-service/__app__/onefuzzlib/extension.py

@@ -95,7 +95,7 @@ def build_pool_config(pool_name: str) -> str:
         onefuzz_url=get_instance_url(),
         instrumentation_key=os.environ.get("APPINSIGHTS_INSTRUMENTATIONKEY"),
         heartbeat_queue=get_queue_sas(
-            "heartbeat",
+            "node-heartbeat",
             account_id=os.environ["ONEFUZZ_FUNC_STORAGE"],
             add=True,
         ),

src/api-service/__app__/onefuzzlib/tasks/config.py

@@ -180,7 +180,7 @@ def build_task_config(
         instrumentation_key=os.environ.get("APPINSIGHTS_INSTRUMENTATIONKEY"),
         telemetry_key=os.environ.get("ONEFUZZ_TELEMETRY"),
         heartbeat_queue=get_queue_sas(
-            "heartbeat",
+            "task-heartbeat",
             account_id=os.environ["ONEFUZZ_FUNC_STORAGE"],
             add=True,
         ),

src/api-service/__app__/pool/__init__.py

@@ -25,7 +25,7 @@ def set_config(pool: Pool) -> Pool:
         instrumentation_key=os.environ.get("APPINSIGHTS_INSTRUMENTATIONKEY"),
         telemetry_key=os.environ.get("ONEFUZZ_TELEMETRY"),
         heartbeat_queue=get_queue_sas(
-            "heartbeat",
+            "node-heartbeat",
             account_id=os.environ["ONEFUZZ_FUNC_STORAGE"],
             add=True,
         ),

src/api-service/__app__/queue_node_heartbeat/__init__.py

@@ -9,7 +9,7 @@ import logging
 import azure.functions as func
 
 from ..onefuzzlib.dashboard import get_event
-from ..onefuzzlib.heartbeat import NodeHeartbeat, TaskHeartbeat
+from ..onefuzzlib.heartbeat import NodeHeartbeat
 
 
 def main(msg: func.QueueMessage, dashboard: func.Out[str]) -> None:
@@ -18,10 +18,7 @@ def main(msg: func.QueueMessage, dashboard: func.Out[str]) -> None:
 
     raw = json.loads(body)
 
-    if TaskHeartbeat.try_add(raw):
-        pass
-    elif NodeHeartbeat.try_add(raw):
-        pass
+    NodeHeartbeat.try_add(raw)
 
     event = get_event()
     if event:

src/api-service/__app__/queue_node_heartbeat/function.json

@@ -5,7 +5,7 @@
 			"name": "msg",
 			"type": "queueTrigger",
 			"direction": "in",
-			"queueName": "heartbeat",
+			"queueName": "node-heartbeat",
 			"connection": "AzureWebJobsStorage"
 		},
 		{

src/api-service/__app__/queue_task_heartbeat/__init__.py

@@ -0,0 +1,25 @@
+#!/usr/bin/env python
+#
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+import json
+import logging
+
+import azure.functions as func
+
+from ..onefuzzlib.dashboard import get_event
+from ..onefuzzlib.heartbeat import TaskHeartbeat
+
+
+def main(msg: func.QueueMessage, dashboard: func.Out[str]) -> None:
+    body = msg.get_body()
+    logging.info("heartbeat: %s", body)
+
+    raw = json.loads(body)
+
+    TaskHeartbeat.try_add(raw)
+
+    event = get_event()
+    if event:
+        dashboard.set(event)

src/api-service/__app__/queue_task_heartbeat/function.json

@@ -0,0 +1,18 @@
+{
+	"scriptFile": "__init__.py",
+	"bindings": [
+		{
+			"name": "msg",
+			"type": "queueTrigger",
+			"direction": "in",
+			"queueName": "task-heartbeat",
+			"connection": "AzureWebJobsStorage"
+		},
+		{
+			"type": "signalR",
+			"direction": "out",
+			"name": "dashboard",
+			"hubName": "dashboard"
+		}
+	]
+}

src/deployment/deploy.py

@@ -367,7 +367,7 @@ class Client:
             account_url=account_url,
             credential={"account_name": name, "account_key": key},
         )
-        for queue in ["file-changes", "heartbeat", "proxy", "update-queue"]:
+        for queue in ["file-changes", "task-heartbeat", "node-heartbeat", "proxy", "update-queue"]:
             try:
                 client.create_queue(queue)
             except ResourceExistsError: