#! /bin/bash
# hpc-testing
# Copyright (C) 2025 SUSE LLC
# 
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
# 
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
# 
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

test_ibv_pingpong()
{
    local testname=$1

    local host1=$2
    local hca1=$3
    local ibport1=$4

    local host2=$5
    local hca2=$6
    local ibport2=$7

    tp $host2 "$testname -g 0 -d $hca2 -i $ibport2 $IBV_EXTRA_OPTS" &
    # Do a random sleep on the remote host so in case SSH is very slow we don't have
    # hardcoded sleep value
    tp $host2 "sleep 2"
    sleep 2
    tp $host1 "$testname -g 0 -d $hca1 -i $ibport1 $IBV_EXTRA_OPTS $host2"

    wait
}

run_rdma_client() {
    local c_host=$1
    local c_hca=$2
    local c_port=$3
    local testcmd=$4
    local opts=$5
    local s_host=$6
    local s_hca=$7
    local s_port=$8
    local log=$9

    set -e
    set -o pipefail

    # Start the Server (Background)
    tp $s_host "$testcmd -x 0 -d $s_hca -i $s_port $opts" &
    local server_pid=$!

    # CRITICAL: Clean up the server if the client fails or errors out
    trap "kill $server_pid 2>/dev/null || true; wait $server_pid 2>/dev/null || true" EXIT

    # Give the server 2 seconds to bind to the port
    sleep 2

    # Run the Client (Foreground) and log output
    tp $c_host "$testcmd -x 0 -d $c_hca -i $c_port $opts $s_host" | tee "$log"
}

rdma_perf(){
    local host1=$1
    local hca1=$2
    local ibport1=$3

    local host2=$4
    local hca2=$5
    local ibport2=$6

    TEST_NAMES=("ib_write_bw" "ib_write_lat" "ib_send_bw" "ib_send_lat" "ib_read_bw" "ib_read_lat")
    QP_TYPES=("RC" "UC" "UD")
    # Expanded to match the full output generated by the perftest '-a' flag
    PACKET_SIZES=(2 4 8 16 32 64 128 256 512 1024 2048 4096 8192 16384 32768 65536 131072 262144 524288 1048576 2097152 4194304 8388608)
    DURATION=5
    CSV_FILE="results/PERFTEST-verbs.csv"

    # --- Initialize CSV Header ---
    # Split into two columns: Test_Name and Data_Type (Unit)
    header="Test_Name,Data_Type"
    for size in "${PACKET_SIZES[@]}"; do
        header="${header},${size}B"
    done
    echo "$header" > "$CSV_FILE"

    HUGEPAGE1=$(tpq $host1 "sysctl vm.nr_hugepages | awk '{print \$NF}'; sysctl -w vm.nr_hugepages=2048 > /dev/null")
    HUGEPAGE2=$(tpq $host2 "sysctl vm.nr_hugepages | awk '{print \$NF}'; sysctl -w vm.nr_hugepages=2048 > /dev/null")


    # --- Execute Test Matrix ---
    for testname in "${TEST_NAMES[@]}"; do
        for qp in "${QP_TYPES[@]}"; do
            
            # --- Determine Unit of Measurement ---
            if [[ "$testname" == *"_bw" ]]; then
                unit="MiB/sec"
            elif [[ "$testname" == *"_lat" ]]; then
                unit="us"
            else
                unit="unknown"
            fi

            # Merge testname and QP for column 1, use the unit for column 2
            csv_row="${testname}_${qp},${unit}"

            # --- Test-Level Compatibility Safeguards ---
            # If the entire QP/Test combination is invalid, skip execution
            if [[ "$qp" == "UD" && "$testname" != *"send"* ]]; then
                continue
            fi
            if [[ "$qp" == "UC" && "$testname" == *"read"* ]]; then
                continue
            fi

            # --- Configure Options for --all (-a) ---
            # Notice we dropped -s $size in favor of -a
            IBV_EXTRA_OPTS="-c $qp -a --use_hugepages"
            if [[ "$testname" == *"_bw" ]]; then
                IBV_EXTRA_OPTS="$IBV_EXTRA_OPTS -F"
            fi

            log_file="results/RDMA-${testname}-${qp}-all.log"
            test_id="rdma_${testname}_${qp}_all"

            # Execute the test ONCE per QP/Test combination
            juLog -name="$test_id" "run_rdma_client $host1 $hca1 $ibport1 $testname '$IBV_EXTRA_OPTS' $host2 $hca2 $ibport2 $log_file"

            # Short sleep to ensure ports are fully released
            sleep 1

            # --- Data Extraction from Master Log File ---
            # Now we loop through the sizes just to parse the log we already generated
            for size in "${PACKET_SIZES[@]}"; do
                metric=""

                if [[ -f "$log_file" ]]; then
                    if [[ "$testname" == *"_bw" ]]; then
                        metric=$(grep -E "^[[:space:]]*$size[[:space:]]+" "$log_file" |
                                     tail -n 1 | awk '{print $4}')
                    elif [[ "$testname" == *"_lat" ]]; then
                        metric=$(grep -E "^[[:space:]]*$size[[:space:]]+" "$log_file" |
                                     tail -n 1 | awk '{print $6}')
                    fi
                fi

                # If the metric wasn't found (due to MTU limit, error, or missing log), default to N/A
                if [[ -z "$metric" ]]; then
                    metric="N/A"
                fi

                csv_row="${csv_row},${metric}"
            done

            # Write the completed row to the CSV
            echo "$csv_row" >> "$CSV_FILE"
        done
    done

    # Restore huge page count to original value
    tp $host1 "sysctl -w vm.nr_hugepages=$HUGEPAGE1"
    tp $host2 "sysctl -w vm.nr_hugepages=$HUGEPAGE2"
}
