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

core/testcontainers/core/container.py

@@ -2,6 +2,8 @@
 import os
 from typing import Iterable, Optional, Tuple
 
+
+from .reaper import Reaper, REAPER_IMAGE
 from .waiting_utils import wait_container_is_ready
 from .docker_client import DockerClient
 from .exceptions import ContainerStartException
@@ -26,6 +28,7 @@ def __init__(self, image: str, docker_client_kw: Optional[dict] = None, **kwargs
         self.env = {}
         self.ports = {}
         self.volumes = {}
+        self.ryuk = False
         self.image = image
         self._docker = DockerClient(**(docker_client_kw or {}))
         self._container = None
@@ -46,6 +49,10 @@ def with_exposed_ports(self, *ports: Iterable[int]) -> 'DockerContainer':
             self.ports[port] = None
         return self
 
+    def with_ryuk(self, enabled: bool) -> 'DockerContainer':
+        self.ryuk = enabled
+        return self
+
     def with_kwargs(self, **kwargs) -> 'DockerContainer':
         self._kwargs = kwargs
         return self
@@ -56,11 +63,14 @@ def maybe_emulate_amd64(self) -> 'DockerContainer':
         return self
 
     def start(self) -> 'DockerContainer':
+        if self.ryuk and not self.image == REAPER_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(
             self.image, command=self._command, detach=True, environment=self.env, ports=self.ports,
-            name=self._name, volumes=self.volumes, **self._kwargs
+            name=self._name, volumes=self.volumes, cleanup_on_exit=not self.ryuk, **self._kwargs
         )
         logger.info("Container started: %s", self._container.short_id)
         return self
@@ -76,9 +86,9 @@ def __exit__(self, exc_type, exc_val, exc_tb) -> None:
 
     def __del__(self) -> None:
         """
-        Try to remove the container in all circumstances
+        Try to remove the container if Ryuk is not active
         """
-        if self._container is not None:
+        if self._container is not None and not self.ryuk:
             try:
                 self.stop()
             except:  # noqa: E722

core/testcontainers/core/docker_client.py

@@ -19,6 +19,8 @@
 from typing import List, Optional, Union
 import urllib
 
+
+from .labels import create_labels, SESSION_ID
 from .utils import default_gateway_ip, inside_container, setup_logger
 
 
@@ -39,19 +41,22 @@ class DockerClient:
     """
     Thin wrapper around :class:`docker.DockerClient` for a more functional interface.
     """
+
     def __init__(self, **kwargs) -> None:
         self.client = docker.from_env(**kwargs)
+        self.client.api.headers["x-tc-sid"] = SESSION_ID
 
     @ft.wraps(ContainerCollection.run)
     def run(self, image: str, command: Union[str, List[str]] = None,
             environment: Optional[dict] = None, ports: Optional[dict] = None,
-            detach: bool = False, stdout: bool = True, stderr: bool = False, remove: bool = False,
+            labels: Optional[dict] = None, detach: bool = False, stdout: bool = True,
+            stderr: bool = False, remove: bool = False, cleanup_on_exit: bool = True,
             **kwargs) -> Container:
         container = self.client.containers.run(
             image, command=command, stdout=stdout, stderr=stderr, remove=remove, detach=detach,
-            environment=environment, ports=ports, **kwargs
+            environment=environment, ports=ports, labels=create_labels(image, labels), **kwargs
         )
-        if detach:
+        if detach and cleanup_on_exit:
             atexit.register(_stop_container, container)
         return container
 

core/testcontainers/core/images.py

@@ -0,0 +1,4 @@
+from os import environ
+
+
+REAPER_IMAGE = environ.get("RYUK_CONTAINER_IMAGE", "testcontainers/ryuk:0.3.4")

core/testcontainers/core/labels.py

@@ -0,0 +1,19 @@
+from uuid import uuid4
+from typing import Optional
+
+from .images import REAPER_IMAGE
+
+
+SESSION_ID: str = str(uuid4())
+LABEL_SESSION_ID = "org.testcontainers.session-id"
+
+
+def create_labels(image: str, labels: Optional[dict]) -> dict:
+    if labels is None:
+        labels = {}
+
+    if image == REAPER_IMAGE:
+        return labels
+
+    labels[LABEL_SESSION_ID] = SESSION_ID
+    return labels

core/testcontainers/core/reaper.py

@@ -0,0 +1,71 @@
+from os import environ
+from socket import socket
+from typing import TYPE_CHECKING, Optional
+
+
+from .utils import setup_logger
+from .images import REAPER_IMAGE
+from .waiting_utils import wait_for_logs
+from .labels import LABEL_SESSION_ID, SESSION_ID
+
+if TYPE_CHECKING:
+    from .container import DockerContainer
+
+
+logger = setup_logger(__name__)
+
+
+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":
+        from .container import DockerContainer
+
+        docker_socket = environ.get(
+            "TESTCONTAINERS_DOCKER_SOCKET_OVERRIDE", "/var/run/docker.sock"
+        )
+        logger.debug(f"Creating new Reaper for session: {SESSION_ID}")
+
+        Reaper._container = (
+            DockerContainer(REAPER_IMAGE)
+            .with_ryuk(True)
+            .with_name(f"testcontainers-ryuk-{SESSION_ID}")
+            .with_exposed_ports(8080)
+            .with_volume_mapping(docker_socket, "/var/run/docker.sock", "rw")
+            .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/tests/test_ryuk.py

@@ -0,0 +1,21 @@
+from testcontainers.core.reaper 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").with_ryuk(True).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 1 .*", timeout=15)
+
+    Reaper.delete_instance()
+
+
+def test_container_without_ryuk():
+    container = DockerContainer("hello-world").start()
+    wait_for_logs(container, "Hello from Docker!")
+    assert Reaper._instance is None