diff --git a/src/api-service/__app__/onefuzzlib/extension.py b/src/api-service/__app__/onefuzzlib/extension.py index 505d13804..e2c554a8b 100644 --- a/src/api-service/__app__/onefuzzlib/extension.py +++ b/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, ), diff --git a/src/api-service/__app__/onefuzzlib/tasks/config.py b/src/api-service/__app__/onefuzzlib/tasks/config.py index 90b445f7d..c471a5e86 100644 --- a/src/api-service/__app__/onefuzzlib/tasks/config.py +++ b/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, ), diff --git a/src/api-service/__app__/pool/__init__.py b/src/api-service/__app__/pool/__init__.py index 1e375fee0..87db6cef4 100644 --- a/src/api-service/__app__/pool/__init__.py +++ b/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, ), diff --git a/src/api-service/__app__/queue_heartbeat/__init__.py b/src/api-service/__app__/queue_node_heartbeat/__init__.py similarity index 72% rename from src/api-service/__app__/queue_heartbeat/__init__.py rename to src/api-service/__app__/queue_node_heartbeat/__init__.py index 296f2fd02..31e81d657 100644 --- a/src/api-service/__app__/queue_heartbeat/__init__.py +++ b/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: diff --git a/src/api-service/__app__/queue_heartbeat/function.json b/src/api-service/__app__/queue_node_heartbeat/function.json similarity index 88% rename from src/api-service/__app__/queue_heartbeat/function.json rename to src/api-service/__app__/queue_node_heartbeat/function.json index 08875c57d..938997ee0 100644 --- a/src/api-service/__app__/queue_heartbeat/function.json +++ b/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" }, { diff --git a/src/api-service/__app__/queue_task_heartbeat/__init__.py b/src/api-service/__app__/queue_task_heartbeat/__init__.py new file mode 100644 index 000000000..ecae1180a --- /dev/null +++ b/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) diff --git a/src/api-service/__app__/queue_task_heartbeat/function.json b/src/api-service/__app__/queue_task_heartbeat/function.json new file mode 100644 index 000000000..6acb7d870 --- /dev/null +++ b/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" + } + ] +} \ No newline at end of file diff --git a/src/deployment/deploy.py b/src/deployment/deploy.py index 4ff3adad8..476e7c131 100644 --- a/src/deployment/deploy.py +++ b/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: