#!/usr/bin/env python3
# pylint: disable=line-too-long

import os
import subprocess
import multiprocessing
from tempfile import NamedTemporaryFile
import pytest


CPU_ID = 4


def run_command_in_container(cmd, *args):
    # /clickhouse is mounted by integration tests runner
    alternative_binary = os.getenv("CLICKHOUSE_BINARY", "/clickhouse")
    if alternative_binary:
        args += (
            "--volume",
            f"{alternative_binary}:/usr/bin/clickhouse",
        )

    return subprocess.check_output(
        [
            "docker",
            "run",
            "--rm",
            *args,
            "ubuntu:20.04",
            "sh",
            "-c",
            cmd,
        ]
    )


def run_with_cpu_limit(cmd, *args):
    with NamedTemporaryFile() as online_cpu:
        # NOTE: this is not the number of CPUs, but specific CPU ID
        online_cpu.write(f"{CPU_ID}".encode())
        online_cpu.flush()

        # replace /sys/devices/system/cpu/online to full _SC_NPROCESSORS_ONLN
        # like LXD/LXC from [1] does.
        #
        #   [1]: https://github.com/ClickHouse/ClickHouse/issues/32806
        args += (
            "--volume",
            f"{online_cpu.name}:/sys/devices/system/cpu/online",
        )

        return run_command_in_container(cmd, *args)


def skip_if_jemalloc_disabled():
    output = run_command_in_container(
        """clickhouse local -q "
        SELECT value FROM system.build_options WHERE name = 'USE_JEMALLOC'"
    """
    ).strip()
    if output != b"ON" and output != b"1":
        pytest.skip(f"Compiled w/o jemalloc (USE_JEMALLOC={output})")


# Ensure that clickhouse works even when number of online CPUs
# (_SC_NPROCESSORS_ONLN) is smaller then available (_SC_NPROCESSORS_CONF).
#
# Refs: https://github.com/jemalloc/jemalloc/pull/2181
def test_jemalloc_percpu_arena():
    skip_if_jemalloc_disabled()

    assert multiprocessing.cpu_count() > CPU_ID

    online_cpus = int(run_with_cpu_limit("getconf _NPROCESSORS_ONLN"))
    assert online_cpus == 1, online_cpus

    all_cpus = int(run_with_cpu_limit("getconf _NPROCESSORS_CONF"))
    assert all_cpus == multiprocessing.cpu_count(), all_cpus

    # implicitly disable percpu arena
    result = run_with_cpu_limit(
        'clickhouse local -q "select 1"',
        # NOTE: explicitly disable, since it is enabled by default in debug build
        # (and even though debug builds are not in CI let's state this).
        "--env",
        "MALLOC_CONF=abort_conf:false",
    )
    assert int(result) == int(1), result

    # should fail because of abort_conf:true
    with pytest.raises(subprocess.CalledProcessError):
        run_with_cpu_limit(
            'clickhouse local -q "select 1"', "--env", "MALLOC_CONF=abort_conf:true"
        )

    # should not fail even with abort_conf:true, due to explicit narenas
    # NOTE: abort:false to make it compatible with debug build
    run_with_cpu_limit(
        'clickhouse local -q "select 1"',
        "--env",
        f"MALLOC_CONF=abort_conf:true,abort:false,narenas:{all_cpus}",
    )


# For manual run.
if __name__ == "__main__":
    test_jemalloc_percpu_arena()