From 4568437487df8b0cf11c160d7751afb8f865907d Mon Sep 17 00:00:00 2001 From: Anas Khan <83116240+anxkhn@users.noreply.github.com> Date: Sat, 4 Jul 2026 15:46:19 +0530 Subject: [PATCH] fix(code_executors): stop ContainerCodeExecutor leaking via atexit ContainerCodeExecutor.__init__ registered the bound method self.__cleanup_container with atexit. A bound method keeps a strong reference to the instance, so the process-global atexit registry retained every ContainerCodeExecutor ever created for the whole interpreter lifetime, and its long-lived Docker container was only stopped at exit rather than when the executor was discarded. A server that builds an executor per app/agent/session would grow both the atexit registry and the number of running containers without bound. Register the exit handler with a weakref.proxy instead, mirroring the existing idiom in bigquery_agent_analytics_plugin.py, so the handler no longer keeps the executor alive; a discarded executor (and its container reference) can be garbage collected. __cleanup_container becomes a staticmethod that guards ReferenceError for the case where the proxy is already dead at exit. Add regression tests asserting a dropped executor is not retained and that cleanup still stops and removes the container. Signed-off-by: Anas Khan <83116240+anxkhn@users.noreply.github.com> --- .../code_executors/container_code_executor.py | 35 ++++++++++++---- .../test_container_code_executor.py | 42 ++++++++++++++++++- 2 files changed, 68 insertions(+), 9 deletions(-) diff --git a/src/google/adk/code_executors/container_code_executor.py b/src/google/adk/code_executors/container_code_executor.py index 671c33ac34..6d2ec11f3b 100644 --- a/src/google/adk/code_executors/container_code_executor.py +++ b/src/google/adk/code_executors/container_code_executor.py @@ -18,6 +18,7 @@ import logging import os from typing import Optional +import weakref import docker from docker.client import DockerClient @@ -139,8 +140,15 @@ def __init__( # Initialize the container. self.__init_container() - # Close the container when the on exit. - atexit.register(self.__cleanup_container) + # Close the container on interpreter exit. Register a weakref.proxy so the + # process-global atexit registry does not keep a strong reference to this + # executor: otherwise every ContainerCodeExecutor ever created (and its + # long-lived Docker container) would be retained until interpreter exit, + # even after the executor is no longer used. A server that builds executors + # per app/agent/session would leak them unboundedly. + atexit.register( + ContainerCodeExecutor.__cleanup_container, weakref.proxy(self) + ) @override def execute_code( @@ -219,12 +227,23 @@ def __init_container(self): # Verify the container is able to run python3. self._verify_python_installation() - def __cleanup_container(self): - """Closes the container on exit.""" - if not self._container: + @staticmethod + def __cleanup_container(self: ContainerCodeExecutor): + """Closes the container on exit. + + Registered with ``atexit`` via a ``weakref.proxy`` so it does not keep the + executor alive. If the executor has already been garbage collected the + proxy is dead and there is nothing left to clean up. + """ + try: + container = self._container + except ReferenceError: + return + + if not container: return logger.info('[Cleanup] Stopping the container...') - self._container.stop() - self._container.remove() - logger.info('Container %s stopped and removed.', self._container.id) + container.stop() + container.remove() + logger.info('Container %s stopped and removed.', container.id) diff --git a/tests/unittests/code_executors/test_container_code_executor.py b/tests/unittests/code_executors/test_container_code_executor.py index 5574135b52..3681457e52 100644 --- a/tests/unittests/code_executors/test_container_code_executor.py +++ b/tests/unittests/code_executors/test_container_code_executor.py @@ -12,9 +12,11 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Tests for the ContainerCodeExecutor container hardening defaults.""" +"""Tests for the ContainerCodeExecutor container hardening and lifecycle.""" +import gc from unittest import mock +import weakref from google.adk.code_executors.container_code_executor import ContainerCodeExecutor @@ -56,3 +58,41 @@ def test_container_network_can_be_explicitly_enabled(mock_docker): _, kwargs = client.containers.run.call_args assert not kwargs['network_disabled'] + + +@mock.patch('google.adk.code_executors.container_code_executor.docker') +def test_executor_is_not_retained_by_atexit(mock_docker): + """The atexit handler must not keep the executor (or its container) alive. + + The exit handler is registered via a ``weakref.proxy`` so that a discarded + executor can be garbage collected instead of surviving (together with its + running Docker container) until interpreter exit. Registering a bound method + would keep a strong reference in the process-global atexit registry and leak + every executor ever created. + """ + client = _mock_docker_client() + mock_docker.from_env.return_value = client + + executor = ContainerCodeExecutor(image='test-image') + ref = weakref.ref(executor) + + del executor + gc.collect() + + assert ref() is None + + +@mock.patch('google.adk.code_executors.container_code_executor.docker') +def test_cleanup_stops_and_removes_container(mock_docker): + """The exit handler still stops and removes a live executor's container.""" + client = _mock_docker_client() + container = client.containers.run.return_value + mock_docker.from_env.return_value = client + + executor = ContainerCodeExecutor(image='test-image') + # Name-mangled staticmethod registered with atexit; invoke it directly to + # emulate interpreter exit while the executor is still referenced. + ContainerCodeExecutor._ContainerCodeExecutor__cleanup_container(executor) + + container.stop.assert_called_once() + container.remove.assert_called_once()