datastore/tests/integration/test_keeper_disks/test.py

#!/usr/bin/env python3
import os

import pytest
from minio.deleteobjects import DeleteObject

import helpers.keeper_utils as keeper_utils
from helpers.cluster import ClickHouseCluster
import time

import logging

CURRENT_TEST_DIR = os.path.dirname(os.path.abspath(__file__))
cluster = ClickHouseCluster(__file__)
node_logs = cluster.add_instance(
    "node_logs",
    main_configs=["configs/enable_keeper.xml"],
    stay_alive=True,
    with_minio=True,
)

node_snapshot = cluster.add_instance(
    "node_snapshot",
    main_configs=["configs/enable_keeper_snapshot.xml"],
    stay_alive=True,
    with_minio=True,
)


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

    finally:
        cluster.shutdown()


def get_fake_zk(nodename, timeout=30.0):
    return keeper_utils.get_fake_zk(cluster, nodename, timeout=timeout)


def stop_zk(zk):
    try:
        if zk:
            zk.stop()
            zk.close()
    except:
        pass


def stop_clickhouse(cluster, node, cleanup_disks):
    node.stop_clickhouse()

    if not cleanup_disks:
        return

    node.exec_in_container(["rm", "-rf", "/var/lib/clickhouse/coordination/logs"])
    node.exec_in_container(["rm", "-rf", "/var/lib/clickhouse/coordination/snapshots"])

    # Wipe everything, including any leftover "tmp_" markers, so the bucket is clean.
    s3_objects = list_s3_objects(cluster, prefix="", exclude_tmp=False)
    if len(s3_objects) == 0:
        return

    assert (
        len(
            list(
                cluster.minio_client.remove_objects(
                    cluster.minio_bucket,
                    [DeleteObject(obj) for obj in s3_objects],
                )
            )
        )
        == 0
    )


def setup_storage(cluster, node, storage_config, cleanup_disks):
    stop_clickhouse(cluster, node, cleanup_disks)
    node.copy_file_to_container(
        os.path.join(CURRENT_TEST_DIR, "configs/enable_keeper.xml"),
        "/etc/clickhouse-server/config.d/enable_keeper.xml",
    )
    node.replace_in_config(
        "/etc/clickhouse-server/config.d/enable_keeper.xml",
        "<!-- DISK DEFINITION PLACEHOLDER -->",
        storage_config,
    )

    node.start_clickhouse()
    # complete readiness checks that the sessions can be established,
    # but it creates sesssion for this, which will create one more record in log,
    # but this test is very strict on number of entries in the log,
    # so let's avoid this extra check and rely on retry policy
    keeper_utils.wait_until_connected(cluster, node, wait_complete_readiness=False)


def setup_local_storage(cluster, node):
    setup_storage(
        cluster,
        node,
        "<log_storage_disk>log_local<\\/log_storage_disk>"
        "<snapshot_storage_disk>snapshot_local<\\/snapshot_storage_disk>",
        cleanup_disks=True,
    )


# Keeper writes an empty "tmp_<name>" marker on the target disk while moving a snapshot/log
# file between disks and removes it once the copy completes (moveFileBetweenDisks in
# KeeperCommon.cpp). Such markers are never real snapshots/logs: the server skips
# "tmp_"-prefixed entries when scanning disks. A listing taken while a move is in flight can
# transiently observe a marker, so exclude it instead of counting it as a snapshot/log file.
tmp_keeper_file_prefix = "tmp_"


def list_s3_objects(cluster, prefix="", exclude_tmp=True):
    minio = cluster.minio_client
    prefix_len = len(prefix)
    return [
        obj.object_name[prefix_len:]
        for obj in minio.list_objects(
            cluster.minio_bucket, prefix=prefix, recursive=True
        )
        if not (
            exclude_tmp
            and os.path.basename(obj.object_name).startswith(tmp_keeper_file_prefix)
        )
    ]


def get_local_files(path, node):
    files = node.exec_in_container(["ls", path]).strip().split("\n")
    files = [f for f in files if not f.startswith(tmp_keeper_file_prefix)]
    files.sort()
    return files


def get_local_logs(node):
    return get_local_files("/var/lib/clickhouse/coordination/logs", node)


def get_local_snapshots(node):
    return get_local_files("/var/lib/clickhouse/coordination/snapshots", node)


def test_logs_with_disks(started_cluster):
    setup_local_storage(started_cluster, node_logs)

    node_zk = get_fake_zk("node_logs")
    try:
        node_zk.create("/test")
        for _ in range(30):
            node_zk.create("/test/somenode", b"somedata", sequence=True)

        stop_zk(node_zk)

        assert len(get_local_logs(node_logs)) > 1

        setup_storage(
            started_cluster,
            node_logs,
            "<log_storage_disk>log_s3_plain<\\/log_storage_disk>"
            "<latest_log_storage_disk>log_local<\\/latest_log_storage_disk>"
            "<snapshot_storage_disk>snapshot_local<\\/snapshot_storage_disk>",
            cleanup_disks=False,
        )

        # all but the latest log should be on S3
        def assert_single_local_log():
            local_log_files = get_local_logs(node_logs)
            start_time = time.time()
            while len(local_log_files) != 1:
                logging.debug(f"Local log files: {local_log_files}")
                assert (
                    time.time() - start_time < 60
                ), "local_log_files size is not equal to 1 after 60s"
                time.sleep(1)
                local_log_files = get_local_logs(node_logs)

        assert_single_local_log()
        assert len(list_s3_objects(started_cluster, "logs/")) > 1

        node_zk = get_fake_zk("node_logs")

        for _ in range(30):
            node_zk.create("/test/somenode", b"somedata", sequence=True)

        assert_single_local_log()
        stop_zk(node_zk)

        setup_storage(
            started_cluster,
            node_logs,
            "<old_log_storage_disk>log_s3_plain<\\/old_log_storage_disk>"
            "<log_storage_disk>log_local<\\/log_storage_disk>"
            "<snapshot_storage_disk>snapshot_local<\\/snapshot_storage_disk>",
            cleanup_disks=False,
        )

        s3_files = list_s3_objects(started_cluster, "logs/")
        start_time = time.time()
        while len(s3_files) != 0:
            logging.debug(f"S3 log files: {s3_files}")
            assert (
                time.time() - start_time < 60
            ), "s3_files size is not equal to 0 after 60s"
            time.sleep(1)
            s3_files = list_s3_objects(started_cluster, "logs/")

        assert len(get_local_logs(node_logs)) > 1

        node_zk = get_fake_zk("node_logs")

        for child in node_zk.get_children("/test"):
            assert node_zk.get(f"/test/{child}")[0] == b"somedata"

    finally:
        stop_zk(node_zk)


def test_snapshots_with_disks(started_cluster):
    setup_local_storage(started_cluster, node_snapshot)

    node_zk = get_fake_zk("node_snapshot")
    try:
        node_zk.create("/test2")
        for _ in range(30):
            node_zk.create("/test2/somenode", b"somedata", sequence=True)

        stop_zk(node_zk)

        snapshot_idx = keeper_utils.send_4lw_cmd(cluster, node_snapshot, "csnp")
        node_snapshot.wait_for_log_line(f"Created persistent snapshot {snapshot_idx}")

        previous_snapshot_files = get_local_snapshots(node_snapshot)

        setup_storage(
            started_cluster,
            node_snapshot,
            "<snapshot_storage_disk>snapshot_s3_plain<\\/snapshot_storage_disk>"
            "<latest_snapshot_storage_disk>snapshot_local<\\/latest_snapshot_storage_disk>"
            "<log_storage_disk>log_local<\\/log_storage_disk>",
            cleanup_disks=False,
        )

        # Relocating non-latest snapshots to the configured disk runs after the
        # "Created persistent snapshot" log line, so poll for the count to settle
        # instead of checking once (mirrors assert_single_local_log above).
        def assert_single_local_snapshot():
            local_snapshot_files = get_local_snapshots(node_snapshot)
            start_time = time.time()
            while len(local_snapshot_files) != 1:
                logging.debug(f"Local snapshot files: {local_snapshot_files}")
                assert (
                    time.time() - start_time < 60
                ), "local_snapshot_files size is not equal to 1 after 60s"
                time.sleep(1)
                local_snapshot_files = get_local_snapshots(node_snapshot)

        ## all but the latest snapshot should be on S3
        assert_single_local_snapshot()
        local_snapshot_files = get_local_snapshots(node_snapshot)
        s3_snapshot_files = list_s3_objects(started_cluster, "snapshots/")
        assert set(s3_snapshot_files) == set(previous_snapshot_files[:-1])
        assert local_snapshot_files[0] == previous_snapshot_files[-1]

        previous_snapshot_files = s3_snapshot_files + local_snapshot_files

        node_zk = get_fake_zk("node_snapshot")

        for _ in range(30):
            node_zk.create("/test2/somenode", b"somedata", sequence=True)

        stop_zk(node_zk)

        snapshot_idx = keeper_utils.send_4lw_cmd(cluster, node_snapshot, "csnp")
        node_snapshot.wait_for_log_line(f"Created persistent snapshot {snapshot_idx}")

        assert_single_local_snapshot()
        local_snapshot_files = get_local_snapshots(node_snapshot)
        snapshot_files = list_s3_objects(started_cluster, "snapshots/")

        snapshot_files.extend(local_snapshot_files)

        previous_snapshot_files = snapshot_files

        setup_storage(
            started_cluster,
            node_snapshot,
            "<old_snapshot_storage_disk>snapshot_s3_plain<\\/old_snapshot_storage_disk>"
            "<snapshot_storage_disk>snapshot_local<\\/snapshot_storage_disk>"
            "<log_storage_disk>log_local<\\/log_storage_disk>",
            cleanup_disks=False,
        )

        local_snapshot_files = get_local_snapshots(node_snapshot)
        assert set(local_snapshot_files) == set(previous_snapshot_files)

        node_zk = get_fake_zk("node_snapshot")

        for child in node_zk.get_children("/test2"):
            assert node_zk.get(f"/test2/{child}")[0] == b"somedata"

    finally:
        stop_zk(node_zk)