fix: Add timeout to IMDS requests to prevent shutdown hang

src/deadline_worker_agent/worker.py

@@ -61,6 +61,11 @@ class Worker:
     """The amount of time to allow the Worker to gracefully shutdown after detecting an auto-scaling
     life-cycle event."""
 
+    _IMDS_REQUEST_TIMEOUT_SECONDS = 0.5
+    """Timeout for EC2 instance metadata service requests. Non-EC2 hosts may blackhole the
+    IMDS address rather than refuse the connection; without a timeout those requests block
+    forever. Matches the timeout used by startup.bootstrap._get_metadata."""
+
     _farm_id: str
     _fleet_id: str
     _worker_id: str
@@ -391,8 +396,12 @@ def _get_ec2_metadata_imdsv2_token(self) -> str | None:
             response = requests.put(
                 "http://169.254.169.254/latest/api/token",
                 headers={"X-aws-ec2-metadata-token-ttl-seconds": "10"},
+                # Non-EC2 hosts may blackhole the IMDS address instead of refusing the
+                # connection. Without a timeout this request can block forever, which
+                # prevents the agent's graceful shutdown from ever running.
+                timeout=Worker._IMDS_REQUEST_TIMEOUT_SECONDS,
             )
-        except requests.ConnectionError:
+        except (requests.ConnectionError, requests.Timeout):
             # Could not connect to the metadata service. Either it's not enabled or we're not
             # on an EC2 instance.
             return None
@@ -422,8 +431,9 @@ def _get_spot_instance_shutdown_action_timeout(self, *, imdsv2_token: str) -> ti
             response = requests.get(
                 "http://169.254.169.254/latest/meta-data/spot/instance-action",
                 headers={"X-aws-ec2-metadata-token": imdsv2_token},
+                timeout=Worker._IMDS_REQUEST_TIMEOUT_SECONDS,
             )
-        except requests.ConnectionError:
+        except (requests.ConnectionError, requests.Timeout):
             # Could not connect to the metadata service. Either it's inactive or we're not
             # on an EC2 instance.
             return None
@@ -473,8 +483,9 @@ def _is_asg_terminated(self, *, imdsv2_token: str) -> bool:
             response = requests.get(
                 "http://169.254.169.254/latest/meta-data/autoscaling/target-lifecycle-state",
                 headers={"X-aws-ec2-metadata-token": imdsv2_token},
+                timeout=Worker._IMDS_REQUEST_TIMEOUT_SECONDS,
             )
-        except requests.ConnectionError:
+        except (requests.ConnectionError, requests.Timeout):
             return False
 
         if response.status_code == 200:

test/unit/test_worker.py

@@ -406,6 +406,7 @@ def test_get_imdsv2_token(self, worker: Worker, requests_put: MagicMock) -> None
         requests_put.assert_called_once_with(
             "http://169.254.169.254/latest/api/token",
             headers={"X-aws-ec2-metadata-token-ttl-seconds": "10"},
+            timeout=Worker._IMDS_REQUEST_TIMEOUT_SECONDS,
         )
 
     def test_get_imdsv2_token_cannot_connect(self, worker: Worker, requests_put: MagicMock) -> None:
@@ -418,6 +419,19 @@ def test_get_imdsv2_token_cannot_connect(self, worker: Worker, requests_put: Mag
         # THEN
         assert result is None
 
+    def test_get_imdsv2_token_request_times_out(
+        self, worker: Worker, requests_put: MagicMock
+    ) -> None:
+        """A blackholed IMDS address must be treated as IMDS-unavailable, not hang."""
+        # GIVEN
+        requests_put.side_effect = worker_mod.requests.Timeout("Error")
+
+        # WHEN
+        result = worker._get_ec2_metadata_imdsv2_token()
+
+        # THEN
+        assert result is None
+
     def test_get_imdsv2_token_imds_inactive(self, worker: Worker, requests_put: MagicMock) -> None:
         # GIVEN
         response_mock = MagicMock()
@@ -463,6 +477,7 @@ def test_spot_shutdown(
         requests_get.assert_called_once_with(
             "http://169.254.169.254/latest/meta-data/spot/instance-action",
             headers={"X-aws-ec2-metadata-token": fake_token},
+            timeout=Worker._IMDS_REQUEST_TIMEOUT_SECONDS,
         )
         if not is_interrupt:
             assert result is None
@@ -532,6 +547,16 @@ def test_spot_shutdown_cannot_connect(self, worker: Worker, requests_get: MagicM
         # THEN
         assert result is None
 
+    def test_spot_shutdown_request_times_out(self, worker: Worker, requests_get: MagicMock) -> None:
+        # GIVEN
+        requests_get.side_effect = worker_mod.requests.Timeout("Error")
+
+        # WHEN
+        result = worker._get_spot_instance_shutdown_action_timeout(imdsv2_token="token")
+
+        # THEN
+        assert result is None
+
     def test_spot_shutdown_imds_inactive(self, worker: Worker, requests_get: MagicMock) -> None:
         # GIVEN
         response_mock = MagicMock()
@@ -575,6 +600,7 @@ def test_asg_terminate(
         requests_get.assert_called_once_with(
             "http://169.254.169.254/latest/meta-data/autoscaling/target-lifecycle-state",
             headers={"X-aws-ec2-metadata-token": fake_token},
+            timeout=Worker._IMDS_REQUEST_TIMEOUT_SECONDS,
         )
 
     def test_asg_terminate_cannot_connect(self, worker: Worker, requests_get: MagicMock) -> None:
@@ -587,6 +613,16 @@ def test_asg_terminate_cannot_connect(self, worker: Worker, requests_get: MagicM
         # THEN
         assert not result
 
+    def test_asg_terminate_request_times_out(self, worker: Worker, requests_get: MagicMock) -> None:
+        # GIVEN
+        requests_get.side_effect = worker_mod.requests.Timeout("Error")
+
+        # WHEN
+        result = worker._is_asg_terminated(imdsv2_token="token")
+
+        # THEN
+        assert not result
+
     def test_asg_terminate_imds_inactive(self, worker: Worker, requests_get: MagicMock) -> None:
         # GIVEN
         response_mock = MagicMock()