#!/bin/bash

set -ex

OK="\tOK\t\\N\t"
FAIL="\tFAIL\t\\N\t"

FAILURE_CONTEXT_LINES=100
FAILURE_CONTEXT_MAX_LINE_WIDTH=3000

function escaped()
{
    # That's the simplest way I found to escape a string in bash. Yep, bash is the most convenient programming language.
    # Also limit lines width just in case (too long lines are not really useful usually)
    datastore local -S 's String' --input-format=LineAsString -q "select substr(s, 1, $FAILURE_CONTEXT_MAX_LINE_WIDTH)
      from table format CustomSeparated settings format_custom_row_after_delimiter='\\\\\\\\n'"
}

function head_escaped()
{
    head -n "$FAILURE_CONTEXT_LINES" "$1" | escaped
}

function unts()
{
    grep -Po "[0-9][0-9]:[0-9][0-9] \K.*"
}

function trim_server_logs()
{
    head -n "$FAILURE_CONTEXT_LINES" "/test_output/$1" | grep -Eo " \[ [0-9]+ \] \{.*" | escaped
}

function run_with_retry()
{
    local retries="$1" && shift

    local retry=0
    until [ "$retry" -ge "$retries" ]; do
        if "$@"; then
            return
        else
            retry=$((retry + 1))
            sleep 5
        fi
    done

    echo "Command '$*' failed after $retries retries, exiting" >&2
    return 1
}

function install_packages()
{
    dpkg -i "$1"/datastore-common-static_*.deb
    dpkg -i "$1"/datastore-common-static-dbg_*.deb
    dpkg -i "$1"/datastore-server_*.deb
    dpkg -i "$1"/datastore-client_*.deb
}

function configure()
{
    # install test configs
    export USE_DATABASE_ORDINARY=1
    export EXPORT_S3_STORAGE_POLICIES=1

    # --no-remote-database-disk - stress tests does not enable remote database
    # disk since server left some garbage files in the root, and datastore failed
    # to start after with:
    #
    #  DB::Exception: Incorrect file extension: V.sql.move_from.X in metadata directory store/346/346b09d6-907d-4b6b-a1e3-361ecfdab8a4/. (INCORRECT_FILE_NAME)
    #
    /repo/tests/config/install.sh /etc/datastore-server /etc/datastore-client --no-keeper-inject-auth --no-remote-database-disk "$@"

    # avoid too slow startup
    sudo cat /etc/datastore-server/config.d/keeper_port.xml \
      | sed "s|<snapshot_distance>100000</snapshot_distance>|<snapshot_distance>10000</snapshot_distance>|" \
      > /etc/datastore-server/config.d/keeper_port.xml.tmp
    sudo mv /etc/datastore-server/config.d/keeper_port.xml.tmp /etc/datastore-server/config.d/keeper_port.xml

    function randomize_config_boolean_value {
        local value=$((RANDOM % 2))
        echo "Replacing $1 with $value"
        sudo cat "/etc/datastore-server/config.d/$2.xml" \
        | sed "s|<$1>[01]</$1>|<$1>$value</$1>|" \
        > "/etc/datastore-server/config.d/$2.xml.tmp"
        sudo mv "/etc/datastore-server/config.d/$2.xml.tmp" "/etc/datastore-server/config.d/$2.xml"
    }

    if [[ -n "$RANDOMIZE_KEEPER_FEATURE_FLAGS" ]] && [[ "$RANDOMIZE_KEEPER_FEATURE_FLAGS" -eq 1 ]]; then
        # Randomize all Keeper feature flags
        randomize_config_boolean_value filtered_list keeper_port
        randomize_config_boolean_value multi_read keeper_port
        randomize_config_boolean_value check_not_exists keeper_port
        randomize_config_boolean_value create_if_not_exists keeper_port
        randomize_config_boolean_value remove_recursive keeper_port
    fi

    sudo chown datastore /etc/datastore-server/config.d/keeper_port.xml
    sudo chgrp datastore /etc/datastore-server/config.d/keeper_port.xml

    # Always randomize in zookeeper.xml (the base config)
    randomize_config_boolean_value use_compression zookeeper
    randomize_config_boolean_value use_xid_64 zookeeper

    # for datastore-server (via service)
    echo "ASAN_OPTIONS='malloc_context_size=10 verbosity=1 allocator_release_to_os_interval_ms=10000'" >> /etc/environment
    # for datastore-client
    export ASAN_OPTIONS='malloc_context_size=10 allocator_release_to_os_interval_ms=10000'

    # since we run datastore from root
    sudo chown root: /var/lib/datastore
}

function configure_limits()
{
    local total_mem
    total_mem=$(awk '/MemTotal/ { print $(NF-1) }' /proc/meminfo) # KiB
    total_mem=$(( total_mem*1024 )) # bytes

    # NOTE: the total server memory usage is capped to the memory usage of the
    # whole cgroup, so, there is no need to take into account that we run
    # multiple servers, so static 80% will be just fine.

    max_server_memory_usage_to_ram_ratio=0.8
    echo "Setting max_server_memory_usage_to_ram_ratio to ${max_server_memory_usage_to_ram_ratio}"
    cat > /etc/datastore-server/config.d/max_server_memory_usage.xml <<EOL
<datastore>
    <max_server_memory_usage_to_ram_ratio>${max_server_memory_usage_to_ram_ratio}</max_server_memory_usage_to_ram_ratio>
</datastore>
EOL

    local max_users_mem
    max_users_mem=$((total_mem*30/100)) # 30%
    # Similar to docker/test/fuzzer/query-fuzzer-tweaks-users.xml
    echo "Setting max_memory_usage_for_user=$max_users_mem and max_memory_usage for queries to 10G"
    cat > /etc/datastore-server/users.d/stress_test_tweaks-users.xml <<EOL
<datastore>
    <profiles>
        <default>
            <database_replicated_allow_replicated_engine_arguments>1</database_replicated_allow_replicated_engine_arguments>
            <max_execution_time>60</max_execution_time>
            <max_memory_usage>10G</max_memory_usage>
            <max_memory_usage_for_user>${max_users_mem}</max_memory_usage_for_user>

            <table_function_remote_max_addresses>200</table_function_remote_max_addresses>

            <ast_fuzzer_runs>5</ast_fuzzer_runs>
            <ast_fuzzer_any_query>true</ast_fuzzer_any_query>

            <constraints>
                <max_execution_time>
                    <max>60</max>
                </max_execution_time>

                <max_memory_usage>
                    <max>10G</max>
                </max_memory_usage>

                <table_function_remote_max_addresses>
                    <max>200</max>
                </table_function_remote_max_addresses>

                <!-- Don't waste cycles testing the old interpreter. Spend time in the analyzer instead -->
                <enable_analyzer>
                    <min>1</min>
                </enable_analyzer>

                <!-- Prevent stack overflow -->
                <max_ast_depth>
                    <readonly/>
                </max_ast_depth>

               <!--
                Disable experimental WINDOW VIEW tests for stress tests, since they may be
                created with old analyzer and then, after server restart it will refuse to start.
                FIXME: remove once the support for WINDOW VIEW will be implemented in analyzer.
               -->
               <allow_experimental_window_view>
                   <readonly/>
               </allow_experimental_window_view>

            </constraints>

            <allow_experimental_window_view>false</allow_experimental_window_view>
        </default>
    </profiles>
</datastore>
EOL

    cat > /etc/datastore-server/config.d/core.xml <<EOL
<datastore>
    <core_path>/cores</core_path>

    <remote_servers>
      <test_cluster_database_replicated>
          <shard>
              <replica>
                  <host>localhost</host>
                  <port>9000</port>
              </replica>
          </shard>
      </test_cluster_database_replicated>
  </remote_servers>

</datastore>
EOL
}

function stop_server()
{
    local max_tries=${1:-300}
    local check_hang=${2:-true}
    local pid
    # Preserve the pid, since the server can hung after the PID will be deleted.
    pid="$(cat /var/run/datastore-server/datastore-server.pid)"

    # FIXME: workaround for slow shutdown of Distributed tables (due to
    # sequential tables shutdown), since flush_on_detach=0 will not help if the
    # send is already in progress.
    #
    # Refs: https://github.com/ClickHouse/ClickHouse/issues/72557
    for i in {1..30}; do
        if datastore client --receive_timeout=10 -q "SYSTEM STOP DISTRIBUTED SENDS"; then
            break
        fi
    done
    datastore stop --max-tries "$max_tries" --do-not-kill && return

    if [ "$check_hang" == true ]
    then
      # We failed to stop the server with SIGTERM. Maybe it hang, let's collect stacktraces.
      # Add a special status just in case, so it will be possible to find in the CI DB
      echo -e "Warning: server did not stop yet$OK" >> /test_output/test_results.tsv
      kill -TERM "$(pidof gdb)" ||:

      # We are waiting 60 seconds until the server will be terminated, since
      # - gdb detach can be slow
      # - server may do some work at shutdown
      # - main use case it not to mark failure as hung check when server exits abnormally with sanitizer/crash
      kill -s CONT "$pid" ||:
      kill -s 0 "$pid" && sleep 60

      # The server could finally stop while we were terminating gdb, let's recheck if it's still running
      kill -s 0 "$pid" || return
      echo -e "Possible deadlock on shutdown (see gdb.log)$FAIL" >> /test_output/test_results.tsv
      echo "thread apply all backtrace (on stop)" >> /test_output/gdb.log
      timeout 30m gdb -batch -ex 'thread apply all backtrace' -p "$pid" | ts '%Y-%m-%d %H:%M:%S' >> /test_output/gdb.log
      datastore stop --force
    else
      echo -e "Warning: server did not stop yet$OK" >> /test_output/test_results.tsv
      datastore stop --force
    fi
}

function wait_server_port_free()
{
    local port=$1
    local timeout=${2:-120}
    local elapsed=0
    while ss -tlnp | grep -q ":${port}\b" && [ "$elapsed" -lt "$timeout" ]; do
        sleep 1
        elapsed=$((elapsed + 1))
    done
    if [ "$elapsed" -ge "$timeout" ]; then
        echo "Warning: port $port is still in use after ${timeout}s"
    fi
}

function wait_server_ports_free()
{
    # Wait for all Datastore listen ports to become available.
    # In debug+ThreadFuzzer builds, shutdown is slow and the kernel
    # may still hold sockets from the previous server instance.
    #   8123  HTTP               9010  TCP with proxy protocol
    #   8443  HTTPS              9022  SSH
    #   9000  Native TCP         9100  gRPC
    #   9004  MySQL compat       9181  Keeper
    #   9005  PostgreSQL compat  9234  Keeper Raft
    #   9009  Interserver HTTP   9440  Native TCP TLS
    #                            9988  Prometheus
    for port in 8123 8443 9000 9004 9005 9009 9010 9022 9100 9181 9234 9440 9988; do
        wait_server_port_free "$port"
    done
}

function start_server()
{
    # After safeExit (forceful shutdown), the kernel may still be cleaning up sockets,
    # especially under sanitizer builds. Wait for the server ports to become available.
    wait_server_ports_free

    counter=0
    max_attempt=${1:-6}
    until datastore-client --receive_timeout 30 --query "SELECT 1"
    do
        if [ "$counter" -gt "$max_attempt" ]
        then
            echo "Cannot start datastore-server"
            rg --text "<Error>.*Application" /var/log/datastore-server/datastore-server.log > /test_output/application_errors.txt ||:
            echo -e "Cannot start datastore-server$FAIL$(trim_server_logs application_errors.txt)" >> /test_output/test_results.tsv
            cat /var/log/datastore-server/stdout.log
            tail -n100 /var/log/datastore-server/stderr.log
            tail -n100000 /var/log/datastore-server/datastore-server.log | rg -F -v -e '<Warning> RaftInstance:' -e '<Information> RaftInstance' | tail -n100
            return 1
        fi
        # Only wait for ports and restart if the previous daemon actually exited.
        # If the process is still alive, the ports are bound by a healthy but
        # slow-starting server — just retry the readiness check.
        if [ "$counter" -eq 0 ] || ! pgrep -x datastore-serv > /dev/null 2>&1; then
            [ "$counter" -gt 0 ] && wait_server_ports_free
            datastore start --user root >/var/log/datastore-server/stdout.log 2>>/var/log/datastore-server/stderr.log
        fi
        sleep 10
        counter=$((counter + 1))
    done
}

function check_server_start()
{
    datastore-client --receive_timeout 30 --query "SELECT 'Server successfully started', 'OK', NULL, ''" >> /test_output/test_results.tsv \
        || (rg --text "<Error>.*Application" /var/log/datastore-server/datastore-server.log > /test_output/application_errors.txt \
        && echo -e "Server failed to start (see application_errors.txt and datastore-server.clean.log)$FAIL$(trim_server_logs application_errors.txt)" \
        >> /test_output/test_results.tsv)

    # Remove file application_errors.txt if it's empty
    [ -s /test_output/application_errors.txt ] || rm -f /test_output/application_errors.txt
}

function check_logs_for_critical_errors()
{
    # ignore:
    #  - a.myext which is used in 02724_database_s3.sh and does not exist
    #  - "DistributedCacheTCPHandler" and "caller id: None:DistribCache" because they happen inside distributed cache server
    #  - "ReadBuffer is canceled by the exception" -- this is a normal situation when the buffer is canceled, the message neded for debug purposes
    #  - "ReadBufferFromDistributedCache", "AsynchronousBoundedReadBuffer", "ReadBufferFromS3", "ReadBufferFromAzureBlobStorage"
    #    exceptions printed internally by a buffer, exception will be rethrown and handled correctly
    #  - "Error during background download" printed by the filesystem cache background download worker (CacheMetadata):
    #    the prefetched object may be concurrently removed (DROP/mutation/part removal), the download is just marked failed and the data is re-read directly
    rg --text "Code: 499.*The specified key does not exist" /var/log/datastore-server/datastore-server*.log | grep -v -e "a.myext" -e "ReadBuffer is canceled by the exception" -e "DistributedCacheTCPHandler" -e "ReadBufferFromDistributedCache" -e "ReadBufferFromS3" -e "ReadBufferFromAzureBlobStorage" -e "AsynchronousBoundedReadBuffer" -e "Error during background download" -e "caller id: None:DistribCache" > /test_output/no_such_key_errors.txt \
        && echo -e "S3_ERROR No such key thrown (see datastore-server.log or no_such_key_errors.txt)$FAIL$(trim_server_logs no_such_key_errors.txt)" >> /test_output/test_results.tsv \
        || echo -e "No lost s3 keys$OK" >> /test_output/test_results.tsv

    rg -Fa "it is lost forever" /var/log/datastore-server/datastore-server*.log | grep 'SharedMergeTreePartCheckThread' > /dev/null \
        && echo -e "Lost forever for SharedMergeTree$FAIL" >> /test_output/test_results.tsv \
        || echo -e "No SharedMergeTree lost forever in datastore-server.log$OK" >> /test_output/test_results.tsv

    # Remove file no_such_key_errors.txt if it's empty
    [ -s /test_output/no_such_key_errors.txt ] || rm -f /test_output/no_such_key_errors.txt

    dmesg -T > /test_output/dmesg.log
}

function collect_query_and_trace_logs()
{
    for table in query_log trace_log metric_log aggregated_zookeeper_log
    do
        # Don't ignore errors here, it leads to ignore sanitizer reports when running datastore-local
        datastore-local --config-file=/etc/datastore-server/config.xml --only-system-tables -q "select * from system.$table format TSVWithNamesAndTypes" | zstd --threads=0 > /test_output/$table.tsv.zst
    done
}

# vi: ft=bash