script.sh - master - davinchi2kxx/cgroups-tests

script.sh
134 строки5 KB
davinchi
добавил скрипт
07 дек 2024, 00:09
07 дек 2024, 00:0997bb3a8
#!/bin/bash

# === Проверка аргументов ===
if [[ $# -lt 5 ]]; then
    echo "Использование: $0 <MEMORY_LIMIT> <CPU_LIMIT> <SWAP_ENABLED> <TEST_TYPE> <CPU_THREADS>"
    echo "  MEMORY_LIMIT  - Лимит памяти, например: 50M, 100M, 1G"
    echo "  CPU_LIMIT     - Лимит CPU в процентах, например: 50 (50% от одного ядра)"
    echo "  SWAP_ENABLED  - Разрешить использование подкачки (1 - да, 0 - нет)"
    echo "  TEST_TYPE     - Тип теста: 'memory' или 'cpu'"
    echo "  CPU_THREADS   - Количество потоков для нагрузки CPU (только для 'cpu' теста)"
    exit 1
fi

# === Настройки ===
CGROUP_PATH="/sys/fs/cgroup/testgroup"
MEMORY_LIMIT="$1"
CPU_LIMIT="$2"
SWAP_ENABLED="$3"
TEST_TYPE="$4"
CPU_THREADS="$5"

# === Проверка окружения ===
if ! mount | grep -q "cgroup2"; then
    echo "Ошибка: Cgroups v2 не поддерживаются на этой системе."
    exit 1
fi

if ! command -v python3 >/dev/null 2>&1; then
    echo "Ошибка: Для работы требуется Python 3. Установите его."
    exit 1
fi

if [[ "$TEST_TYPE" == "cpu" && ! $(command -v stress) ]]; then
    echo "Ошибка: Для CPU теста требуется утилита 'stress'. Установите её: sudo apt install stress"
    exit 1
fi

# === Функции ===
function setup_cgroup() {
    echo "Создаем cgroup..."
    mkdir -p "$CGROUP_PATH"

    # Включаем контроль ресурсов
    echo "+memory" > /sys/fs/cgroup/cgroup.subtree_control
    echo "+cpu" > /sys/fs/cgroup/cgroup.subtree_control

    # Устанавливаем лимиты
    echo "$MEMORY_LIMIT" > "$CGROUP_PATH/memory.max"
    if [[ "$SWAP_ENABLED" -eq 0 ]]; then
        echo 0 > "$CGROUP_PATH/memory.swap.max"
    else
        echo "$((2 * $(echo $MEMORY_LIMIT | tr -d M | tr -d G)))M" > "$CGROUP_PATH/memory.swap.max"
    fi
    echo $((CPU_LIMIT * 1000)) > "$CGROUP_PATH/cpu.max"

    echo "Ограничения: память=$MEMORY_LIMIT, CPU=$CPU_LIMIT%, swap=$SWAP_ENABLED"
}

function create_python_script() {
    echo "Создаем Python-скрипт для тестирования памяти..."
    cat <<EOF > consume_memory.py
a = []
try:
    while True:
        a.append(" " * 10**6)  # Запрашиваем память блоками по 1 МБ
except MemoryError:
    print("MemoryError: Превышение лимита памяти")
EOF
}

function run_memory_test() {
    echo "Запуск теста Python для демонстрации OOM..."

    # Запускаем Python-скрипт
    python3 consume_memory.py &
    local PYTHON_PID=$!

    # Устанавливаем обработчик для SIGINT
    trap "echo 'Получен сигнал SIGINT. Завершаю процесс Python...'; kill -TERM $PYTHON_PID 2>/dev/null; exit" SIGINT

    # Перемещаем процесс в cgroup
    echo "$PYTHON_PID" > "$CGROUP_PATH/cgroup.procs"
    echo "Процесс Python ($PYTHON_PID) запущен и ограничен."

    # Ожидаем завершения
    wait "$PYTHON_PID"

    # Сброс обработчика
    trap - SIGINT
}

function run_cpu_test() {
    echo "Запуск теста CPU с использованием stress..."

    # Запускаем stress для нагрузки на CPU
    stress --cpu "$CPU_THREADS" &
    local STRESS_PID=$!

    # Устанавливаем обработчик для SIGINT
    trap "echo 'Получен сигнал SIGINT. Завершаю процесс stress...'; kill -TERM $STRESS_PID 2>/dev/null; exit" SIGINT

    # Перемещаем процесс в cgroup
    echo "$STRESS_PID" > "$CGROUP_PATH/cgroup.procs"
    echo "Процесс stress ($STRESS_PID) запущен и ограничен."

    # Ожидаем завершения
    wait "$STRESS_PID"

    # Сброс обработчика
    trap - SIGINT
}

function cleanup() {
    echo "Очистка cgroup..."
    rm -f consume_memory.py
    rmdir "$CGROUP_PATH"
    echo "Группа $CGROUP_PATH удалена."
}

# === Основной скрипт ===
trap cleanup EXIT

echo "=== Настройка Cgroups и запуск теста ==="
setup_cgroup

if [[ "$TEST_TYPE" == "memory" ]]; then
    create_python_script
    run_memory_test
elif [[ "$TEST_TYPE" == "cpu" ]]; then
    run_cpu_test
else
    echo "Ошибка: Неверный тип теста. Используйте 'memory' или 'cpu'."
    exit 1
fi