#!/bin/bash
# shellcheck disable=SC1091 # Runtime-resolved path; intentionally not followed.
. "$(dirname "$(realpath "${BASH_SOURCE[0]}")")/../local-snap.sh" || return 1
set -eu

# Setup swap if none exist
setup_swap

# Install dependencies
install_deps jq iperf3

# Install LXD
install_lxd

# Configure LXD
lxd init --auto

IMAGE="${TEST_IMG:-ubuntu-minimal-daily:24.04}"

LXCFS_SUPPORTS_SWAP_ACCOUNTING=0
if journalctl --quiet --no-hostname --no-pager --boot=0 --unit=snap.lxd.daemon.service --grep "Kernel supports swap accounting"; then
    LXCFS_SUPPORTS_SWAP_ACCOUNTING=1
fi

# Test
set -x

# Start a container with no limits
echo "=> Start a container with no limits"
lxc launch "${IMAGE}" c1

echo "==> Validate default values"
[ "$(lxc exec c1 -- nproc)" = "$(nproc)" ]
[ "$(lxc exec c1 -- grep ^MemTotal /proc/meminfo)" = "$(grep ^MemTotal /proc/meminfo)" ]
if [ -e "/sys/fs/cgroup/memory/memory.memsw.limit_in_bytes" ] ||
   { [ -e "/sys/fs/cgroup/system.slice/memory.swap.max" ] && [ "${LXCFS_SUPPORTS_SWAP_ACCOUNTING}" -eq 1 ]; }; then
    [ "$(lxc exec c1 -- grep ^SwapTotal /proc/meminfo)" = "$(grep ^SwapTotal /proc/meminfo)" ]
else
    [ "$(lxc exec c1 -- grep ^SwapTotal /proc/meminfo)" = "SwapTotal:             0 kB" ]
fi

if [ -e "/sys/fs/cgroup/cpu" ]; then
    [ "$(lxc exec c1 -- cat /sys/fs/cgroup/cpu/cpu.shares)" = "1024" ]
else
    [ "$(lxc exec c1 -- cat /sys/fs/cgroup/cpu.weight)" = "100" ]
fi

if [ -e "/sys/fs/cgroup/pids" ]; then
    [ "$(lxc exec c1 -- cat /sys/fs/cgroup/pids/pids.max)" = "max" ]
else
    [ "$(lxc exec c1 -- cat /sys/fs/cgroup/pids.max)" = "max" ]
fi

if [ -e "/sys/fs/cgroup/memory" ]; then
    [ "$(lxc exec c1 -- cat /sys/fs/cgroup/memory/memory.swappiness)" = "$(cat /sys/fs/cgroup/memory/memory.swappiness)" ]
fi

if [ -e "/sys/fs/cgroup/net_prio" ]; then
    [ "$(lxc exec c1 -- cat /sys/fs/cgroup/net_prio/net_prio.ifpriomap)" = "$(cat /sys/fs/cgroup/net_prio/net_prio.ifpriomap)" ]
fi

NPROC="$(nproc)"


echo "==> Testing CPU limits"
lxc config set c1 limits.cpu=2
[ "$(lxc exec c1 -- nproc)" = "2" ]

# trying to allocate more CPUs than there are available should silently be capped at NPROC
lxc config set c1 limits.cpu="$((NPROC+2))"
[ "$(lxc exec c1 -- nproc)" = "${NPROC}" ]

# The number 3 comes from `limits.cpu=0,2` which takes the 1st and 3rd host's CPUs
if [ "${NPROC}" -ge 3 ]; then
    lxc config set c1 limits.cpu=0,2
    [ "$(lxc exec c1 -- nproc)" = "2" ]
else
    echo "==> SKIP: Not enough CPUs to test limits.cpu=0,2"
fi

if [ "${NPROC}" -ge 3 ]; then
  lxc config set c1 limits.cpu=2-2
  [ "$(lxc exec c1 -- nproc)" = "1" ]
else
    echo "==> SKIP: Not enough CPUs to test limits.cpu=2-2"
fi

lxc config unset c1 limits.cpu
# XXX: avoid a race (https://github.com/canonical/lxd/issues/12659)
sleep 2
[ "$(lxc exec c1 -- nproc)" = "${NPROC}" ]

lxc config set c1 limits.cpu.priority 5
if [ -e "/sys/fs/cgroup/cpu" ]; then
    [ "$(lxc exec c1 -- cat /sys/fs/cgroup/cpu/cpu.shares)" = "1019" ]
else
    [ "$(lxc exec c1 -- cat /sys/fs/cgroup/cpu.weight)" = "95" ]
fi

lxc config unset c1 limits.cpu.priority
if [ -e "/sys/fs/cgroup/cpu" ]; then
    [ "$(lxc exec c1 -- cat /sys/fs/cgroup/cpu/cpu.shares)" = "1024" ]
else
    [ "$(lxc exec c1 -- cat /sys/fs/cgroup/cpu.weight)" = "100" ]
fi

lxc config set c1 limits.cpu.allowance 10%
if [ -e "/sys/fs/cgroup/cpu" ]; then
    [ "$(lxc exec c1 -- cat /sys/fs/cgroup/cpu/cpu.shares)" = "102" ]
else
    [ "$(lxc exec c1 -- cat /sys/fs/cgroup/cpu.weight)" = "10" ]
fi

lxc config set c1 limits.cpu.allowance 10ms/100ms
if [ -e "/sys/fs/cgroup/cpu" ]; then
    [ "$(lxc exec c1 -- cat /sys/fs/cgroup/cpu/cpu.shares)" = "1024" ]
    [ "$(lxc exec c1 -- cat /sys/fs/cgroup/cpu/cpu.cfs_period_us)" = "100000" ]
    [ "$(lxc exec c1 -- cat /sys/fs/cgroup/cpu/cpu.cfs_quota_us)" = "10000" ]
else
    [ "$(lxc exec c1 -- cat /sys/fs/cgroup/cpu.weight)" = "100" ]
    [ "$(lxc exec c1 -- cat /sys/fs/cgroup/cpu.max)" = "10000 100000" ]
fi
lxc config unset c1 limits.cpu.allowance

echo "==> Testing memory limits"
MEM_LIMIT_MIB=512
lxc config set c1 limits.memory="${MEM_LIMIT_MIB}MiB"
[ "$(lxc exec c1 -- grep ^MemTotal /proc/meminfo)" = "MemTotal:         $((MEM_LIMIT_MIB * 1024)) kB" ]
if [ -e "/sys/fs/cgroup/memory/memory.memsw.limit_in_bytes" ] ||
   { [ -e "/sys/fs/cgroup/system.slice/memory.swap.max" ] && [ "${LXCFS_SUPPORTS_SWAP_ACCOUNTING}" -eq 1 ]; }; then
    [ "$(lxc exec c1 -- grep ^SwapTotal /proc/meminfo)" = "SwapTotal:        $((MEM_LIMIT_MIB * 1024)) kB" ]
else
    [ "$(lxc exec c1 -- grep ^SwapTotal /proc/meminfo)" = "SwapTotal:             0 kB" ]
fi

# ensure that we don't set memory.high when limits.memory.enforce=hard (default value)
if [ -e "/sys/fs/cgroup/system.slice/memory.high" ]; then
    [ "$(lxc exec c1 -- cat /sys/fs/cgroup/memory.max)" = $((MEM_LIMIT_MIB * 1024 * 1024)) ]
    [ "$(lxc exec c1 -- cat /sys/fs/cgroup/memory.high)" = "max" ]
fi

if [ -e "/sys/fs/cgroup/memory" ]; then
    lxc config set c1 limits.memory.swap=false
    [ "$(lxc exec c1 -- cat /sys/fs/cgroup/memory/memory.swappiness)" = "0" ]
    [ "$(lxc exec c1 -- grep ^SwapTotal /proc/meminfo)" = "SwapTotal:             0 kB" ]

    lxc config set c1 limits.memory.swap=true
    lxc config set c1 limits.memory.swap.priority=5
    # swappiness is 70 - 5 (priority)
    [ "$(lxc exec c1 -- cat /sys/fs/cgroup/memory/memory.swappiness)" = "65" ]
fi

if [ -e "/sys/fs/cgroup/memory/memory.memsw.limit_in_bytes" ] ||
   { [ -e "/sys/fs/cgroup/system.slice/memory.swap.max" ] && [ "${LXCFS_SUPPORTS_SWAP_ACCOUNTING}" -eq 1 ]; }; then
    lxc config set c1 limits.memory 128MiB
    [ "$(lxc exec c1 -- grep ^MemTotal /proc/meminfo)" = "MemTotal:         131072 kB" ]

    lxc exec c1 -- mkdir -p /root/dd
    lxc exec c1 -- mount -t tmpfs tmpfs /root/dd -o size=2G
    lxc exec c1 -- dd if=/dev/zero of=/root/dd/blob bs=4M count=16
    dmesg -c >/dev/null 2>&1
    ! lxc exec c1 -- dd if=/dev/zero of=/root/dd/blob bs=4M count=64 || false
    dmesg | grep -F "Memory cgroup out of memory: Killed process"
    dmesg -c >/dev/null 2>&1

    lxc stop c1 --force

    # Wait for post-stop to complete
    sleep 2s

    lxc config set c1 limits.memory.enforce soft
    lxc start c1
    lxc exec c1 -- mkdir -p /root/dd
    lxc exec c1 -- mount -t tmpfs tmpfs /root/dd -o size=2G
    dmesg -c >/dev/null 2>&1
    lxc exec c1 -- dd if=/dev/zero of=/root/dd/blob bs=4M count=64
    dmesg | grep -F "Memory cgroup out of memory: Killed process" && false
    dmesg -c >/dev/null 2>&1

    lxc stop c1 --force

    # Wait for post-stop to complete
    sleep 2s

    lxc config set c1 limits.memory= limits.memory.enforce=
    lxc start c1
fi

echo "==> Testing process limits"
lxc config set c1 limits.processes=2000
if [ -e "/sys/fs/cgroup/pids" ]; then
    [ "$(lxc exec c1 -- cat /sys/fs/cgroup/pids/pids.max)" = "2000" ]
else
    [ "$(lxc exec c1 -- cat /sys/fs/cgroup/pids.max)" = "2000" ]
fi

# limits.network.priority was dropped after 5.0
if [[ -e "/sys/fs/cgroup/net_prio" ]] && [[ "${LXD_SNAP_CHANNEL}" =~ ^5\.0/ ]]; then
    echo "==> Testing limits.network.priority"
    # NOTE: No such thing as net_prio under cgroup2
    lxc config set c1 limits.network.priority=10
    lxc exec c1 -- grep -xF "10" /sys/fs/cgroup/net_prio/net_prio.ifpriomap
fi

lxc restart c1 --force
waitInstanceReady c1
IP=$(lxc network get lxdbr0 ipv4.address | cut -d/ -f1)
lxc exec c1 -- apt-get update
lxc exec c1 --env DEBIAN_FRONTEND=noninteractive -- apt-get install --no-install-recommends --yes iperf3

iperf3 --server > /dev/null 2>&1 &
iperf_pid=$!

# Wait for iperf3 to be ready
if ! waitTCPPort 127.0.0.1 5201; then
    echo "ERROR: iperf3 server did not start listening on 127.0.0.1:5201" >&2
    exit 1
fi

lxc exec c1 -- iperf3 -t 3 -c "${IP}" -J > iperf.json
E1=$(($(jq .end.sum_sent.bits_per_second < iperf.json | cut -d. -f1)/1024/1024))
lxc exec c1 -- iperf3 -t 3 -R -c "${IP}" -J > iperf.json
I1=$(($(jq .end.sum_sent.bits_per_second < iperf.json | cut -d. -f1)/1024/1024))
echo "Throughput before limits: ${E1}Mbps / ${I1}Mbps"

# linux-image-kvm doesn't support setting limits.egress/limits.ingress
if uname -r | grep -- "-kvm$"; then
    echo "The -kvm kernel flavor does not have the required extra modules to test limits.egress/limits.ingress"
else
    lxc config device override c1 eth0 limits.egress=50Mbit limits.ingress=200Mbit
    lxc exec c1 -- iperf3 -t 3 -c "${IP}" -J > iperf.json
    E2=$(($(jq .end.sum_sent.bits_per_second < iperf.json | cut -d. -f1)/1024/1024))
    lxc exec c1 -- iperf3 -t 3 -R -c "${IP}" -J > iperf.json
    I2=$(($(jq .end.sum_sent.bits_per_second < iperf.json | cut -d. -f1)/1024/1024))
    echo "Throughput after limits: ${E2}Mbps / ${I2}Mbps"
    [ "${E2}" -le "50" ]
    [ "${I2}" -le "200" ]

    lxc config device set c1 eth0 limits.egress= limits.ingress=
fi

kill -9 "${iperf_pid}"
wait "${iperf_pid}" 2>/dev/null || true
rm iperf.json

echo "==> Testing disk limits"
if [ -e /sys/fs/cgroup/init.scope/io.weight ]; then
    lxc config set c1 limits.disk.priority=5
    [ "$(lxc exec c1 -- cat /sys/fs/cgroup/io.weight)" = "default 500" ]
elif [ -e /sys/fs/cgroup/blkio/blkio.weight ]; then
    lxc config set c1 limits.disk.priority=5
    [ "$(lxc exec c1 -- cat /sys/fs/cgroup/blkio/blkio.weight)" = "500" ]
fi

lxc config device override c1 root limits.read=10iops limits.write=20iops
if [ -e /sys/fs/cgroup/init.scope/io.max ]; then
    lxc exec c1 -- grep -wF "riops=10 wiops=20" /sys/fs/cgroup/io.max
else
    lxc exec c1 -- grep -w "10$" /sys/fs/cgroup/blkio/blkio.throttle.read_iops_device
    lxc exec c1 -- grep -w "20$" /sys/fs/cgroup/blkio/blkio.throttle.write_iops_device
fi

lxc config device set c1 root limits.read=10MB limits.write=20MB
if [ -e /sys/fs/cgroup/init.scope/io.max ]; then
    lxc exec c1 -- grep -wF "rbps=10000000 wbps=20000000" /sys/fs/cgroup/io.max
else
    lxc exec c1 -- grep -w "10000000$" /sys/fs/cgroup/blkio/blkio.throttle.read_bps_device
    lxc exec c1 -- grep -w "20000000$" /sys/fs/cgroup/blkio/blkio.throttle.write_bps_device
fi

lxc config device set c1 root limits.read= limits.write=

echo "==> Testing the freezer"
lxc pause c1
! lxc exec c1 bash || false
lxc start c1

echo "==> Cleaning up"
lxc delete -f c1

# shellcheck disable=SC2034
FAIL=0
