datastore/tests/integration/test_cluster_discovery/test.py

import functools
import time

import pytest

from helpers.cluster import ClickHouseCluster

from .common import check_on_cluster

cluster = ClickHouseCluster(__file__)

shard_configs = {
    "node0": "config/config.xml",
    "node1": "config/config_shard1.xml",
    "node2": "config/config.xml",
    "node3": "config/config_shard3.xml",
    "node4": "config/config.xml",
    "node_observer": "config/config_observer.xml",
}

nodes = {
    node_name: cluster.add_instance(
        node_name,
        main_configs=[shard_config, "config/config_discovery_path.xml", "config/macros.xml"],
        stay_alive=True,
        with_zookeeper=True,
    )
    for node_name, shard_config in shard_configs.items()
}


@pytest.fixture(scope="module")
def start_cluster():
    try:
        cluster.start()
        yield cluster
    finally:
        cluster.shutdown()


def test_cluster_discovery_startup_and_stop(start_cluster):
    """
    Start cluster, check nodes count in system.clusters,
    then stop/start some nodes and check that it (dis)appeared in cluster.
    """

    check_nodes_count = functools.partial(
        check_on_cluster, what="count()", msg="Wrong nodes count in cluster"
    )
    check_shard_num = functools.partial(
        check_on_cluster,
        what="count(DISTINCT shard_num)",
        msg="Wrong shard_num count in cluster",
    )

    # `- 1` because one node is an observer
    total_shards = len(set(shard_configs.values())) - 1
    total_nodes = len(nodes) - 1

    check_nodes_count(
        [nodes["node0"], nodes["node2"], nodes["node_observer"]], total_nodes
    )
    check_shard_num(
        [nodes["node0"], nodes["node2"], nodes["node_observer"]], total_shards
    )

    # test ON CLUSTER query
    nodes["node0"].query(
        "CREATE TABLE tbl ON CLUSTER 'test_auto_cluster' (x UInt64) ENGINE = MergeTree ORDER BY x"
    )
    nodes["node0"].query("INSERT INTO tbl VALUES (1)")
    nodes["node1"].query("INSERT INTO tbl VALUES (2)")

    assert (
        int(
            nodes["node_observer"]
            .query(
                "SELECT sum(x) FROM clusterAllReplicas(test_auto_cluster, default.tbl) settings serialize_query_plan=0"
            )
            .strip()
        )
        == 3
    )

    # Query SYSTEM CLEAR DNS CACHE may reload cluster configuration
    # check that it does not affect cluster discovery
    nodes["node1"].query("SYSTEM CLEAR DNS CACHE")
    nodes["node0"].query("SYSTEM CLEAR DNS CACHE")

    check_shard_num(
        [nodes["node0"], nodes["node2"], nodes["node_observer"]], total_shards
    )

    nodes["node1"].stop_clickhouse(kill=True)
    check_nodes_count(
        [nodes["node0"], nodes["node2"], nodes["node_observer"]], total_nodes - 1
    )

    # node1 was the only node in shard '1'
    check_shard_num(
        [nodes["node0"], nodes["node2"], nodes["node_observer"]], total_shards - 1
    )

    nodes["node3"].stop_clickhouse()
    check_nodes_count(
        [nodes["node0"], nodes["node2"], nodes["node_observer"]], total_nodes - 2
    )

    nodes["node1"].start_clickhouse()
    check_nodes_count(
        [nodes["node0"], nodes["node2"], nodes["node_observer"]], total_nodes - 1
    )

    nodes["node3"].start_clickhouse()
    check_nodes_count(
        [nodes["node0"], nodes["node2"], nodes["node_observer"]], total_nodes
    )

    # regular cluster is not affected
    check_nodes_count(
        [nodes["node1"], nodes["node2"]], 2, cluster_name="two_shards", retries=1
    )

    # cleanup
    nodes["node0"].query("DROP TABLE tbl ON CLUSTER 'test_auto_cluster' SYNC")


def test_cluster_discovery_no_crash_on_empty_static_cluster(start_cluster):
    """
    Regression test: server must not crash when a static cluster temporarily has no nodes.

    Crash path (without fix):
    1. All member nodes stop → ZK ephemeral nodes deleted → observer detects empty cluster.
    2. upsertCluster sees nodes_info.empty() → calls removeCluster() → erases watch callback
       from get_nodes_callbacks.
    3. A member node restarts → old ZK watch fires → cluster re-queued for update.
    4. upsertCluster is called again → on_exit calls getNodeNames(set_callback=True).
    5. Callback not found, zk_root_index==0 → multicluster_discovery_paths[0-1] OOB → crash.
    """
    import time

    observer = nodes["node_observer"]
    member_nodes = [nodes[n] for n in ("node0", "node1", "node2", "node3", "node4")]

    # Wait for cluster to be fully up before the test
    check_on_cluster(
        [observer], len(member_nodes), what="count()", msg="Pre-test cluster count wrong"
    )

    # Stop all member nodes gracefully — closes ZK sessions immediately,
    # deleting their ephemeral registration nodes in ZK.
    for node in member_nodes:
        node.stop_clickhouse()

    # Wait for the observer to detect the now-empty cluster.
    # This triggers upsertCluster → nodes_info.empty() → removeCluster() erasing the
    # watch callback (the step that sets up the crash on the subsequent call).
    check_on_cluster(
        [observer], 0, what="count()", msg="Cluster should appear empty to observer"
    )

    assert observer.query("SELECT 1").strip() == "1", "Observer crashed after cluster became empty"

    # Restart one member node. It registers in ZK, which triggers the old watch callback
    # that was stored in the ZK client (independently of get_nodes_callbacks).
    # On unfixed code this triggers the second upsertCluster call that crashes.
    nodes["node0"].start_clickhouse()
    time.sleep(10)

    assert observer.query("SELECT 1").strip() == "1", "Observer crashed after member node restarted"

    # Restore all nodes for subsequent tests
    for node in member_nodes[1:]:
        node.start_clickhouse()

    check_on_cluster(
        [observer], len(member_nodes), what="count()", msg="Cluster should recover after restart"
    )


def test_cluster_discovery_macros(start_cluster):
    # wait for all nodes to be started
    check_nodes_count = functools.partial(
        check_on_cluster, what="count()", msg="Wrong nodes count in cluster"
    )
    total_nodes = len(nodes) - 1
    check_nodes_count([nodes["node_observer"]], total_nodes, retries=6)

    # Poll the distributed query instead of asserting once: it needs every replica
    # reachable, and the cluster may still be settling when this test starts, so a
    # single read can flake. Retry (tolerating transient errors) until it converges.
    expected = str(total_nodes)
    res = None
    last_exception = None
    for _ in range(30):
        try:
            res = nodes["node_observer"].query(
                "SELECT sum(number) FROM clusterAllReplicas('{autocluster}', system.numbers) WHERE number=1"
            ).strip()
            if res == expected:
                break
        except Exception as e:
            last_exception = e
        time.sleep(2)
    assert res == expected, (
        f"Wrong macro result: {res}, expected {expected}, last exception: {last_exception!r}"
    )