start-dlp-us-sv-win1

#!/bin/bash
set -e


#################
# Configuration #
#################

#VM_DOMAIN="local"
VM_DOMAIN="dlp-us-sv-win1"

#
# List of devices that need to be passed through to the VM
#
# 01:00.0 VGA compatible controller: NVIDIA Corporation GP104 [GeForce GTX 1070] (rev a1)
# 01:00.1 Audio device: NVIDIA Corporation GP104 High Definition Audio Controller (rev a1)
# 00:1f.3 Audio device: Intel Corporation 200 Series PCH HD Audio
# 00:14.0 USB device: Intel Corporation 200 Series/Z370 Chipset Family USB 3.0 xHCI Controller
#
# Pass through but do not force through stub driver
# 0b:00.0 Non-Volatile memory controller: Phison Electronics Corporation E12 NVMe Controller (rev 01)
#
PCI_DEVICEIDS="0000:01:00.0 0000:01:00.1 0000:00:1f.3 0000:00:14.0"

# Blacklised services and modules are stopped before devices are passed through and restarted
# after the VM is stopped
BLACKLISTED_SERVICES="gdm bumblebeed"
BLACKLISTED_MODULES="nvidia_drm nvidia_modeset nvidia_uvm nvidia snd_hda_intel xhci_pci"
#BLACKLISTED_MODULES="nvidia_modeset nvidia"

# Ensured services and modules are started before devices are passed through, but are not
# stopped after devices are reattached
ENSURE_SERVICES="libvirtd"
#ENSURE_MODULES="virtio virtio_net virtio_blk virtio_console virtio_balloon virtio_ring 9pnet_virtio virtio_crypto iptable_nat ip_conntrack_netbios_ns nf_nat nf_tables nft_masq_ipv4 nft_masq_ipv6 nft_chain_nat_ipv6 nft_chain_nat_ipv4"
ENSURE_MODULES="vfio-pci virtio virtio_net virtio_blk virtio_console virtio_balloon virtio_ring 9pnet_virtio virtio_crypto iptable_nat ip_conntrack_netbios_ns nf_nat nf_tables nft_masq nft_chain_nat"

# Start services and modules are started before devices are passed through and stopped after
# devices are reattached
START_SERVICES="nmb smb" # winbind"
START_MODULES="virtio-pci"

# Number of HugePages to request.
# Should be approximately memory allocation divided by 1GiB
HUGEPAGE_COUNT=16

# Cpu Isolation Config
TOTAL_CORES='0-7'
TOTAL_CORES_MASK=FF             # 0-7, bitmask 0b11111111
HOST_CORES='0,4'                # Cores reserved for host
HOST_CORES_MASK=11              # 0,4, bitmask 0b00010001
VIRT_CORES='1-3,5-7'            # Cores reserved for virtual machine(s). This prevents other processes from touching these cores. This is not the same as the cpu list of the VM

################
# DO NOT TOUCH #
################

# Check root access
if [[ $EUID -ne 0 ]]; then
    echo "this script must be run as the super-user"
    exit 1
fi

case "$1" in
    "")
        ;&

    start)
        echo "[INFO] Verifying Configuration"

        if [[ ! $(virsh domstate ${VM_DOMAIN}) =~ "shut off" ]]; then
            echo "[ERROR] The virtual machine ${VM_DOMAIN} is already started"
            exit 1
        fi

        echo "[INFO] Rescanning PCI devices in case prior run failed"
        echo 1 > /sys/bus/pci/rescan
        sleep 0.5

        echo "[INFO] Verifying that devices exist"
        for i in ${PCI_DEVICEIDS}; do
            if [[ ! -e "/sys/bus/pci/devices/${dev}" ]]; then
                echo "[ERROR] Device ${dev} was not found"
                echo "[ERROR] /sys/bus/pci/devices/${dev} was not found"
                exit 1
            fi
        done

        echo "[INFO] Verifying that kernel modules exist"
        modprobe -n ${ENSURE_MODULES} ${START_MODULES}

        echo "[INFO] Configuration Verified"
        ;&

    prep-services)
        echo "[INFO] Preparing services and drivers"

        ! ps axco command | grep gnome | xargs killall 2>/dev/null
        ! ps axco command | grep gdm- | xargs killall 2>/dev/null
        ! virsh net-start default 2>/dev/null

        sleep 3

        for i in ${BLACKLISTED_SERVICES}; do
            echo "[INFO] -> Stopping service ${i}"
            systemctl is-active --quiet ${i} && systemctl stop ${i}
        done

        echo "[INFO] -> Unloading kernel module ${BLACKLISTED_MODULES}"
        modprobe -r ${BLACKLISTED_MODULES}

        for i in ${ENSURE_SERVICES} ${START_SERVICES}; do
            echo "[INFO] -> Starting service ${i}"
            systemctl start ${i}
        done

        for i in ${ENSURE_MODULES} ${START_MODULES}; do
            echo "[INFO] -> Loading kernel module ${i}"
            modprobe ${i}
        done

        sleep 0.25

        ;&

    passthrough-devices)
        echo "[INFO] Passing through PCI devices"
        
        # Trying to unbind the GPU results in general protection faults.
        for dev in ${PCI_DEVICEIDS}
        do
            for it in {1..5}; do
                if [[ ! -e "/sys/bus/pci/drivers/vfio-pci/${dev}" ]]; then
                    vendor=$(cat /sys/bus/pci/devices/${dev}/vendor)
                    device=$(cat /sys/bus/pci/devices/${dev}/device)
                    echo "[INFO] -> Passing through device ${dev} (${vendor} ${device}) Attempt ${it}"

                    echo "[INFO] ----> Issuing PCI hotplug remove command"
                    echo 1 > /sys/bus/pci/devices/${dev}/remove
                    while [[ -e "/sys/bus/pci/devices/${dev}" ]]; do sleep 0.1; done

                    echo "[INFO] ----> Rescan PCI bus to rediscover device"
                    echo 1 > /sys/bus/pci/rescan
                    while [[ ! -e "/sys/bus/pci/devices/${dev}" ]]; do sleep 0.1; done

                    echo "[INFO] ----> Adding device to vfio-pci id list"
                    ! echo ${vendor} ${device} > /sys/bus/pci/drivers/vfio-pci/new_id
                    sleep 2

                    echo "[INFO] ----> Verifying device is attached to vfio-pci stub driver"
                    if [[ ! -e "/sys/bus/pci/drivers/vfio-pci/${dev}" ]]; then
                        if [[ "${it}" -eq 5 ]]; then
                            echo "[ERROR] ----> DEVICE IS NOT ATTACHED TO VFIO-PCI"
                            exit 1
                        else
                            echo "[WARN] ----> DEVICE IS NOT ATTACHED TO VFIO-PCI"
                            sleep 5
                            continue
                        fi
                    fi
                    break
                else
                    echo "[INFO] -> Device ${dev} is already attached to vfio-pci"
                    break
                fi
            done
        done
        ;&

    setup-hugepages)
        echo "[INFO] Dropping kernel caches & compacting memory"
        echo 3 > /proc/sys/vm/drop_caches
        echo 1 > /proc/sys/vm/compact_memory
        sleep 3

        echo "[INFO] Attempting to allocate ${HUGEPAGE_COUNT} HugePages"
        echo ${HUGEPAGE_COUNT} > /sys/kernel/mm/hugepages/hugepages-1048576kB/nr_hugepages

        ALLOC_PAGES=`cat /sys/kernel/mm/hugepages/hugepages-1048576kB/nr_hugepages`
        echo "[INFO] Tried to allocate hugepages. Got pages ${ALLOC_PAGES} / ${HUGEPAGE_COUNT}"

        if [ "${ALLOC_PAGES}" -ne "${HUGEPAGE_COUNT}" ]
        then
            echo "[ERROR] Unable to allocate sufficient hugepages"
            echo 0 > /proc/sys/vm/nr_hugepages
            exit 1
        fi
        
        ;&

    isolate-cores)
        echo "[INFO] Reserving CPUs ${VIRT_CORES} for VM"
        systemctl set-property --runtime -- user.slice AllowedCPUs=$HOST_CORES     
        systemctl set-property --runtime -- system.slice AllowedCPUs=$HOST_CORES
        systemctl set-property --runtime -- init.scope AllowedCPUs=$HOST_CORES

        # the kernel's dirty page writeback mechanism uses kthread workers. They introduce
        # massive arbitrary latencies when doing disk writes on the host and aren't
        # migrated by systemd. Restrict the workqueue to use only cpu 0.
        echo $HOST_CORES_MASK > /sys/bus/workqueue/devices/writeback/cpumask
        echo 0 > /sys/bus/workqueue/devices/writeback/numa
        ;&

    start-vm)
        echo "[INFO] Starting virtual machine (${VM_DOMAIN})"
        #ulimit -r 99
        #sysctl -w kernel.sched_rt_runtime_us=-1
        sysctl vm.stat_interval=120
        sysctl -w kernel.watchdog=0
        
        virsh start ${VM_DOMAIN}
        sleep 15

        # Set the nice of all vCPUs
        for f in /sys/fs/cgroup/machine.slice/machine-qemu*.scope/cgroup.threads
        do
            echo "[INFO] libvirt-qemu nice: Setting $($f)'s nice level to -1"
            for pid in $(cat $f)
            do
                renice -n "-1" -p "$pid"
            done
        done

        ;&

    wait-vm)
        #read -p "[WAIT] Press enter when VM has completed"

        echo "[INFO] Waiting for virtual machine to exit"
        while true; do
            [[ $(virsh domstate ${VM_DOMAIN}) =~ "shut off" ]] && break
            sleep 15
        done
        ;&

    restore-settings)
        #ulimit -r 0
        #sysctl -w kernel.sched_rt_runtime_us=950000
        sysctl vm.stat_interval=1
        sysctl -w kernel.watchdog=1
        ;&

    unisolate-cores)
        echo "[INFO] Unreserving CPUs ${VIRT_CORES}"
        systemctl set-property --runtime -- user.slice AllowedCPUs=$TOTAL_CORES 
        systemctl set-property --runtime -- system.slice AllowedCPUs=$TOTAL_CORES
        systemctl set-property --runtime -- init.scope AllowedCPUs=$TOTAL_CORES

        echo $TOTAL_CORES_MASK > /sys/bus/workqueue/devices/writeback/cpumask
        echo 1 > /sys/bus/workqueue/devices/writeback/numa
        ;&

    teardown-hugepages)
        echo "[INFO] Attempting to deallocate ${HUGEPAGE_COUNT} HugePages"
        echo 0 > /sys/kernel/mm/hugepages/hugepages-1048576kB/nr_hugepages

        ;&

    restore-modules)
        for i in ${START_MODULES}; do
            echo "[INFO] -> Unloading kernel module ${i}"
            ! modprobe -r ${i} 2>/dev/null
        done

        for i in ${BLACKLISTED_MODULES}; do
            echo "[INFO] -> Trying to load kernel module ${i}"
            ! modprobe ${i}
        done
    
        ;&
    restore-devices)
        for dev in ${PCI_DEVICEIDS}
        do
            for it in {1..5}; do
                if [[ -e "/sys/bus/pci/drivers/vfio-pci/${dev}" ]]; then
                    vendor=$(cat /sys/bus/pci/devices/${dev}/vendor)
                    device=$(cat /sys/bus/pci/devices/${dev}/device)
                    echo "[INFO] -> Returning device ${dev} to host (${vendor} ${device}) Attempt ${it}"

                    echo "[INFO] ----> Removing device from vfio-pci id list"
                    ! echo ${vendor} ${device} > /sys/bus/pci/drivers/vfio-pci/remove_id
                    sleep 2

                    echo "[INFO] ----> Issuing PCI hotplug remove command"
                    echo 1 > /sys/bus/pci/devices/${dev}/remove
                    while [[ -e "/sys/bus/pci/devices/${dev}" ]]; do sleep 0.1; done

                    echo "[INFO] ----> Rescan PCI bus to rediscover device"
                    echo 1 > /sys/bus/pci/rescan
                    while [[ ! -e "/sys/bus/pci/devices/${dev}" ]]; do sleep 0.1; done

                    echo "[INFO] ----> Verifying device is no longer attached to vfio-pci stub driver"
                    if [[ -e "/sys/bus/pci/drivers/vfio-pci/${dev}" ]]; then
                        if [[ "${it}" -eq 5 ]]; then
                            echo "[ERROR] ----> DEVICE IS STILL ATTACHED TO VFIO-PCI"
                            echo "Run $0 restore-devices again"
                            exit 1
                        else
                            echo "[WARN] ----> DEVICE IS STILL ATTACHED TO VFIO-PCI"
                            sleep 5
                            continue
                        fi
                    fi
                    break
                else
                    echo "[INFO] -> Device ${dev} is not attached to vfio-pci"
                    break
                fi
            done
        done

        # Try reloading modules again after hardware dependencies have loaded
        for i in ${BLACKLISTED_MODULES}; do
            echo "[INFO] -> Trying to load kernel module ${i}"
            ! modprobe ${i}
        done
        ;&

    restore-services)
        echo "[INFO] Restoring services and drivers"

        for i in ${START_SERVICES}; do
            echo "[INFO] -> Stopping service ${i}"
            systemctl is-active --quiet ${i} && systemctl stop ${i}
        done

        for i in ${BLACKLISTED_SERVICES}; do
            echo "[INFO] -> Starting service ${i}"
            ! systemctl start ${i}
        done

        sleep 0.25
        ;&

    restore-services)
        echo "[INFO] Done!"
        ;&
esac