address code review findings

Author: 97nitt

apollo/egress/agent/service/base_egress_service.py

@@ -2,6 +2,7 @@
 import os
 import signal
 import sys
+import threading
 import uuid
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
@@ -195,7 +196,7 @@ def __init__(
                 config_manager=config_manager,
                 push_metrics_handler=self._push_metrics,
             )
-        self._shutting_down = False
+        self._shutdown_lock = threading.Lock()
         self._operations_mapping = [
             OperationMapping(
                 path="/api/v1/agent/execute/",
@@ -271,12 +272,11 @@ def _trigger_graceful_shutdown(self):
         """Notify orchestrator, stop all threads, and terminate the process.
 
         Safe to call from any thread. Cleanup runs exactly once (guarded by
-        _shutting_down flag). In-flight operations are abandoned — the
-        orchestrator requeues them via the shutdown notification.
+        _shutdown_lock). In-flight operations are abandoned — the orchestrator
+        requeues them via the shutdown notification.
         """
-        if self._shutting_down:
+        if not self._shutdown_lock.acquire(blocking=False):
             return
-        self._shutting_down = True
         try:
             self._backend_client.notify_shutdown()
             logger.info("Notified orchestrator of shutdown")
@@ -302,8 +302,7 @@ def register_signal_handlers(self):
         def _signal_handler(signum: int, frame: Any):
             sig_name = signal.Signals(signum).name
             logger.info(f"Received {sig_name}, shutting down")
-            if not self._shutting_down:
-                self._trigger_graceful_shutdown()
+            self._trigger_graceful_shutdown()
             sys.exit(0)
 
         signal.signal(signal.SIGTERM, _signal_handler)

tests/test_backend_client.py

@@ -1,7 +1,78 @@
 from unittest import TestCase
 from unittest.mock import Mock, patch
 
-from apollo.egress.agent.backend.backend_client import BackendClient
+from apollo.egress.agent.backend.backend_client import BackendClient, INSTANCE_ID_HEADER
+
+
+class BackendClientTests(TestCase):
+    def setUp(self):
+        self._login_token_provider = Mock()
+        self._login_token_provider.get_token.return_value = {
+            "x-mcd-id": "test-id",
+            "x-mcd-token": "test-token",
+        }
+        self._client = BackendClient(
+            backend_service_url="https://orchestrator.test",
+            login_token_provider=self._login_token_provider,
+        )
+
+    def test_instance_id_is_generated(self):
+        """Instance ID should be a non-empty string, unique per client."""
+        self.assertIsInstance(self._client.instance_id, str)
+        self.assertTrue(len(self._client.instance_id) > 0)
+
+        other = BackendClient(
+            backend_service_url="https://orchestrator.test",
+            login_token_provider=self._login_token_provider,
+        )
+        self.assertNotEqual(self._client.instance_id, other.instance_id)
+
+    def test_headers_include_instance_id(self):
+        """All requests should include the instance ID header."""
+        headers = self._client._headers()
+        self.assertEqual(headers[INSTANCE_ID_HEADER], self._client.instance_id)
+        self.assertEqual(headers["x-mcd-id"], "test-id")
+        self.assertEqual(headers["x-mcd-token"], "test-token")
+
+    def test_headers_include_extra(self):
+        """Extra headers should be merged."""
+        headers = self._client._headers(**{"Content-Type": "application/json"})
+        self.assertEqual(headers["Content-Type"], "application/json")
+        self.assertIn(INSTANCE_ID_HEADER, headers)
+
+    @patch("requests.post")
+    def test_send_heartbeat(self, mock_post):
+        """send_heartbeat should POST to /api/v1/agent/heartbeat with correct headers."""
+        mock_post.return_value.status_code = 200
+
+        self._client.send_heartbeat()
+
+        mock_post.assert_called_once()
+        args, kwargs = mock_post.call_args
+        self.assertEqual(args[0], "https://orchestrator.test/api/v1/agent/heartbeat")
+        self.assertIn(INSTANCE_ID_HEADER, kwargs["headers"])
+        self.assertEqual(kwargs["timeout"], 10)
+
+    @patch("requests.post")
+    def test_notify_shutdown(self, mock_post):
+        """notify_shutdown should POST to /api/v1/agent/shutdown with correct headers."""
+        mock_post.return_value.status_code = 200
+
+        self._client.notify_shutdown()
+
+        mock_post.assert_called_once()
+        args, kwargs = mock_post.call_args
+        self.assertEqual(args[0], "https://orchestrator.test/api/v1/agent/shutdown")
+        self.assertIn(INSTANCE_ID_HEADER, kwargs["headers"])
+        self.assertEqual(kwargs["timeout"], 15)
+
+    @patch("requests.post")
+    def test_notify_shutdown_raises_on_failure(self, mock_post):
+        """notify_shutdown should raise on HTTP errors."""
+        mock_post.return_value.raise_for_status.side_effect = Exception("500")
+
+        with self.assertRaises(Exception):
+            self._client.notify_shutdown()
 
 
 class BackendClientURLTests(TestCase):

tests/test_operations_poller.py

@@ -220,6 +220,25 @@ def can_accept_work():
         self.assertGreater(call_count[0], 1)
         poller.stop()
 
+    def test_backpressure_sends_heartbeat(self):
+        """Test that poller sends heartbeat to orchestrator during backpressure."""
+        self._backend_client.get_next_operation.return_value = None
+
+        self._config_manager.get_int_value.return_value = 0.1
+
+        poller = OperationsPoller(
+            backend_client=self._backend_client,
+            config_manager=self._config_manager,
+            operation_handler=self._operation_handler,
+            can_accept_work=lambda: False,  # always backpressured
+        )
+
+        poller.start()
+        time.sleep(0.3)
+
+        self._backend_client.send_heartbeat.assert_called()
+        poller.stop()
+
     def test_no_backpressure_when_can_accept_work_is_none(self):
         """Test that poller fetches normally when can_accept_work is not provided."""
         op1 = {"operation_id": "op-1", "path": "/test1", "operation": {}}