feat(reliability): integrate the ryuk container for better container cleanup

INDEX.rst

@@ -92,6 +92,21 @@ When trying to launch a testcontainer from within a Docker container, e.g., in c
 1. The container has to provide a docker client installation. Either use an image that has docker pre-installed (e.g. the `official docker images <https://hub.docker.com/_/docker>`_) or install the client from within the `Dockerfile` specification.
 2. The container has to have access to the docker daemon which can be achieved by mounting `/var/run/docker.sock` or setting the `DOCKER_HOST` environment variable as part of your `docker run` command.
 
+Configuration
+-------------
+
++-------------------------------------------+-------------------------------+------------------------------------------+
+| Env Variable                              | Example                       | Description                              |
++===========================================+===============================+==========================================+
+| ``TESTCONTAINERS_DOCKER_SOCKET_OVERRIDE`` | ``/var/run/docker.sock``      | Path to Docker's socket used by ryuk     |
++-------------------------------------------+-------------------------------+------------------------------------------+
+| ``TESTCONTAINERS_RYUK_PRIVILEGED``        | ``false``                     | Run ryuk as a privileged container       |
++-------------------------------------------+-------------------------------+------------------------------------------+
+| ``TESTCONTAINERS_RYUK_DISABLED``          | ``false``                     | Disable ryuk                             |
++-------------------------------------------+-------------------------------+------------------------------------------+
+| ``RYUK_CONTAINER_IMAGE``                  | ``testcontainers/ryuk:0.5.1`` | Custom image for ryuk                    |
++-------------------------------------------+-------------------------------+------------------------------------------+
+
 Development and Contributing
 ----------------------------
 

README.md

@@ -22,3 +22,12 @@ For more information, see [the docs][readthedocs].
 ```
 
 The snippet above will spin up a postgres database in a container. The `get_connection_url()` convenience method returns a `sqlalchemy` compatible url we use to connect to the database and retrieve the database version.
+
+## Configuration
+
+| Env Variable                              | Example                       | Description                              |
+| ----------------------------------------- | ----------------------------- | ---------------------------------------- |
+| `TESTCONTAINERS_DOCKER_SOCKET_OVERRIDE`   | `/var/run/docker.sock`        | Path to Docker's socket used by ryuk     |
+| `TESTCONTAINERS_RYUK_PRIVILEGED`          | `false`                       | Run ryuk as a privileged container       |
+| `TESTCONTAINERS_RYUK_DISABLED`            | `false`                       | Disable ryuk                             |
+| `RYUK_CONTAINER_IMAGE`                    | `testcontainers/ryuk:0.5.1`   | Custom image for ryuk                    |

core/testcontainers/core/config.py

@@ -3,3 +3,8 @@
 MAX_TRIES = int(environ.get("TC_MAX_TRIES", 120))
 SLEEP_TIME = int(environ.get("TC_POOLING_INTERVAL", 1))
 TIMEOUT = MAX_TRIES * SLEEP_TIME
+
+RYUK_IMAGE: str = environ.get("RYUK_CONTAINER_IMAGE", "testcontainers/ryuk:0.5.1")
+RYUK_PRIVILEGED: bool = environ.get("TESTCONTAINERS_RYUK_PRIVILEGED", "false") == "true"
+RYUK_DISABLED: bool = environ.get("TESTCONTAINERS_RYUK_DISABLED", "false") == "true"
+RYUK_DOCKER_SOCKET: str = environ.get("TESTCONTAINERS_DOCKER_SOCKET_OVERRIDE", "/var/run/docker.sock")

core/testcontainers/core/container.py

@@ -1,13 +1,16 @@
-import contextlib
 from platform import system
-from typing import Optional
-
-from docker.models.containers import Container
+from socket import socket
+from typing import TYPE_CHECKING, Optional
 
+from testcontainers.core.config import RYUK_DISABLED, RYUK_DOCKER_SOCKET, RYUK_IMAGE, RYUK_PRIVILEGED
 from testcontainers.core.docker_client import DockerClient
 from testcontainers.core.exceptions import ContainerStartException
+from testcontainers.core.labels import LABEL_SESSION_ID, SESSION_ID
 from testcontainers.core.utils import inside_container, is_arm, setup_logger
-from testcontainers.core.waiting_utils import wait_container_is_ready
+from testcontainers.core.waiting_utils import wait_container_is_ready, wait_for_logs
+
+if TYPE_CHECKING:
+    from docker.models.containers import Container
 
 logger = setup_logger(__name__)
 
@@ -25,7 +28,12 @@ class DockerContainer:
         ...    delay = wait_for_logs(container, "Hello from Docker!")
     """
 
-    def __init__(self, image: str, docker_client_kw: Optional[dict] = None, **kwargs) -> None:
+    def __init__(
+        self,
+        image: str,
+        docker_client_kw: Optional[dict] = None,
+        **kwargs,
+    ) -> None:
         self.env = {}
         self.ports = {}
         self.volumes = {}
@@ -58,7 +66,10 @@ def maybe_emulate_amd64(self) -> "DockerContainer":
             return self.with_kwargs(platform="linux/amd64")
         return self
 
-    def start(self) -> "DockerContainer":
+    def start(self):
+        if not RYUK_DISABLED and self.image != RYUK_IMAGE:
+            logger.debug("Creating Ryuk container")
+            Reaper.get_instance()
         logger.info("Pulling image %s", self.image)
         docker_client = self.get_docker_client()
         self._container = docker_client.run(
@@ -69,7 +80,7 @@ def start(self) -> "DockerContainer":
             ports=self.ports,
             name=self._name,
             volumes=self.volumes,
-            **self._kwargs
+            **self._kwargs,
         )
         logger.info("Container started: %s", self._container.short_id)
         return self
@@ -78,21 +89,12 @@ def stop(self, force=True, delete_volume=True) -> None:
         self._container.remove(force=force, v=delete_volume)
         self.get_docker_client().client.close()
 
-    def __enter__(self) -> "DockerContainer":
+    def __enter__(self):
         return self.start()
 
     def __exit__(self, exc_type, exc_val, exc_tb) -> None:
         self.stop()
 
-    def __del__(self) -> None:
-        """
-        __del__ runs when Python attempts to garbage collect the object.
-        In case of leaky test design, we still attempt to clean up the container.
-        """
-        with contextlib.suppress(Exception):
-            if self._container is not None:
-                self.stop()
-
     def get_container_host_ip(self) -> str:
         # infer from docker host
         host = self.get_docker_client().host()
@@ -140,7 +142,7 @@ def with_volume_mapping(self, host: str, container: str, mode: str = "ro") -> "D
         self.volumes[host] = mapping
         return self
 
-    def get_wrapped_container(self) -> Container:
+    def get_wrapped_container(self) -> "Container":
         return self._container
 
     def get_docker_client(self) -> DockerClient:
@@ -155,3 +157,54 @@ def exec(self, command) -> tuple[int, str]:
         if not self._container:
             raise ContainerStartException("Container should be started before executing a command")
         return self._container.exec_run(command)
+
+
+class Reaper:
+    _instance: "Optional[Reaper]" = None
+    _container: Optional[DockerContainer] = None
+    _socket: Optional[socket] = None
+
+    @classmethod
+    def get_instance(cls) -> "Reaper":
+        if not Reaper._instance:
+            Reaper._instance = Reaper._create_instance()
+
+        return Reaper._instance
+
+    @classmethod
+    def delete_instance(cls) -> None:
+        if Reaper._socket is not None:
+            Reaper._socket.close()
+            Reaper._socket = None
+
+        if Reaper._container is not None:
+            Reaper._container.stop()
+            Reaper._container = None
+
+        if Reaper._instance is not None:
+            Reaper._instance = None
+
+    @classmethod
+    def _create_instance(cls) -> "Reaper":
+        logger.debug(f"Creating new Reaper for session: {SESSION_ID}")
+
+        Reaper._container = (
+            DockerContainer(RYUK_IMAGE)
+            .with_name(f"testcontainers-ryuk-{SESSION_ID}")
+            .with_exposed_ports(8080)
+            .with_volume_mapping(RYUK_DOCKER_SOCKET, "/var/run/docker.sock", "rw")
+            .with_kwargs(privileged=RYUK_PRIVILEGED)
+            .start()
+        )
+        wait_for_logs(Reaper._container, r".* Started!")
+
+        container_host = Reaper._container.get_container_host_ip()
+        container_port = int(Reaper._container.get_exposed_port(8080))
+
+        Reaper._socket = socket()
+        Reaper._socket.connect((container_host, container_port))
+        Reaper._socket.send(f"label={LABEL_SESSION_ID}={SESSION_ID}\r\n".encode())
+
+        Reaper._instance = Reaper()
+
+        return Reaper._instance

core/testcontainers/core/docker_client.py

@@ -10,7 +10,6 @@
 #    WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
 #    License for the specific language governing permissions and limitations
 #    under the License.
-import atexit
 import functools as ft
 import os
 import urllib
@@ -19,10 +18,10 @@
 from typing import Optional, Union
 
 import docker
-from docker.errors import NotFound
 from docker.models.containers import Container, ContainerCollection
 
-from .utils import default_gateway_ip, inside_container, setup_logger
+from testcontainers.core.labels import SESSION_ID, create_labels
+from testcontainers.core.utils import default_gateway_ip, inside_container, setup_logger
 
 LOGGER = setup_logger(__name__)
 TC_FILE = ".testcontainers.properties"
@@ -42,6 +41,7 @@ def __init__(self, **kwargs) -> None:
             self.client = docker.DockerClient(base_url=docker_host)
         else:
             self.client = docker.from_env(**kwargs)
+        self.client.api.headers["x-tc-sid"] = SESSION_ID
 
     @ft.wraps(ContainerCollection.run)
     def run(
@@ -50,6 +50,7 @@ def run(
         command: Optional[Union[str, list[str]]] = None,
         environment: Optional[dict] = None,
         ports: Optional[dict] = None,
+        labels: Optional[dict[str, str]] = None,
         detach: bool = False,
         stdout: bool = True,
         stderr: bool = False,
@@ -65,10 +66,9 @@ def run(
             detach=detach,
             environment=environment,
             ports=ports,
+            labels=create_labels(image, labels),
             **kwargs,
         )
-        if detach:
-            atexit.register(_stop_container, container)
         return container
 
     def port(self, container_id: str, port: int) -> int:
@@ -145,12 +145,3 @@ def read_tc_properties() -> dict[str, str]:
             tuples = [line.split("=") for line in contents.readlines() if "=" in line]
             settings = {**settings, **{item[0]: item[1] for item in tuples}}
     return settings
-
-
-def _stop_container(container: Container) -> None:
-    try:
-        container.stop()
-    except NotFound:
-        pass
-    except Exception as ex:
-        LOGGER.warning("failed to shut down container %s with image %s: %s", container.id, container.image, ex)

core/testcontainers/core/labels.py

@@ -0,0 +1,20 @@
+from typing import Optional
+from uuid import uuid4
+
+from testcontainers.core.config import RYUK_IMAGE
+
+SESSION_ID: str = str(uuid4())
+LABEL_SESSION_ID = "org.testcontainers.session-id"
+LABEL_LANG = "org.testcontainers.lang"
+
+
+def create_labels(image: str, labels: Optional[dict[str, str]]) -> dict[str, str]:
+    if labels is None:
+        labels = {}
+    labels[LABEL_LANG] = "python"
+
+    if image == RYUK_IMAGE:
+        return labels
+
+    labels[LABEL_SESSION_ID] = SESSION_ID
+    return labels

core/tests/test_ryuk.py

@@ -0,0 +1,24 @@
+from testcontainers.core import container
+from testcontainers.core.container import Reaper
+from testcontainers.core.container import DockerContainer
+from testcontainers.core.waiting_utils import wait_for_logs
+
+
+def test_wait_for_reaper():
+    container = DockerContainer("hello-world").start()
+    wait_for_logs(container, "Hello from Docker!")
+
+    assert Reaper._socket is not None
+    Reaper._socket.close()
+
+    assert Reaper._container is not None
+    wait_for_logs(Reaper._container, r".* Removed \d .*", timeout=30)
+
+    Reaper.delete_instance()
+
+
+def test_container_without_ryuk(monkeypatch):
+    monkeypatch.setattr(container, "RYUK_DISABLED", True)
+    with DockerContainer("hello-world") as cont:
+        wait_for_logs(cont, "Hello from Docker!")
+        assert Reaper._instance is None

pyproject.toml

@@ -197,6 +197,9 @@ ignore = [
     "INP001"
 ]
 
+[tool.ruff.lint.pyupgrade]
+keep-runtime-typing = true
+
 [tool.ruff.lint.flake8-type-checking]
 strict = true