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

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.3.4")
+RYUK_PRIVILEGED: bool = environ.get("TESTCONTAINERS_RYUK_PRIVILEGED", "false") == "true"
+RYUK_DOCKER_SOCKET: str = environ.get("TESTCONTAINERS_DOCKER_SOCKET_OVERRIDE",
+                                      "/var/run/docker.sock")

core/testcontainers/core/container.py

@@ -2,6 +2,9 @@
 import os
 from typing import Iterable, Optional, Tuple
 
+
+from .reaper import Reaper
+from .config import RYUK_IMAGE
 from .waiting_utils import wait_container_is_ready
 from .docker_client import DockerClient
 from .exceptions import ContainerStartException
@@ -26,6 +29,7 @@ def __init__(self, image: str, docker_client_kw: Optional[dict] = None, **kwargs
         self.env = {}
         self.ports = {}
         self.volumes = {}
+        self.auto_remove = True
         self.image = image
         self._docker = DockerClient(**(docker_client_kw or {}))
         self._container = None
@@ -46,6 +50,10 @@ def with_exposed_ports(self, *ports: Iterable[int]) -> 'DockerContainer':
             self.ports[port] = None
         return self
 
+    def with_auto_remove(self, enabled: bool) -> 'DockerContainer':
+        self.auto_remove = enabled
+        return self
+
     def with_kwargs(self, **kwargs) -> 'DockerContainer':
         self._kwargs = kwargs
         return self
@@ -56,6 +64,9 @@ def maybe_emulate_amd64(self) -> 'DockerContainer':
         return self
 
     def start(self) -> 'DockerContainer':
+        if self.auto_remove 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(
@@ -74,16 +85,6 @@ def __enter__(self) -> 'DockerContainer':
     def __exit__(self, exc_type, exc_val, exc_tb) -> None:
         self.stop()
 
-    def __del__(self) -> None:
-        """
-        Try to remove the container in all circumstances
-        """
-        if self._container is not None:
-            try:
-                self.stop()
-            except:  # noqa: E722
-                pass
-
     def get_container_host_ip(self) -> str:
         # infer from docker host
         host = self.get_docker_client().host()

core/testcontainers/core/docker_client.py

@@ -10,49 +10,40 @@
 #    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 docker
-from docker.errors import NotFound
 from docker.models.containers import Container, ContainerCollection
 import functools as ft
 import os
 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
 
 
 LOGGER = setup_logger(__name__)
 
 
-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)
-
-
 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,
             **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:
-            atexit.register(_stop_container, container)
         return container
 
     def port(self, container_id: str, port: int) -> int:

core/testcontainers/core/labels.py

@@ -0,0 +1,21 @@
+from uuid import uuid4
+from typing import Optional
+
+from .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]) -> dict:
+    if labels is None:
+        labels = {}
+    labels[LABEL_LANG] = "python"
+
+    if image == RYUK_IMAGE:
+        return labels
+
+    labels[LABEL_SESSION_ID] = SESSION_ID
+    return labels

core/testcontainers/core/reaper.py

@@ -0,0 +1,69 @@
+from __future__ import annotations
+
+from socket import socket
+from typing import TYPE_CHECKING, Optional
+
+
+from .utils import setup_logger
+from .config import RYUK_IMAGE, RYUK_DOCKER_SOCKET, RYUK_PRIVILEGED
+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
+
+        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/tests/test_ryuk.py

@@ -0,0 +1,22 @@
+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").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():
+    with DockerContainer("hello-world").with_auto_remove(False) as container:
+        wait_for_logs(container, "Hello from Docker!")
+        assert Reaper._instance is None