Repository: PrefectHQ/prefect

Issue:
Reconnect Docker Agent on daemon restart
## Description
When a Docker Agent is running and connected to a docker daemon it will raise a broken pipe error when the daemon restarts and a new flow run is found.



## Expected Behavior
Docker Agent either gracefully reconnects to the daemon or exits when the pipe is broken. I'm leaning on the side of reconnecting as all it takes is reinstantiating the Docker low level API client




Discussion:
Yea I like the idea of pinging with each query and then attempting to reconnect if a ping fails; if the reconnect fails I think the Agent should actually exit and / or "deregister" itself so that it will be clear this Agent is not functional.

Resolving patch:
diff --git a/src/prefect/agent/docker/agent.py b/src/prefect/agent/docker/agent.py
--- a/src/prefect/agent/docker/agent.py
+++ b/src/prefect/agent/docker/agent.py
@@ -55,6 +55,7 @@ def __init__(
         self.no_pull = no_pull or context.get("no_pull", False)
         self.logger.debug("no_pull set to {}".format(self.no_pull))
 
+        self.failed_connections = 0
         self.docker_client = docker.APIClient(base_url=self.base_url, version="auto")
 
         # Ping Docker daemon for connection issues
@@ -67,6 +68,23 @@ def __init__(
             )
             raise exc
 
+    def heartbeat(self) -> None:
+        try:
+            if not self.docker_client.ping():
+                raise RuntimeError("Unexpected Docker ping result")
+            if self.failed_connections > 0:
+                self.logger.info("Reconnected to Docker daemon")
+            self.failed_connections = 0
+        except Exception as exc:
+            self.logger.warning("Failed heartbeat: {}".format(repr(exc)))
+            self.failed_connections += 1
+
+        if self.failed_connections >= 6:
+            self.logger.error(
+                "Cannot reconnect to Docker daemon. Agent is shutting down."
+            )
+            self.is_running = False
+
     def deploy_flow(self, flow_run: GraphQLResult) -> str:
         """
         Deploy flow runs on your local machine as Docker containers


Test patch:
diff --git a/tests/agent/test_docker_agent.py b/tests/agent/test_docker_agent.py
--- a/tests/agent/test_docker_agent.py
+++ b/tests/agent/test_docker_agent.py
@@ -322,3 +322,73 @@ def test_docker_agent_deploy_flow_no_registry_does_not_pull(monkeypatch, runner_
     assert not api.pull.called
     assert api.create_container.called
     assert api.start.called
+
+
+def test_docker_agent_heartbeat_gocase(monkeypatch, runner_token):
+    api = MagicMock()
+    api.ping.return_value = True
+    monkeypatch.setattr(
+        "prefect.agent.docker.agent.docker.APIClient", MagicMock(return_value=api)
+    )
+
+    agent = DockerAgent()
+    agent.heartbeat()
+    assert api.ping.call_count == 2
+
+
+def test_docker_agent_heartbeat_exits_on_failure(monkeypatch, runner_token, caplog):
+    api = MagicMock()
+    api.ping.return_value = True
+    monkeypatch.setattr(
+        "prefect.agent.docker.agent.docker.APIClient", MagicMock(return_value=api)
+    )
+
+    agent = DockerAgent()
+    agent.is_running = True  # this is only set on start()
+    api.ping.return_value = False
+    agent.heartbeat()
+    agent.heartbeat()
+    agent.heartbeat()
+    agent.heartbeat()
+    agent.heartbeat()
+    assert agent.is_running
+    agent.heartbeat()
+    assert not agent.is_running
+    assert "Cannot reconnect to Docker daemon. Agent is shutting down." in caplog.text
+    assert api.ping.call_count == 7
+
+
+def test_docker_agent_heartbeat_logs_reconnect(monkeypatch, runner_token, caplog):
+    api = MagicMock()
+    api.ping.return_value = True
+    monkeypatch.setattr(
+        "prefect.agent.docker.agent.docker.APIClient", MagicMock(return_value=api)
+    )
+
+    agent = DockerAgent()
+    agent.is_running = True  # this is only set on start()
+    api.ping.return_value = False
+    agent.heartbeat()
+    agent.heartbeat()
+    api.ping.return_value = True
+    agent.heartbeat()
+    assert api.ping.call_count == 4
+    assert "Reconnected to Docker daemon" in caplog.text
+
+
+def test_docker_agent_heartbeat_resets_fail_count(monkeypatch, runner_token, caplog):
+    api = MagicMock()
+    api.ping.return_value = True
+    monkeypatch.setattr(
+        "prefect.agent.docker.agent.docker.APIClient", MagicMock(return_value=api)
+    )
+
+    agent = DockerAgent()
+    api.ping.return_value = False
+    agent.heartbeat()
+    agent.heartbeat()
+    assert agent.failed_connections == 2
+    api.ping.return_value = True
+    agent.heartbeat()
+    assert agent.failed_connections == 0
+    assert api.ping.call_count == 4