#!/bin/bash

set -ue
PATH="/usr/sbin:/usr/bin:/sbin:/bin"
export PATH

unset CHROOT NETNS
cleanup() {
    if [ "${CHROOT+x}" ]; then
        schroot --end-session --chroot="$CHROOT"
    fi
    if [ "${NETNS+x}" ]; then
        ip netns del "$NETNS"
    fi
}
trap cleanup EXIT INT TERM

# create new CHROOT
DEB_BUILD_ARCH="$(dpkg-architecture -qDEB_BUILD_ARCH)"
CHROOT="$(schroot -c "unstable-$DEB_BUILD_ARCH-sbuild" -b)"
ROOTDIR="/run/schroot/mount/$CHROOT"

# create new network namespace and place counters to measure network usage
ip netns add "${NETNS:="interimap-benchmark"}"
ip netns exec "$NETNS" nft -f- <<- EOF
	flush ruleset
	table inet filter {
	    counter interimap-in { }
	    counter interimap-out { }
	    counter offlineimap-in { }
	    counter offlineimap-out { }
	    chain input {
	        type filter hook input priority 0
	        iif "lo" ip daddr 127.0.0.1 tcp dport 10143 counter name interimap-in
	        iif "lo" ip daddr 127.0.0.1 tcp dport 10144 counter name offlineimap-in
	    }
	    chain output {
	        type filter hook output priority 0
	        oif "lo" ip saddr 127.0.0.1 tcp sport 10143 counter name interimap-out
	        oif "lo" ip saddr 127.0.0.1 tcp sport 10144 counter name offlineimap-out
	    }
	}
EOF
ip netns exec "$NETNS" ip addr add "127.0.0.1" dev "lo"
ip netns exec "$NETNS" ip link set "lo" up

# resize the partition so it can hold the mail stores (you may want to
# turn swap off too)
mount -o"remount,size=15G" "/run/schroot/mount/$CHROOT/dev/shm"

# install dependencies
schroot --directory="/" --chroot="$CHROOT" -r -- \
    env DEBIAN_FRONTEND="noninteractive" apt-get install \
        --no-install-recommends --assume-yes \
            time dovecot-imapd offlineimap \
            libconfig-tiny-perl libdbd-sqlite3-perl libnet-ssleay-perl \
            libcrypt-urandom-perl procps

# run a command in the chroot
jail() {
    local user="" home=""
    case "${u:-local}" in
        local) user="nobody"; home="/dev/shm/nobody";;
        remote) user="user"; home="/dev/shm/vmail/user";;
    esac
    ip netns exec "$NETNS" \
        schroot --directory="/dev/shm/nobody" --user="nobody" --chroot="$CHROOT" -r \
            -- env -i PATH="/usr/bin:/bin" USER="$user" HOME="$home" "$@"
}

# run a command in the chroot in a monitored fashion
jail_stat() {
    local e U S M P u="local"
    local counters="$ROOTDIR/tmp/counters.$1.json"
    ip netns exec "$NETNS" nft reset counter inet filter "$1-in" >/dev/null
    ip netns exec "$NETNS" nft reset counter inet filter "$1-out" >/dev/null

    jail time --format="%e\\t%U\\t%S\\t%M\\t%P" --output="/tmp/time.$1" \
        -- "$@" >/dev/null || true
    IFS=$'\t' read e U S M P <"$ROOTDIR/tmp/time.$1"

    local i="${NAME:-$1}"
    local a="${i#* }"
    [ "$a" = "$i" ] && a="   " || a=" $a"

    printf "%11s%s" "${i%% *}" "$a"
    printf "  %5.2fs  %6.2fs" "$U" "$S"
    if [ "${IDLE:-n}" = "n" ]; then
        printf "  %5.2fs  %4s" "$e" "$P"
    fi
    printf "  %8s" "${M}k"

    ip netns exec "$NETNS" nft -j list counters | jq ".nftables
        | map(select(.counter) | .counter | { key: .name, value: {packets, bytes} })
        | from_entries" >"$counters"

    local ib ip ob op
    ib="$(bytes "$(jq ".\"$1-in\".bytes"  <"$counters")")"
    ob="$(bytes "$(jq ".\"$1-out\".bytes" <"$counters")")"
    ip="$(_units "$(jq ".\"$1-in\".packets"  <"$counters")" 1000)"
    op="$(_units "$(jq ".\"$1-out\".packets" <"$counters")" 1000)"
    printf "  %8s / %-7s" "$ob" "$ib" # inverse for the client's perspective
    printf "  %8s / %-7s" "$op" "$ip"
    printf "\\n"
}

# display metrics headers
headers() {
    declare -a h=("  user" " system")
    if [ "${IDLE:-n}" = "n" ]; then
        h+=( "  real" " CPU" )
    fi
    h+=( " max RSS" " traffic (in/out) " " packets (in/out) " )
    local x="offlineimap -q" i

    printf "%s" "${x//?/ }"
    for i in "${h[@]}"; do
        printf "  %s" "$i"
    done
    printf "\\n"

    printf "%s" "${x//?/-}"
    for i in "${h[@]}"; do
        printf "  %s" "${i//?/-}"
    done
    printf "\\n"
}

# install Dovecot's "system" configuration and start the server
install -onobody -gnogroup -m0700 --directory \
    "$ROOTDIR/dev/shm/dovecot" \
    "$ROOTDIR/dev/shm/vmail" \
    "$ROOTDIR/dev/shm/nobody"

install -onobody -gnogroup -m0644 \
    "./benchmark/dovecot.conf" \
    "$ROOTDIR/dev/shm/dovecot/config"
jail /usr/sbin/dovecot -c"/dev/shm/dovecot/config"

install -onobody -gnogroup -m0600 /dev/null \
    "$ROOTDIR/dev/shm/dovecot/users"
PASSWORD="$(xxd -l16 -p </dev/urandom)"
printf "%s:%s:::::\\n" "user" "$PASSWORD" \
    >"$ROOTDIR/dev/shm/dovecot/users"

# install user configuration for Dovecot, interimap, and offlineimap
cat >"$ROOTDIR/dev/shm/nobody/.dovecot.conf" <<-EOF
	log_path = /dev/null
	mail_home = /dev/shm/nobody
	mail_location = maildir:~/Maildir
	ssl = no
EOF

install -onobody -gnogroup -Dm0700 --directory \
    "$ROOTDIR/dev/shm/nobody/.config/interimap" \
    "$ROOTDIR/dev/shm/nobody/.local/share"

cat >"$ROOTDIR/dev/shm/nobody/.config/interimap/config" <<-EOF
	database = bench.db

	[local]
	type = tunnel
	command = doveadm -c/dev/shm/nobody/.dovecot.conf exec imap
	null-stderr = YES

	[remote]
	type = imap
	host = 127.0.0.1
	port = 10143
	STARTTLS = no
	username = user
	password = $PASSWORD
EOF

cat >"$ROOTDIR/dev/shm/nobody/.offlineimaprc" <<-EOF
	[general]
	accounts = bench

	[Account bench]
	localrepository = local
	remoterepository = remote

	[Repository local]
	type = Maildir
	localfolders = ~/Maildir2

	[Repository remote]
	type = IMAP
	remotehost = 127.0.0.1
	remotepass = $PASSWORD
	remoteport = 10144
	remoteuser = user
	ssl = no
	starttls = no
	# keep the default (no) as it doesn't seem to work with large mailboxes, perhaps
	# due to https://dovecot.org/pipermail/dovecot/2019-November/117522.html
	#usecompression = yes
EOF

# install interimap's development version
install -oroot -groot -m0755 -Dt "$ROOTDIR/usr/bin" \
    ./interimap ./benchmark/random_maildir.pl
sed -ri "0,/^(use\\s+\lib\\s+)([\"'])[^\"']*\\2\\s*;/ s||#&|" -- "$ROOTDIR/usr/bin/interimap"
install -oroot -groot -m0644 -DT \
    ./lib/Net/IMAP/InterIMAP.pm "$ROOTDIR/usr/share/perl5/Net/IMAP/InterIMAP.pm"

# create a random mail store at mdbox:~/mail.back
prepare() {
    local u="remote" d m n seqs
    local maildir="/dev/shm/vmail/user/maildir"

    clear
    jail rm -rf -- "$maildir" "/dev/shm/vmail/user/mail.back"

    for m in "${!MAILBOXES[@]}"; do
        [ "${m^^[a-z]}" = "INBOX" ] && d="$maildir" || d="$maildir/.$m"
        jail mkdir -p -- "$d"
        # create 20% more; will be deleted afterwards (having only
        # contiguous UIDs might bias the metrics)
        n="${MAILBOXES["$m"]}"
        jail random_maildir.pl "$d" $((n+n/5))
    done

    # convert to mdbox
    jail doveadm -c"/dev/shm/dovecot/config" -omail_location="maildir:~/maildir" \
        sync "mdbox:~/mail.back"
    jail rm -rf -- "$maildir"

    # expunge 20% and purge
    for m in "${!MAILBOXES[@]}"; do
        n="${MAILBOXES["$m"]}"
        seqs="$(shuf -n $((n/5)) -i"1-$n")"
        jail doveadm -c"/dev/shm/dovecot/config" -omail_location="mdbox:~/mail.back" \
            expunge mailbox "$m" "${seqs//$'\n'/,}"
    done
    jail doveadm -c"/dev/shm/dovecot/config" -omail_location="mdbox:~/mail.back" purge
}

# populate a clientn from backup mailstore mdbox:~/mail.back (copied to
# avoid recreating / conversion)
populate() {
    local m u="remote" cmd
    clear

    if [ "${KEEP_BACKUP:-y}" = "n" ]; then
        jail mv -T "/dev/shm/vmail/user/mail.back" "/dev/shm/vmail/user/mail"
    else
        jail cp -aT "/dev/shm/vmail/user/mail.back" "/dev/shm/vmail/user/mail"
    fi

    # force dovecot to index and compute the state, otherwise the first
    # thing to query might be disadvantaged
    jail doveadm -c"/dev/shm/dovecot/config" index "INBOX"
    jail doveadm -c"/dev/shm/dovecot/config" mailbox status "all" "*" >/dev/null

    u="local"
    # initial configuration
    for cmd in "$@"; do
        case "$cmd" in
            interimap) jail interimap --quiet 2>/dev/null;;
            offlineimap) jail offlineimap -u quiet 2>/dev/null;;
            *) exit 1;;
        esac
    done
}

# remove interimap / offlineimap database and mail store (but keep
# mdbox:~/mail.back)
clear() {
    jail rm -rf -- \
        "/dev/shm/vmail/user/mail" \
        "/dev/shm/nobody/.local/share/interimap/bench.db" \
        "/dev/shm/nobody/.offlineimap" \
        "/dev/shm/nobody/Maildir" \
        "/dev/shm/nobody/Maildir2"
}

# pretty print a number in k/M/G/T etc
_units() {
    local n=$(( $1 )) b="$2" s u=""
    [ $b -eq 1024 ] && s="i" || s=""
    while [ ${#n} -gt 4 ]; do
        case "$u" in
            "") u="k";;
            k) u="M";;
            M) u="G";;
            G) u="T";;
            *) break;;
        esac
        n=$((n/b))
    done
    printf "%d%s" "$n" "${u:+$u$s}"
}
bytes() { printf "%sB" "$(_units "$1" 1024)"; }

# generate and deliver a random message
sample_message() {
    local date="$(date +"%s.%N")"
    cat <<-EOF
		From: <sender@example.net>
		To: <recipient@example.net>
		Date: $(date -R -d@"$date")
		Message-ID: <$date@example.net>

	EOF
    xxd -ps -c30 -l2048 /dev/urandom # 4165 bytes
}
deliver() {
    local m="$1" u="remote"
    jail doveadm -c"/dev/shm/dovecot/config" exec dovecot-lda -e -m "$m"
}

# write down markdown title
title() {
    local x="$1" h="$2"
    printf "\\n%s\\n%s\\n" "$h" "${h//?/$x}"
}

# run benchmark for `interimap` / `offlineimap -q` / `offlineimap`:
# populate, run optional actions (such as delivery), then sync again in
# a monitored fashion
run-all() {
    local a cmd u q="" NAME=""
    for cmd in "interimap" "offlineimap -q" "offlineimap"; do
        populate "${cmd%% *}"
        case "${cmd%% *}" in
            interimap) q="--quiet";;
            offlineimap) q="-u quiet";;
            *) exit 1;
        esac
        for a in "$@"; do "$a"; done
        NAME="$cmd" jail_stat $cmd $q 2>/dev/null
    done
}


echo; echo
title "=" "Single mailbox"

cat <<-EOF

	We create a mailbox on the remote server, populate it with a number of
	messages, and synchronize it locally.  We then collect metrics for no-op
	synchronization (i.e., of mailboxes that are already in sync), and
	reconciliation after receiving a *single* message on the remote server.
EOF

# generate a message to be used in *all* "Single mailbox" tests
sample_message >"$ROOTDIR/tmp/msg1"
activity1() {
    deliver "inbox" <"$ROOTDIR/tmp/msg1"
}


declare -A MAILBOXES
for n in 100 1000 10000 100000; do
    title "-" "$n messages"
    MAILBOXES=( ["inbox"]="$n" )
    prepare

    printf "\\n### %s ###\\n\\n" "No-op (in sync)"
    headers
    run-all

    printf "\\n### %s ###\\n\\n" "Reconciliation"
    headers
    run-all activity1
done


m=75
echo; echo
title "=" "$m mailboxes"

cat <<-EOF

	We create $m mailboxes on the remote server, populate them with an equal
	number of messages, and synchronize them locally.  We then collect
	metrics for no-op synchronization (i.e., of mailboxes that are already
	in sync), and reconciliation after the following changes are being
	applied to the remote server:

	  - 3 *new* messages (two on mailbox #2, one on mailbox #3); and
	  - 5 existing messages *EXPUNGEd* (two on mailboxes #3 and #4, one on
	    mailbox #5).
EOF

# generate more messages to be used in *all* "$m mailboxes" tests
sample_message >"$ROOTDIR/tmp/msg2"
sample_message >"$ROOTDIR/tmp/msg3"

activity2() {
    local u="remote"
    deliver "mailbox2" <"$ROOTDIR/tmp/msg1"
    deliver "mailbox2" <"$ROOTDIR/tmp/msg2"
    deliver "mailbox3" <"$ROOTDIR/tmp/msg3"
    # intentionally modify the remote only because not all local backend speak IMAP
    jail doveadm -c"/dev/shm/dovecot/config" expunge mailbox "mailbox3" "1:2"
    jail doveadm -c"/dev/shm/dovecot/config" expunge mailbox "mailbox4" "1,3"
    jail doveadm -c"/dev/shm/dovecot/config" expunge mailbox "mailbox5" "*"
}

for n in 100 1000 10000; do
    title "-" "$n messages per mailbox"

    MAILBOXES=( ["inbox"]="$n" )
    for ((i=2; i<=$m; i++)); do
        MAILBOXES["mailbox$i"]="$n"
    done
    prepare

    printf "\\n### %s ###\\n\\n" "No-op (in sync)"
    headers
    run-all

    printf "\\n### %s ###\\n\\n" "Reconciliation"
    headers
    run-all activity2
done


title "=" "Live synchronization"
timeout=$((6 * 3600))
step=5

MAILBOXES=( ["inbox"]=100000 ["xlarge"]=100000 )
for ((i=0; i<10; i++)); do
    MAILBOXES["large$i"]=10000
done
for ((i=0; i<20; i++)); do
    MAILBOXES["medium$i"]=5000
done
for ((i=0; i<45; i++)); do
    MAILBOXES["small$i"]=2000
done
for ((i=0; i<20; i++)); do
    MAILBOXES["xsmall$i"]=500
done

n=0
for i in "${MAILBOXES[@]}"; do
    n=$(( n + i ))
done

cat <<-EOF

	${#MAILBOXES[@]} mailboxes, $n messages in total:

	  - 2 with 100000 messages;
	  - 10 with 10000 messages;
	  - 20 with 5000 messages;
	  - 45 with 2000 messages; and
	  - 20 with 500 messages.

	The two local mail stores (respectively for [InterIMAP] and
	[OfflineIMAP]) are initially in sync with the remote server, and we keep
	long-running “autorefresh” synchronization processes alive for 6h, with
	updates being regularly applied to the remote server: every $step seconds,

	  - a new message is delivered to a random mailbox with 5% probability
	    (once every $((20*step))s on average);
	  - a random message is EXPUNGEd with 5% probability (once every $((20*step))s on
	    average); and
	  - a random message is marked as seen with 10% probability (once every
	    $((10*step))s on average).

	\`interimap\` is configured to sync every *30s*.  \`offlineimap\` is
	configured to quick sync very *30s*, with a regular sync every *1h*.

EOF

IDLE="y" headers
prepare
KEEP_BACKUP="n" populate "interimap" "offlineimap"

IDLE="y" jail_stat interimap --quiet --watch=30 2>/dev/null &
IDLE="y" jail_stat offlineimap -u quiet -k "Account_bench:autorefresh=0.5" \
    -k "Account_bench:quick=120" 2>/dev/null &

u="remote"
timeout=$(( $(date +%s) + timeout ))
while [ $(date +%s) -lt $timeout ]; do
    n="$(shuf -n1 -i1-100)"
    if [ $n -le 5 ]; then
        # deliver to a random mailbox on the remote
        m="$(shuf -n1 -e -- "${!MAILBOXES[@]}")"
        sample_message | deliver "$m"
    fi
    n="$(shuf -n1 -i1-100)"
    if [ $n -le 5 ]; then
        # expunge a random message on the remote
        read guid uid < <(jail doveadm -c"/dev/shm/dovecot/config" search all | shuf -n1)
        jail doveadm -c"/dev/shm/dovecot/config" expunge mailbox-guid "$guid" uid "$uid"
    fi
    n="$(shuf -n1 -i1-100)"
    if [ $n -le 10 ]; then
        # mark a random message as seen
        read guid uid < <(jail doveadm -c"/dev/shm/dovecot/config" search all | shuf -n1)
        jail doveadm -c"/dev/shm/dovecot/config" flags add "\\Seen" mailbox-guid "$guid" uid "$uid"
    fi
    sleep $step
done

jail pkill -TERM -u"nobody" -s0 interimap
sleep 0.2 # give a chance to print the stats
jail pkill -SIGABRT -u"nobody" -s0 offlineimap
wait