litellm/tests/integration/_support/redis_process.py

import os
import shutil
import signal
import subprocess
import threading
import time
import uuid
from collections.abc import Callable, Iterator, Mapping
from contextlib import contextmanager
from dataclasses import dataclass
from pathlib import Path
from queue import SimpleQueue
from typing import Final, TextIO

from pydantic import JsonValue, TypeAdapter
from redis import Redis
from redis.exceptions import ConnectionError as RedisConnectionError

from tests.integration._support.ports import allocate_port

_MONITOR_RECORD: Final = TypeAdapter(dict[str, JsonValue])


@dataclass
class OwnedRedis:
    host: str
    port: int
    command: tuple[str, ...]
    log: TextIO
    pid_file: str
    process: subprocess.Popen | None = None
    server_pid: int | None = None

    def serves(self) -> bool:
        with Redis(host=self.host, port=self.port, socket_connect_timeout=0.2, socket_timeout=0.2) as client:
            try:
                return bool(client.ping())
            except RedisConnectionError:
                return False

    def owned_pid(self) -> int | None:
        assert self.process is not None
        if self.command[0] != "docker":
            return self.process.pid
        result: Final = subprocess.run(
            ["docker", "exec", "redis-cache", "cat", self.pid_file],
            capture_output=True,
            text=True,
            timeout=2,
        )
        if result.returncode != 0:
            return None
        return int(result.stdout)

    def start(self) -> None:
        assert self.process is None
        assert not self.serves(), f"Another Redis already serves {self.host}:{self.port} before the owned one starts"
        process: Final = subprocess.Popen(
            self.command, stdout=self.log, stderr=subprocess.STDOUT, start_new_session=True
        )
        self.process = process
        deadline: Final = time.monotonic() + 8
        with Redis(host=self.host, port=self.port, socket_connect_timeout=0.2, socket_timeout=0.2) as client:
            while True:
                code: Final = process.poll()
                if code is not None:
                    self.process = None
                    raise AssertionError(
                        f"Owned Redis exited with code {code} before readiness on {self.host}:{self.port}"
                    )
                try:
                    if client.ping():
                        actual: Final = int(client.info("server")["process_id"])
                        expected: Final = self.owned_pid()
                        if actual == expected:
                            self.server_pid = actual
                            return
                        try:
                            process.wait(timeout=max(deadline - time.monotonic(), 0))
                        except subprocess.TimeoutExpired as error:
                            raise AssertionError("Redis readiness reached a different process") from error
                        self.process = None
                        raise AssertionError(f"A different Redis answered on {self.host}:{self.port} during startup")
                except RedisConnectionError:
                    pass
                assert time.monotonic() < deadline, "Owned Redis readiness deadline exceeded"
                time.sleep(0.05)

    def stop(self) -> None:
        assert self.process is not None
        failure = None
        forced = False
        try:
            if self.process.poll() is None:
                with Redis(host=self.host, port=self.port, socket_connect_timeout=1, socket_timeout=1) as client:
                    assert int(client.info("server")["process_id"]) == self.server_pid, (
                        "Redis ownership changed before shutdown"
                    )
                    client.shutdown(nosave=True)
        except Exception as error:
            failure = error
        finally:
            try:
                self.process.wait(timeout=5)
            except subprocess.TimeoutExpired:
                forced = True
                self.signal(signal.SIGTERM)
                try:
                    self.process.wait(timeout=5)
                except subprocess.TimeoutExpired:
                    self.signal(signal.SIGKILL)
                    self.process.wait(timeout=3)
        self.process = None
        self.server_pid = None
        assert not self.serves(), "Owned Redis still serves after shutdown"
        assert failure is None and not forced, f"Owned Redis required shutdown recovery: {failure!r}"

    def signal(self, action: signal.Signals) -> None:
        assert self.process is not None
        if self.command[0] != "docker":
            self.process.send_signal(action)
            return
        pid: Final = int(subprocess.check_output(["docker", "exec", "redis-cache", "cat", self.pid_file], timeout=2))
        command: Final = subprocess.check_output(
            ["docker", "exec", "redis-cache", "cat", f"/proc/{pid}/cmdline"], timeout=2
        )
        assert self.pid_file.encode() in command, "Redis process ownership changed"
        subprocess.run(["docker", "exec", "redis-cache", "kill", f"-{int(action)}", str(pid)], check=True, timeout=2)

    @contextmanager
    def paused(self) -> Iterator[None]:
        self.signal(signal.SIGSTOP)
        try:
            yield
        finally:
            self.signal(signal.SIGCONT)


@contextmanager
def owned_redis(
    directory: Path,
    choose_port: Callable[[], int] = allocate_port,
    inherited: Mapping[str, str] = os.environ,
) -> Iterator[OwnedRedis]:
    binary: Final = shutil.which("redis-server")
    if binary:
        host: Final = "127.0.0.1"
        prefix: Final = (binary,)
    else:
        host: Final = subprocess.check_output(
            [
                "docker",
                "inspect",
                "--format",
                "{{range .NetworkSettings.Networks}}{{.IPAddress}}{{end}}",
                "redis-cache",
            ],
            text=True,
        ).strip()
        assert host, "CircleCI owned Redis container has no address"
        prefix: Final = ("docker", "exec", "redis-cache", "redis-server")
    output: Final = Path(inherited.get("INTEGRATION_RESULTS_DIR", str(directory)))
    output.mkdir(parents=True, exist_ok=True)
    directory.mkdir(parents=True, exist_ok=True)
    with (output / f"owned-redis-{uuid.uuid4().hex}.log").open("w") as log:
        port: Final = choose_port()
        pid_file: Final = (
            str(directory / f"owned-redis-{port}.pid") if binary else f"/tmp/integration-redis-{uuid.uuid4().hex}.pid"
        )
        command: Final = (
            *prefix,
            "--port",
            str(port),
            "--set-proc-title",
            "no",
            "--pidfile",
            pid_file,
            "--bind",
            "0.0.0.0" if not binary else "127.0.0.1",
            "--protected-mode",
            "no",
            "--save",
            "",
            "--appendonly",
            "no",
        )
        server: Final = OwnedRedis(host, port, command, log, pid_file)
        try:
            server.start()
            yield server
        finally:
            if server.process is not None:
                server.stop()


@dataclass(frozen=True, slots=True)
class RedisCommands:
    queue: SimpleQueue[str]

    def drain(self) -> tuple[tuple[str, ...], ...]:
        """Every command MONITOR saw since the last drain, split on spaces into its name and arguments"""
        lines: Final = tuple(self.queue.get_nowait() for _ in range(self.queue.qsize()))
        return tuple(tuple(line.split(" ")) for line in lines)


@contextmanager
def redis_commands(host: str, port: int) -> Iterator[RedisCommands]:
    """Records every command the Redis server at `host:port` executes while the block runs"""
    commands: Final = RedisCommands(SimpleQueue())
    started: Final = threading.Event()
    armed: Final = threading.Event()
    ready_marker: Final = f"monitor-ready-{uuid.uuid4().hex}"
    stop_marker: Final = f"monitor-stop-{uuid.uuid4().hex}"

    def capture() -> None:
        # No socket timeout: a timeout raised inside `listen()` ends the generator, so capture would stop at the
        # first quiet second. The stop marker sent on exit is what wakes and ends this read.
        with Redis(host=host, port=port, decode_responses=True) as client:
            with client.monitor() as monitor:
                started.set()
                for raw in monitor.listen():
                    record: Final = _MONITOR_RECORD.validate_python(raw)
                    command: Final = record.get("command")
                    if not isinstance(command, str):
                        continue
                    if stop_marker in command:
                        return
                    if ready_marker in command:
                        armed.set()
                        continue
                    commands.queue.put(command)

    thread: Final = threading.Thread(target=capture, daemon=True)
    thread.start()
    try:
        assert started.wait(timeout=5), "Redis MONITOR did not start"
        with Redis(host=host, port=port, socket_timeout=1, decode_responses=True) as client:
            client.set(ready_marker, "ready", ex=1)
        assert armed.wait(timeout=5), "Redis MONITOR did not capture its readiness command"
        yield commands
    finally:
        with Redis(host=host, port=port, socket_timeout=1, decode_responses=True) as client:
            client.set(stop_marker, "stop", ex=1)
        thread.join(timeout=5)
        assert not thread.is_alive(), "Redis MONITOR did not reach its stop command"