milvus-io--milvus/tests/python_client/cdc/conftest.py

import logging
import time
from concurrent.futures import ThreadPoolExecutor, wait

import pytest
from pymilvus import MilvusClient

# Configure logging
logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s")
logger = logging.getLogger(__name__)

CDC_UPDATE_REPLICATE_TIMEOUT_SECONDS = 600


def apply_replicate_configuration(tasks, timeout=CDC_UPDATE_REPLICATE_TIMEOUT_SECONDS):
    # Fan out in parallel: the server blocks non-primary clusters in
    # waitUntilPrimaryChangeOrConfigurationSame until the primary's broadcast
    # propagates via CDC, so a sequential call where the first client happens
    # to be a replica deadlocks on the client's RPC timeout.
    with ThreadPoolExecutor(max_workers=len(tasks)) as executor:
        futures = [
            executor.submit(client.update_replicate_configuration, timeout=timeout, **config)
            for client, config in tasks
        ]
        wait(futures)
    for f in futures:
        f.result()


def pytest_addoption(parser):
    """Add command line options for pytest."""
    parser.addoption(
        "--upstream-uri",
        action="store",
        default="http://10.104.17.154:19530",
        help="Upstream Milvus uri",
    )
    parser.addoption(
        "--upstream-token",
        action="store",
        default="root:Milvus",
        help="Upstream Milvus token",
    )
    parser.addoption(
        "--downstream-uri",
        action="store",
        default="http://10.104.17.156:19530",
        help="Downstream Milvus uri",
    )
    parser.addoption(
        "--downstream-token",
        action="store",
        default="root:Milvus",
        help="Downstream Milvus token",
    )
    parser.addoption("--sync-timeout", action="store", default="30", help="Sync timeout in seconds")
    parser.addoption(
        "--source-cluster-id",
        action="store",
        default="cdc-test-source-0930",
        help="Source cluster ID for CDC topology",
    )
    parser.addoption(
        "--target-cluster-id",
        action="store",
        default="cdc-test-target-0930",
        help="Target cluster ID for CDC topology",
    )
    parser.addoption(
        "--pchannel-num",
        action="store",
        default="16",
        help="Number of physical channels for CDC",
    )
    parser.addoption(
        "--request-duration",
        action="store",
        default="30m",
        help="Duration for test operations (e.g., 30m, 1h, 60s)",
    )
    parser.addoption(
        "--is-check",
        action="store",
        default="true",
        help="Whether to assert on checker statistics",
    )
    parser.addoption(
        "--milvus-ns",
        action="store",
        default="chaos-testing",
        help="Kubernetes namespace for Milvus deployment",
    )
    parser.addoption(
        "--import-2pc-workload",
        action="store",
        default="true",
        help="Whether CDC stability suites should include Import 2PC workload",
    )
    parser.addoption(
        "--import-2pc-minio-host",
        action="store",
        default="",
        help="Upstream MinIO host or host:port used by Import 2PC workload",
    )
    parser.addoption(
        "--import-2pc-minio-bucket",
        action="store",
        default="",
        help="Upstream MinIO bucket used by Import 2PC workload",
    )
    parser.addoption(
        "--import-2pc-downstream-minio-host",
        action="store",
        default="",
        help="Downstream MinIO host or host:port used by CDC Import 2PC workload",
    )
    parser.addoption(
        "--import-2pc-downstream-minio-bucket",
        action="store",
        default="",
        help="Downstream MinIO bucket used by CDC Import 2PC workload",
    )
    parser.addoption(
        "--import-2pc-rows",
        action="store",
        default="20",
        help="Rows per Import 2PC checker operation",
    )


def _as_bool(value):
    if isinstance(value, bool):
        return value
    return str(value).lower() in ("1", "true", "yes", "y")


def _minio_endpoint(host):
    host = (host or "").strip()
    if not host:
        return ""
    if ":" in host:
        return host
    return f"{host}:9000"


@pytest.fixture(scope="session")
def upstream_client(request):
    """Create upstream MilvusClient."""
    uri = request.config.getoption("--upstream-uri")
    token = request.config.getoption("--upstream-token")
    client = MilvusClient(uri=uri, token=token)
    yield client
    client.close()


@pytest.fixture(scope="session")
def downstream_client(request):
    """Create downstream MilvusClient."""
    uri = request.config.getoption("--downstream-uri")
    token = request.config.getoption("--downstream-token")
    client = MilvusClient(uri=uri, token=token)
    yield client
    client.close()


@pytest.fixture(scope="session")
def sync_timeout(request):
    """Get sync timeout from command line."""
    return int(request.config.getoption("--sync-timeout"))


@pytest.fixture(scope="session")
def upstream_uri(request):
    """Get upstream uri from command line."""
    return request.config.getoption("--upstream-uri")


@pytest.fixture(scope="session")
def upstream_token(request):
    """Get upstream token from command line."""
    return request.config.getoption("--upstream-token")


@pytest.fixture(scope="session")
def downstream_uri(request):
    """Get downstream uri from command line."""
    return request.config.getoption("--downstream-uri")


@pytest.fixture(scope="session")
def downstream_token(request):
    """Get downstream token from command line."""
    return request.config.getoption("--downstream-token")


@pytest.fixture(scope="session")
def source_cluster_id(request):
    """Get source cluster id from command line."""
    return request.config.getoption("--source-cluster-id")


@pytest.fixture(scope="session")
def target_cluster_id(request):
    """Get target cluster id from command line."""
    return request.config.getoption("--target-cluster-id")


@pytest.fixture(scope="session")
def pchannel_num(request):
    """Get pchannel num from command line."""
    return int(request.config.getoption("--pchannel-num"))


@pytest.fixture(scope="session")
def request_duration(request):
    """Get request duration from command line."""
    return request.config.getoption("--request-duration")


@pytest.fixture(scope="session")
def is_check(request):
    # The root tests/python_client/conftest.py registers --is_check (underscore)
    # with type=bool, which argparse maps to the same dest (is_check) as our
    # --is-check (hyphen). The root's bool wins in chaos runs, so accept either.
    val = request.config.getoption("--is-check")
    return val if isinstance(val, bool) else str(val).lower() == "true"


@pytest.fixture(scope="session")
def milvus_ns(request):
    return request.config.getoption("--milvus-ns")


@pytest.fixture(scope="session")
def import_2pc_workload(request):
    return _as_bool(request.config.getoption("--import-2pc-workload"))


@pytest.fixture(scope="session")
def import_2pc_minio_endpoint(request):
    host = request.config.getoption("--import-2pc-minio-host") or request.config.getoption("--minio_host")
    return _minio_endpoint(host)


@pytest.fixture(scope="session")
def import_2pc_minio_bucket(request):
    return request.config.getoption("--import-2pc-minio-bucket") or request.config.getoption("--minio_bucket")


@pytest.fixture(scope="session")
def import_2pc_downstream_minio_endpoint(request):
    host = (
        request.config.getoption("--import-2pc-downstream-minio-host")
        or request.config.getoption("--import-2pc-minio-host")
        or request.config.getoption("--minio_host")
    )
    return _minio_endpoint(host)


@pytest.fixture(scope="session")
def import_2pc_downstream_minio_bucket(request):
    return (
        request.config.getoption("--import-2pc-downstream-minio-bucket")
        or request.config.getoption("--import-2pc-minio-bucket")
        or request.config.getoption("--minio_bucket")
    )


@pytest.fixture(scope="session")
def import_2pc_rows(request):
    return int(request.config.getoption("--import-2pc-rows"))


@pytest.fixture(scope="session")
def switchover_helper(request, upstream_client, downstream_client):
    """Returns a callable that performs CDC topology switchover."""
    upstream_uri = request.config.getoption("--upstream-uri")
    upstream_token = request.config.getoption("--upstream-token")
    downstream_uri = request.config.getoption("--downstream-uri")
    downstream_token = request.config.getoption("--downstream-token")
    pchannel_num = int(request.config.getoption("--pchannel-num"))
    original_source = request.config.getoption("--source-cluster-id")
    original_target = request.config.getoption("--target-cluster-id")

    # Map cluster IDs to their URIs/tokens
    cluster_map = {
        original_source: {"uri": upstream_uri, "token": upstream_token},
        original_target: {"uri": downstream_uri, "token": downstream_token},
    }

    def do_switchover(new_source_id, new_target_id):
        logger.info(f"Performing switchover: {new_source_id} -> {new_target_id}")
        config = {
            "clusters": [
                {
                    "cluster_id": new_source_id,
                    "connection_param": cluster_map[new_source_id],
                    "pchannels": [f"{new_source_id}-rootcoord-dml_{i}" for i in range(pchannel_num)],
                },
                {
                    "cluster_id": new_target_id,
                    "connection_param": cluster_map[new_target_id],
                    "pchannels": [f"{new_target_id}-rootcoord-dml_{i}" for i in range(pchannel_num)],
                },
            ],
            "cross_cluster_topology": [{"source_cluster_id": new_source_id, "target_cluster_id": new_target_id}],
        }
        # Dedicated short-lived clients so switchover RPCs don't share a
        # gRPC channel with concurrent DML on the session-scoped clients.
        # pymilvus's connection manager closes a channel on UNAVAILABLE /
        # STREAMING_CODE_REPLICATE_VIOLATION to trigger recovery; if a DML
        # on the session client triggers that close while our sibling
        # update_replicate_configuration RPC is in flight on the same
        # channel, the latter surfaces "Cannot invoke RPC on closed
        # channel!". Separate clients = separate channels = no race.
        up_tmp = MilvusClient(uri=upstream_uri, token=upstream_token)
        dn_tmp = MilvusClient(uri=downstream_uri, token=downstream_token)
        try:
            apply_replicate_configuration([(up_tmp, config), (dn_tmp, config)])
        finally:
            up_tmp.close()
            dn_tmp.close()
        logger.info("Switchover completed, waiting 10s for stabilization...")
        time.sleep(10)

    return do_switchover


@pytest.fixture(scope="session")
def kubectl_helper(milvus_ns):
    """Helpers for pod-kill failover scenarios."""
    import subprocess as _sp
    import time as _time

    def delete_pods(instance_label):
        cmd = [
            "kubectl",
            "delete",
            "pods",
            "-l",
            f"app.kubernetes.io/instance={instance_label}",
            "-n",
            milvus_ns,
            "--grace-period=0",
            "--force",
        ]
        result = _sp.run(cmd, capture_output=True, text=True, check=False)
        logger.info(
            f"[KUBECTL] delete pods {instance_label}: rc={result.returncode}, "
            f"stdout={result.stdout!r}, stderr={result.stderr!r}"
        )
        return result

    def wait_for_pods_ready(instance_label, timeout=300):
        """Wait for pods matching the label to be Ready.

        Two phases because right after `kubectl delete pods`, the operator
        hasn't recreated pods yet — `kubectl wait` would error with "no
        matching resources found". So:
          1. Poll until at least one pod matches.
          2. Then `kubectl wait` for Ready on the remaining time budget.
        """
        deadline = _time.time() + timeout
        existence = None
        while _time.time() < deadline:
            existence = _sp.run(
                [
                    "kubectl",
                    "get",
                    "pods",
                    "-l",
                    f"app.kubernetes.io/instance={instance_label}",
                    "-n",
                    milvus_ns,
                    "--no-headers",
                ],
                capture_output=True,
                text=True,
                check=False,
            )
            if existence.stdout.strip():
                break
            _time.sleep(5)
        else:
            raise TimeoutError(f"no pods matched {instance_label} within {timeout}s")

        remaining = max(1, int(deadline - _time.time()))
        cmd = [
            "kubectl",
            "wait",
            "--for=condition=Ready",
            "pods",
            "-l",
            f"app.kubernetes.io/instance={instance_label}",
            "-n",
            milvus_ns,
            f"--timeout={remaining}s",
        ]
        result = _sp.run(cmd, capture_output=True, text=True, check=False)
        logger.info(f"[KUBECTL] wait pods {instance_label}: rc={result.returncode}")
        if result.returncode != 0:
            raise TimeoutError(
                f"pods matching {instance_label} did not become Ready within {remaining}s: "
                f"stdout={result.stdout!r}, stderr={result.stderr!r}"
            )
        return result

    class KubectlHelper:
        pass

    KubectlHelper.delete_pods = staticmethod(delete_pods)
    KubectlHelper.wait_for_pods_ready = staticmethod(wait_for_pods_ready)

    return KubectlHelper()


@pytest.fixture(scope="session", autouse=True)
def cdc_topology_setup(request, upstream_client, downstream_client):
    """Setup CDC topology at the beginning of test session."""
    upstream_uri = request.config.getoption("--upstream-uri")
    downstream_uri = request.config.getoption("--downstream-uri")
    source_cluster_id = request.config.getoption("--source-cluster-id")
    target_cluster_id = request.config.getoption("--target-cluster-id")
    pchannel_num = int(request.config.getoption("--pchannel-num"))

    logger.info(f"Setting up CDC topology: {source_cluster_id} -> {target_cluster_id} (channels: {pchannel_num})...")

    # Create CDC replication configuration
    config = {
        "clusters": [
            {
                "cluster_id": source_cluster_id,
                "connection_param": {
                    "uri": upstream_uri,
                    "token": request.config.getoption("--upstream-token"),
                },
                "pchannels": [f"{source_cluster_id}-rootcoord-dml_{i}" for i in range(pchannel_num)],
            },
            {
                "cluster_id": target_cluster_id,
                "connection_param": {
                    "uri": downstream_uri,
                    "token": request.config.getoption("--downstream-token"),
                },
                "pchannels": [f"{target_cluster_id}-rootcoord-dml_{i}" for i in range(pchannel_num)],
            },
        ],
        "cross_cluster_topology": [
            {
                "source_cluster_id": source_cluster_id,
                "target_cluster_id": target_cluster_id,
            }
        ],
    }

    try:
        # Dedicated clients for the control-plane update_replicate_configuration
        # RPC, mirroring switchover_helper. Keeps the session-scoped clients'
        # channels clean of any recovery side effects from the initial setup.
        up_tmp = MilvusClient(uri=upstream_uri, token=request.config.getoption("--upstream-token"))
        dn_tmp = MilvusClient(uri=downstream_uri, token=request.config.getoption("--downstream-token"))
        try:
            apply_replicate_configuration([(up_tmp, config), (dn_tmp, config)])
        finally:
            up_tmp.close()
            dn_tmp.close()
        logger.info("CDC topology setup completed successfully")

        # Allow some time for CDC to initialize
        time.sleep(5)

    except Exception as e:
        logger.error(f"Failed to setup CDC topology: {e}")
        raise

    yield

    # Cleanup can be added here if needed
    logger.info("CDC topology teardown completed")