engine

#!/bin/bash
# shellcheck disable=SC1091,SC2155,SC2086
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
CONFIG_DIR="$HOME/.config/engine-stressor"
CONSTANTS_FILE="$CONFIG_DIR/constants"

if [ ! -f "$CONSTANTS_FILE" ]; then
  echo "Error: File $CONSTANTS_FILE does not exist."
  exit 1
fi

source "$CONSTANTS_FILE"
source "$SHARE_DIR/cgroup"
source "$SHARE_DIR/common"
source "$SHARE_DIR/stress"
source "$SHARE_DIR/systemd"
source "$SHARE_DIR/selinux"
source "$SHARE_DIR/monitor"

# Install extra packages in the container
# Argument: container name
install_packages() {
    local container_name=$1

    if [[ "$VERBOSE" = "${FEATURE_ENABLED}" ]]; then
        echo -e "INFO: Installing packages in container ${container_name} - packages ${EXTRA_PACKAGES_TO_INSTALL}..."
    fi

    # Define the command as a variable
    cmd="sudo podman exec \"${container_name}\" ${PACKAGER_INSTALLER}"

    # Add PACKAGER_INSTALLER_EXTRA_FLAGS if not empty
    if [ -n "${PACKAGER_INSTALLER_EXTRA_FLAGS}" ]; then
        cmd+=" ${PACKAGER_INSTALLER_EXTRA_FLAGS}"
    fi

    # Complete the command
    cmd+=" install -y -q ${EXTRA_PACKAGES_TO_INSTALL}"

    # Echo the command for debugging
    echo "Executing command: $cmd"

    # Execute the command
    if ! eval $cmd; then
        if [[ "$VERBOSE" = "${FEATURE_ENABLED}" ]]; then
            echo -e "FAIL: container ${container_name} fail to install package(s) ${EXTRA_PACKAGES_TO_INSTALL}"
            return 1
        fi
    fi
    return 0
}

# Function to check if the OS is CentOS Stream 9 or higher and modify the Podman command
# FIX-ME: We found in CentOS 9 Stream some Read only filesystem cases, this is a probably temp. workaround
check_os_and_modify_podman_cmd() {
    local os_name
    local os_version

    echo "========== DEBUG: OS =================="
    # Debug: Display the contents of /etc/os-release
    sudo cat /etc/os-release
    echo "========== DEBUG: OS =================="

    os_name=$(grep -oP '(?<=^NAME=).+' /etc/os-release | tr -d '"')
    os_version=$(grep -oP '(?<=^VERSION_ID=).+' /etc/os-release | tr -d '"')

    # Check if the OS is CentOS Stream 9 or higher
    if [[ "$os_name" == "CentOS Stream" && "$os_version" -ge 9 ]]; then
        # Create necessary directories
        sudo mkdir -p /var/tmp/varlog /var/tmp/dnf-cache /var/tmp/dnf-lib /var/tmp/rpm

        # Add volume mounts to the Podman command
        podman_cmd+=" --volume /var/tmp/varlog:/var/log:Z"
        podman_cmd+=" --volume /var/tmp/dnf-cache:/var/cache/dnf:Z"
        podman_cmd+=" --volume /var/tmp/dnf-lib:/var/lib/dnf:Z"
        podman_cmd+=" --volume /var/tmp/rpm:/var/lib/rpm:Z"
    fi
}

# Function to create a container
create_container() {
    local container_name=$1
    local start_time=$(date +%s%3N) # Log start time in milliseconds
    local ostree_deployment=""

    # Running Podman containers under a systemd scope integrates them more
    # tightly with the systemd service manager.
    if stat /run/ostree-booted > /dev/null 2>&1 || [[ $(uname -r | grep -w -q el[0-9]*iv) -eq 0 ]]; then
        ostree_deployment="--read-only=false"
    fi

    podman_cmd="podman run ${ostree_deployment} -d --replace --name ${container_name}"

    if [ -n "${VOLUME_NAME}" ]; then
        podman_cmd+=" --volume ${VOLUME_NAME}:${VOLUME_PATH}"
    fi

    if [ -n "${NETWORK_NAME}" ]; then
        podman_cmd+=" --network ${NETWORK_NAME}"
    fi

    # Check if NETWORK_DNS_INFO is provided and construct the DNS parameters
    if [ -n "$NETWORK_DNS_SERVER" ]; then
        DNS_PARAMS=""
        for DNS in "${NETWORK_DNS_SERVER}"; do
            DNS_PARAMS+="--dns=$DNS "
        done
        podman_cmd+=" ${DNS_PARAMS}"
    fi

    podman_cmd+=" ${IMAGE_NAME_CONTAINER} ${IMAGE_COMMAND}"

    cmd="sudo systemd-run --quiet --scope \
              -p Delegate=yes \
              --slice=${CGROUP_NAME}.slice \
              ${podman_cmd}"

    # Run the command
    systemd_scope=$(eval $cmd)

    # Add the scope thing for the validation
    systemd_scope="libpod-${systemd_scope}.scope"

    # Check if it's running using systemctl is-active or
    # systemctl list-units --type=scope --state=running
    if ! sudo systemctl is-active "${systemd_scope}" &> /dev/null; then
        echo -e "FAIL: Container ${container_name} failed to start"
        exit 255
    fi

    if [[ "$VERBOSE" = "${FEATURE_ENABLED}" ]]; then
        echo -e "INFO: Container ${container_name} started"
    fi

    local end_time=$(date +%s%3N) # Log end time in milliseconds

    ################################
    #   Adding Extra Packages      #
    ################################
    if [[ ! -z "$EXTRA_PACKAGES_TO_INSTALL" ]]; then
        install_packages "${container_name}"
    fi
}

# Function to stop and remove a container
stop_and_remove_container() {
    local container_name=$1

    if [[ "$VERBOSE" = "${FEATURE_ENABLED}" ]]; then
        echo -e "INFO: stopping and removing container ${container_name}"
    fi

    if podman ps -a --format "{{.Names}}" | grep -q "^${BASE_NAME_FOR_CONTAINERS}"; then
        # Attempt to stop the container gracefully with a custom timeout
        podman stop -t "${TIMEOUT_PODMAN_STOP_CONTAINER}" "$container_name" &> /dev/null
        # Forcefully remove the container if it still exists
        if ! podman rm "$container_name" &> /dev/null; then
	    echo -e "FAIL: to remove container $container_name."
            exit 255
        fi
    else
        if [[ "$VERBOSE" = "${FEATURE_ENABLED}" ]]; then
            echo -e "INFO: no container ${container_name} found to be removed."
	fi
	exit 0
    fi
}

# Function to execute specific action in containers
execute_action_in_containers_in_parallel() {
    local caller_function=$1

    for ((i=1; i<=${NUMBER_OF_CONTAINERS}; i+=BATCH_SIZE)); do
        declare -A pid_to_container
        for ((j=i; j<i+BATCH_SIZE && j<=NUMBER_OF_CONTAINERS; j++)); do
            container_name="${BASE_NAME_FOR_CONTAINERS}$j"
            if [[ "$VERBOSE" = "${FEATURE_ENABLED}" ]]; then
                echo -e "INFO: creating container ${container_name}"
            fi

            ${caller_function} "$container_name" &
            pid_to_container[$!]=$container_name
        done

        wait_and_check_status pid_to_container
    done
}

# Wait and Check the Status
wait_and_check_status() {
    local -n pid_to_container_ref=$1
    local wait_fail=0

    for pid in "${!pid_to_container_ref[@]}"; do
        start_time=$(date +%s)
        while kill -0 "$pid" 2>/dev/null; do
            current_time=$(date +%s)
            elapsed_time=$((current_time - start_time))
            if (( elapsed_time > TIMEOUT_WAITING_ACTIONS )); then
                echo -e "FAIL: Timeout waiting for container ${pid_to_container_ref[$pid]} with PID $pid."
                kill -TERM "$pid" 2>/dev/null
                wait_fail=1
                break
            fi
            sleep 1
        done

        if ! wait $pid; then
            echo -e "FAIL: Failed to run container ${pid_to_container_ref[$pid]} with PID $pid."
            wait_fail=1
        fi
    done

    if [[ $wait_fail -ne 0 ]]; then
        exit 255
    fi
}

cleanup_containers_in_parallel() {
    execute_action_in_containers_in_parallel stop_and_remove_container
}

# Main function to run the script
main() {
    if [ $# -lt 1 ]; then
        echo "Usage: $0 <create|remove> <additional_parameters>"
        exit 1
    fi
    local action=$1

    case $action in
        create)
            local start_time=$(date +%s)
            local msg="creating ${NUMBER_OF_CONTAINERS} containers"

	    # Checking if we need to install stress-ng in the containers
	    is_stress_ng_set_to_run

	    if [ -n "${DISK_EXHAUSTION}" ]; then
		source "$SHARE_DIR/disk-exhaustion"
	        disk_exhaustion_test "disk_exhaustion_test"
	    elif [ -n "${MEMORY_EXHAUSTION}" ]; then
		source "$SHARE_DIR/memory-exhaustion"
	        memory_exhaustion_test
            else
		# No specific test, just creation of containers
	        execute_action_in_containers_in_parallel "create_container"
	    fi

	    # user requested to check if service is enabled
	    if [ -n "${SERVICE_MUST_BE_ENABLED}" ]; then
	        services_must_be_enabled "${container_name}"
            fi

	    if [[ "${SYSTEMD_TIMEOUTSTOPSEC}" = "INFINITY" ]]; then
	        is_systemd_TimeoutStopSec_infinity_works "${container_name}"
	    fi

	    # user requested to check if service is disabled
	    if [ -n "${SERVICE_MUST_BE_DISABLED}" ]; then
	        services_must_be_disabled "${container_name}"
            fi

	    if [ -n "${SELINUX_STATUS_MUST_BE}" ]; then
	        selinux_check_status "${container_name}"
            fi

	    # If users requested to stress the containers, it's time to stress
	    # all containers via stress-ng
	    if [ "$STRESS_NG_SET_TO_RUN" = true ]; then
                # Call system_baseline before running the stress test
                #system_baseline 60

                msg+=" and stressing it"
                run_stress_cmd
                if [ $? -ne 0 ]; then
                    echo "Failed to run stress-ng. Exiting."
                    exit 1
                fi

	        # Find the process IDs (PIDs) of running stress-ng processes,
	        # remove duplicates, and store the output in the stress_pids
	        # variable
	        #stress_pids=$(pgrep stress-ng | uniq | while read pid; do
	        # Extract the second argument of the stress-ng command to
	        # determine the stressor type
	        #stress_type=$(ps -p $pid -o args= | awk '{print $2}')

	        # If no specific stressor type is found, output the PID with
	        # 'stress-ng'
	        #if [ -z "$stress_type" ]; then
	        #  echo "$pid stress-ng"
	        #else
	          # Otherwise, append the stressor type to 'stress-ng' and
	          # output it with the PID
	        #  echo "$pid stress-ng-$stress_type"
	        #fi
	        #done)

                # Run monitor with stress-ng pids
                #stress_pids=$(pgrep stress-ng)
                #MONITOR_INTERVAL_SEC=5
                #system_ongoing ${MONITOR_INTERVAL_SEC} ${stress_pids[@]}
            fi

            ;;
        remove)
            cleanup_containers_in_parallel
            ;;
        *)
            echo "Unknown action: $action"
            exit 1
            ;;
    esac
}

main "$@"