seaweedfs/test/s3vectors/conftest.py

import contextlib
import os
from pathlib import Path
import socket
import subprocess
import time
import urllib.error
import urllib.request

import boto3
from botocore.config import Config
import pytest

HERE = Path(__file__).resolve().parent


def combinations():
    volumes = os.environ.get("VOLUME_SERVER_IMPL", "go,rust").split(",")
    default_modes = "single,distributed" if (HERE / "../../rust/sw-vector-distributed/Cargo.toml").is_file() else "single"
    modes = os.environ.get("VECTOR_QUERY_MODE", default_modes).split(",")
    if not set(volumes) <= {"go", "rust"} or not set(modes) <= {"single", "distributed"}:
        raise ValueError("invalid VOLUME_SERVER_IMPL or VECTOR_QUERY_MODE")
    return [(volume, mode) for volume in volumes for mode in modes]


def binary(variable, name):
    path = Path(os.environ.get(variable, str(HERE / name))).resolve()
    if not path.is_file() or not os.access(path, os.X_OK):
        pytest.fail(f"Missing executable {path}; run make build or set {variable}")
    return str(path)


def wait_for(check, processes, message):
    deadline = time.monotonic() + 60
    while time.monotonic() < deadline:
        for name, process in processes.items():
            if process.poll() is not None:
                pytest.fail(f"{name} exited with {process.returncode}; {message}")
        try:
            if check():
                return
        except (OSError, urllib.error.URLError):
            pass
        time.sleep(0.2)
    pytest.fail(message)


def listening(port):
    with socket.create_connection(("127.0.0.1", port), timeout=0.5):
        return True


def http_ready(url):
    try:
        with urllib.request.urlopen(url, timeout=0.5):
            return True
    except urllib.error.HTTPError:
        return True


def stop(process):
    if process.poll() is None:
        process.terminate()
        try:
            process.wait(timeout=10)
        except subprocess.TimeoutExpired:
            process.kill()
            process.wait(timeout=10)


@pytest.fixture(scope="module", params=combinations(), ids=lambda p: "-".join(p))
def cluster(request, tmp_path_factory):
    yield from run_cluster(request.param, tmp_path_factory, identities=True)


@pytest.fixture(scope="module", params=combinations(), ids=lambda p: "-".join(p))
def open_cluster(request, tmp_path_factory):
    """A cluster with no S3 identities, and an engine given no credentials."""
    yield from run_cluster(request.param, tmp_path_factory, identities=False)


def run_cluster(param, tmp_path_factory, identities):
    volume, mode = param
    directory = tmp_path_factory.mktemp(f"vectors-{volume}-{mode}" + ("" if identities else "-open"))
    print(f"\nVector integration logs: {directory}")
    weed = binary("WEED_BINARY", "weed")
    engine = binary("VECTOR_BINARY", "weed-vector")
    worker = binary("VECTOR_WORKER_BINARY", "weed-vector-worker") if mode == "distributed" else None
    volume_binary = binary("RUST_VOLUME_SERVER_BINARY", "weed-volume") if volume == "rust" else weed
    processes = {}
    with contextlib.ExitStack() as stack:
        with contextlib.ExitStack() as reservations:
            ports = []
            for _ in range(12):
                sock = reservations.enter_context(socket.socket())
                sock.bind(("127.0.0.1", 0))
                ports.append(sock.getsockname()[1])
        master, master_grpc, vol, vol_grpc, filer, filer_grpc, s3, s3_grpc, lance, vector, worker1, worker2 = ports
        master_addr = f"127.0.0.1:{master}.{master_grpc}"
        filer_addr = f"127.0.0.1:{filer}.{filer_grpc}"

        def start(name, args, unset=()):
            log = stack.enter_context((directory / f"{name}.log").open("w"))
            env = dict(os.environ, AWS_ACCESS_KEY_ID="test", AWS_SECRET_ACCESS_KEY="test",
                       AWS_DEFAULT_REGION="us-east-1", AWS_EC2_METADATA_DISABLED="true",
                       WEED_S3_SSE_KEY="s3vectors-test")
            env.pop("WEED_S3_SSE_KEK", None)
            for key in unset:
                env.pop(key)
            process = subprocess.Popen(args, cwd=directory, stdout=log, stderr=log, env=env)
            processes[name] = process
            stack.callback(stop, process)
            return process

        for name in ("master", "volume", "filer"):
            (directory / name).mkdir()
        start("master", [weed, "-v=1", "master", "-ip=127.0.0.1", "-ip.bind=127.0.0.1",
                         f"-port={master}", f"-port.grpc={master_grpc}",
                         f"-mdir={directory / 'master'}", "-peers=none", "-volumeSizeLimitMB=32"])
        wait_for(lambda: http_ready(f"http://127.0.0.1:{master}/dir/status"), processes, "master not ready")
        prefix = [volume_binary] if volume == "rust" else [weed, "volume"]
        flag = "--" if volume == "rust" else "-"
        start("volume", prefix + [f"{flag}{key}={value}" for key, value in {
            "master": master_addr, "ip": "127.0.0.1", "ip.bind": "127.0.0.1",
            "port": vol, "port.grpc": vol_grpc, "dir": directory / "volume",
            "max": 16, "preStopSeconds": 0,
        }.items()])
        wait_for(lambda: listening(vol_grpc), processes, "volume not ready")
        start("filer", [weed, "filer", f"-master={master_addr}", "-ip=127.0.0.1",
                        "-ip.bind=127.0.0.1", f"-port={filer}", f"-port.grpc={filer_grpc}",
                        f"-defaultStoreDir={directory / 'filer'}"])
        wait_for(lambda: listening(filer_grpc), processes, "filer not ready")
        endpoint = f"http://127.0.0.1:{s3}"
        # weed s3 turns AWS_ACCESS_KEY_ID into an admin identity, and the engine
        # would find it as a credential.
        if identities:
            config, unset = [f"-config={HERE / 's3.json'}"], ()
            credentials = ["--storage-option", "aws_access_key_id=test",
                           "--storage-option", "aws_secret_access_key=test"]
        else:
            config, unset, credentials = [], ("AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY"), []
        start("s3", [weed, "-v=1", "s3", f"-filer={filer_addr}", *config,
                     "-ip=127.0.0.1", "-ip.bind=127.0.0.1", f"-port={s3}",
                     f"-port.grpc={s3_grpc}", "-port.iceberg=0", f"-port.lance={lance}"], unset)
        wait_for(lambda: http_ready(f"http://127.0.0.1:{lance}/v1/namespace/root/list"), processes, "catalog not ready")
        start("vector", [engine, "--namespace", f"http://127.0.0.1:{lance}",
                         "--masters", master_addr, "--listen", f"127.0.0.1:{vector}", *credentials], unset)
        wait_for(lambda: listening(vector), processes, "vector engine not ready")

        wait_for(
            lambda: f"Cluster: vector node 127.0.0.1:{vector} added to group ''"
            in (directory / "master.log").read_text(),
            processes,
            "vector engine never registered with the master",
        )

        if worker:
            for i, port in enumerate((worker1, worker2)):
                start(f"worker-{i}", [worker, "--masters", master_addr, "--listen", f"127.0.0.1:{port}"])
                wait_for(lambda: listening(port), processes, "vector worker not ready")
            wait_for(
                lambda: all(
                    f"S3 Vectors: discovered search worker 127.0.0.1:{port}"
                    in (directory / "s3.log").read_text()
                    for port in (worker1, worker2)
                ),
                processes,
                "gateway did not discover both search workers",
            )
        yield {"endpoint": endpoint, "mode": mode, "processes": processes, "directory": directory}


@pytest.fixture(scope="module")
def vectors(cluster):
    return boto3.client("s3vectors", endpoint_url=cluster["endpoint"],
                        aws_access_key_id="test", aws_secret_access_key="test", region_name="us-east-1",
                        config=Config(connect_timeout=3, read_timeout=20, retries={"max_attempts": 0}))


@pytest.fixture(scope="module")
def vector_clients(vectors):
    return {
        name: boto3.client("s3vectors", endpoint_url=vectors.meta.endpoint_url,
                           aws_access_key_id=name, aws_secret_access_key=name + "-secret",
                           region_name="us-east-1", config=vectors.meta.config)
        for name in ("reader", "writer")
    }