Configure Linux process scheduling and CPU affinity for performance optimization

#!/usr/bin/env bash

set -euo pipefail

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m'

# Global variables
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
BACKUP_DIR="/root/cpu-scheduling-backup-$(date +%Y%m%d-%H%M%S)"

# Usage function
usage() {
    cat << EOF
Usage: $0 [OPTIONS]
Configure Linux process scheduling and CPU affinity for performance optimization

OPTIONS:
    -c, --cores CORES       CPU cores to assign for high-priority processes (default: 2-3)
    -s, --system-cores      CPU cores reserved for system processes (default: 0-1)
    -h, --help             Show this help message

Examples:
    $0                      # Use default CPU affinity settings
    $0 -c "4-7" -s "0-3"   # Custom core assignments
EOF
}

# Logging functions
log_info() { echo -e "${GREEN}[INFO]${NC} $1" >&2; }
log_warn() { echo -e "${YELLOW}[WARN]${NC} $1" >&2; }
log_error() { echo -e "${RED}[ERROR]${NC} $1" >&2; }

# Cleanup function
cleanup() {
    if [ -n "${BACKUP_DIR:-}" ] && [ -d "$BACKUP_DIR" ]; then
        log_warn "Script failed. Backup files available in: $BACKUP_DIR"
        log_warn "To restore: cp -r $BACKUP_DIR/* /"
    fi
}

trap cleanup ERR

# Check if running as root
check_root() {
    if [[ $EUID -ne 0 ]]; then
        log_error "This script must be run as root"
        exit 1
    fi
}

# Detect distribution
detect_distro() {
    if [ -f /etc/os-release ]; then
        . /etc/os-release
        case "$ID" in
            ubuntu|debian) 
                PKG_MGR="apt"
                PKG_INSTALL="apt install -y"
                PKG_UPDATE="apt update"
                ;;
            almalinux|rocky|centos|rhel|ol|fedora) 
                PKG_MGR="dnf"
                PKG_INSTALL="dnf install -y"
                PKG_UPDATE="dnf check-update || true"
                ;;
            amzn) 
                PKG_MGR="yum"
                PKG_INSTALL="yum install -y"
                PKG_UPDATE="yum check-update || true"
                ;;
            *) 
                log_error "Unsupported distribution: $ID"
                exit 1 
                ;;
        esac
        export PKG_MGR PKG_INSTALL PKG_UPDATE
        log_info "Detected distribution: $ID"
    else
        log_error "Cannot detect distribution - /etc/os-release not found"
        exit 1
    fi
}

# Parse command line arguments
parse_args() {
    HIGH_PRIORITY_CORES="2-3"
    SYSTEM_CORES="0-1"

    while [[ $# -gt 0 ]]; do
        case $1 in
            -c|--cores)
                HIGH_PRIORITY_CORES="$2"
                shift 2
                ;;
            -s|--system-cores)
                SYSTEM_CORES="$2"
                shift 2
                ;;
            -h|--help)
                usage
                exit 0
                ;;
            *)
                log_error "Unknown option: $1"
                usage
                exit 1
                ;;
        esac
    done
}

# Create backup directory
create_backup() {
    mkdir -p "$BACKUP_DIR"
    log_info "Created backup directory: $BACKUP_DIR"
}

# Install required packages
install_packages() {
    echo "[1/8] Installing CPU scheduling and monitoring tools..."
    
    $PKG_UPDATE > /dev/null 2>&1
    
    case "$PKG_MGR" in
        apt)
            $PKG_INSTALL util-linux schedtool htop numactl cgroup-tools
            ;;
        dnf|yum)
            $PKG_INSTALL util-linux schedutils htop numactl libcgroup-tools
            ;;
    esac
    
    log_info "Required packages installed successfully"
}

# Display system topology
show_system_info() {
    echo "[2/8] Analyzing system CPU topology..."
    
    log_info "CPU Information:"
    echo "Total CPU cores: $(nproc)"
    echo "Architecture: $(lscpu | grep Architecture | awk '{print $2}')"
    echo "CPU(s): $(lscpu | grep '^CPU(s):' | awk '{print $2}')"
    echo "Thread(s) per core: $(lscpu | grep 'Thread(s) per core' | awk '{print $4}')"
    echo "Core(s) per socket: $(lscpu | grep 'Core(s) per socket' | awk '{print $4}')"
}

# Configure cgroup CPU scheduling
setup_cgroups() {
    echo "[3/8] Setting up CPU cgroups..."
    
    # Backup existing cgroup configurations
    if [ -d /sys/fs/cgroup/cpu ]; then
        cp -r /sys/fs/cgroup/cpu "$BACKUP_DIR/" 2>/dev/null || true
    fi
    
    # Create cgroup directories
    mkdir -p /sys/fs/cgroup/cpu/high-priority
    mkdir -p /sys/fs/cgroup/cpu/low-priority
    
    # Configure CPU shares (high priority gets more CPU time)
    echo 2048 > /sys/fs/cgroup/cpu/high-priority/cpu.shares
    echo 512 > /sys/fs/cgroup/cpu/low-priority/cpu.shares
    
    # Set CPU quota for high priority group (200% = 2 cores)
    echo 200000 > /sys/fs/cgroup/cpu/high-priority/cpu.cfs_quota_us
    echo 100000 > /sys/fs/cgroup/cpu/high-priority/cpu.cfs_period_us
    
    # Set CPU affinity for cgroups
    echo "$HIGH_PRIORITY_CORES" > /sys/fs/cgroup/cpuset/high-priority/cpuset.cpus 2>/dev/null || true
    echo "$SYSTEM_CORES" > /sys/fs/cgroup/cpuset/low-priority/cpuset.cpus 2>/dev/null || true
    
    log_info "CPU cgroups configured successfully"
}

# Create systemd override configurations
configure_systemd_overrides() {
    echo "[4/8] Configuring systemd CPU scheduling overrides..."
    
    # Example configuration for common services
    services=("sshd" "chronyd" "systemd-logind")
    
    for service in "${services[@]}"; do
        if systemctl list-unit-files "$service.service" > /dev/null 2>&1; then
            override_dir="/etc/systemd/system/$service.service.d"
            mkdir -p "$override_dir"
            
            # Backup existing override if it exists
            if [ -f "$override_dir/cpu-affinity.conf" ]; then
                cp "$override_dir/cpu-affinity.conf" "$BACKUP_DIR/$service-cpu-affinity.conf"
            fi
            
            cat > "$override_dir/cpu-affinity.conf" << EOF
[Service]
# System service CPU affinity - use system cores
CPUAffinity=$SYSTEM_CORES
Nice=0
CPUSchedulingPolicy=0
EOF
            chmod 644 "$override_dir/cpu-affinity.conf"
        fi
    done
    
    log_info "Systemd CPU scheduling overrides configured"
}

# Create sample high-priority service template
create_sample_service() {
    echo "[5/8] Creating sample high-priority service template..."
    
    cat > "/etc/systemd/system/high-priority-app.service" << EOF
[Unit]
Description=High Priority Application Template
After=network.target
Documentation=man:systemd.service(5)

[Service]
Type=simple
User=root
Group=root
ExecStart=/bin/sleep 3600

# CPU Scheduling Configuration
Nice=-10
CPUAffinity=$HIGH_PRIORITY_CORES
CPUSchedulingPolicy=1
CPUSchedulingPriority=50

# Resource Limits
MemoryMax=1G
CPUQuota=150%

# Security
PrivateTmp=true
ProtectSystem=strict
NoNewPrivileges=true

[Install]
WantedBy=multi-user.target
EOF
    
    chmod 644 "/etc/systemd/system/high-priority-app.service"
    log_info "Sample high-priority service template created"
}

# Create CPU affinity management scripts
create_management_scripts() {
    echo "[6/8] Creating CPU affinity management scripts..."
    
    # Create script directory
    mkdir -p /usr/local/bin
    
    # CPU affinity helper script
    cat > "/usr/local/bin/cpu-affinity-tool" << 'EOF'
#!/bin/bash
# CPU Affinity Management Tool

case "$1" in
    show)
        echo "Current CPU topology:"
        lscpu | grep -E "CPU\(s\)|Thread|Core|Socket"
        echo
        echo "Process CPU affinity (top 10 by CPU usage):"
        ps -eo pid,ppid,cmd,nice,psr --sort=-pcpu | head -11
        ;;
    set)
        if [ $# -lt 3 ]; then
            echo "Usage: $0 set <pid> <cpu_list>"
            exit 1
        fi
        taskset -cp "$3" "$2"
        ;;
    monitor)
        watch -n 1 'ps -eo pid,ppid,cmd,nice,psr --sort=-pcpu | head -20'
        ;;
    *)
        echo "Usage: $0 {show|set|monitor}"
        echo "  show     - Display current CPU topology and process affinity"
        echo "  set      - Set CPU affinity for a process"
        echo "  monitor  - Monitor process CPU usage and affinity"
        exit 1
        ;;
esac
EOF
    
    chmod 755 "/usr/local/bin/cpu-affinity-tool"
    
    log_info "CPU affinity management scripts created"
}

# Apply configurations
apply_configurations() {
    echo "[7/8] Applying CPU scheduling configurations..."
    
    # Reload systemd daemon
    systemctl daemon-reload
    
    # Set kernel parameters for better scheduling
    cat > "/etc/sysctl.d/99-cpu-scheduling.conf" << EOF
# CPU scheduling optimizations
kernel.sched_latency_ns = 6000000
kernel.sched_min_granularity_ns = 2000000
kernel.sched_wakeup_granularity_ns = 3000000
kernel.sched_migration_cost_ns = 500000
kernel.sched_nr_migrate = 32

# NUMA balancing
kernel.numa_balancing = 1
EOF
    
    chmod 644 "/etc/sysctl.d/99-cpu-scheduling.conf"
    sysctl -p "/etc/sysctl.d/99-cpu-scheduling.conf" > /dev/null
    
    log_info "CPU scheduling configurations applied"
}

# Verify installation
verify_installation() {
    echo "[8/8] Verifying CPU scheduling configuration..."
    
    # Check if tools are installed
    local tools=("taskset" "schedtool" "htop" "numactl")
    for tool in "${tools[@]}"; do
        if command -v "$tool" > /dev/null 2>&1; then
            log_info "✓ $tool is installed"
        else
            log_warn "✗ $tool is not installed or not in PATH"
        fi
    done
    
    # Check cgroup setup
    if [ -d "/sys/fs/cgroup/cpu/high-priority" ]; then
        log_info "✓ CPU cgroups are configured"
    else
        log_warn "✗ CPU cgroups not found"
    fi
    
    # Check systemd overrides
    if [ -f "/etc/systemd/system/high-priority-app.service" ]; then
        log_info "✓ Sample high-priority service created"
    fi
    
    # Check management script
    if [ -x "/usr/local/bin/cpu-affinity-tool" ]; then
        log_info "✓ CPU affinity management tool installed"
    fi
    
    echo
    log_info "CPU scheduling and affinity configuration completed!"
    echo
    echo "Next steps:"
    echo "1. Use 'cpu-affinity-tool show' to view current CPU topology"
    echo "2. Use 'cpu-affinity-tool monitor' to monitor process CPU usage"
    echo "3. Customize service CPU affinity in /etc/systemd/system/SERVICE.service.d/"
    echo "4. Restart services to apply CPU affinity: systemctl restart SERVICE"
    echo "5. Use 'taskset -cp CORES PID' to set CPU affinity for running processes"
}

# Main function
main() {
    parse_args "$@"
    check_root
    detect_distro
    create_backup
    install_packages
    show_system_info
    setup_cgroups
    configure_systemd_overrides
    create_sample_service
    create_management_scripts
    apply_configurations
    verify_installation
}

main "$@"