/* Aggressive pressure: N hogs each mmap+touch a large region and hold it;
 * concurrently M spammers fire vquotactl("set limit uid") for new uid chunks
 * to maximize the chance the slab allocator blocks (vm_wait) while a cpu
 * holds ac_spin, tripping the lwkt_switch() KASSERT on INVARIANTS kernels. */
#include <sys/types.h>
#include <sys/vfs_quota.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <errno.h>
#include <unistd.h>
#include <signal.h>
#include <sys/wait.h>
#include <sys/mman.h>
#include <libprop/proplib.h>

static int
set_limit_uid(const char *path, uid_t uid, uint64_t limit)
{
	prop_dictionary_t dict, args, res = NULL;
	struct plistref pref;
	int error;

	dict = prop_dictionary_create();
	if (dict == NULL) return ENOMEM;
	prop_dictionary_set_cstring(dict, "command", "set limit uid");
	args = prop_dictionary_create();
	prop_dictionary_set_uint32(args, "uid", uid);
	prop_dictionary_set_uint64(args, "limit", limit);
	prop_dictionary_set(dict, "arguments", args);
	prop_object_release(args);
	error = prop_dictionary_send_syscall(dict, &pref);
	if (error == 0)
		error = vquotactl(path, &pref);
	if (error == 0)
		error = prop_dictionary_recv_syscall(&pref, &res);
	prop_object_release(dict);
	if (res) prop_object_release(res);
	return error;
}

static void
hog(void)
{
	size_t sz = (size_t)512 << 20;	/* 512 MiB each */
	char *p; size_t i;
	for (;;) {
		p = mmap(NULL, sz, PROT_READ | PROT_WRITE,
		    MAP_ANON | MAP_PRIVATE | MAP_NOSYNC, -1, 0);
		if (p == MAP_FAILED) {
			sz >>= 1;
			if (sz < (1u<<20)) _exit(0);
			continue;
		}
		/* touch every page, keep resident */
		for (i = 0; i < sz; i += 4096) p[i] = (char)i;
		/* leak intentionally to hold pages; re-loop to keep growing */
	}
}

/* Each spammer creates a CONTINUOUSLY GROWING set of NEW uid chunks so that
 * RB_FIND in cmd_set_limit_uid always misses -> unode_insert -> kmalloc under
 * ac_spin. Every ~60 new chunks exhausts one slab zone and forces
 * kmem_slab_alloc(M_WAITOK); under memory+swap shortage that calls vm_wait()
 * -> tsleep -> lwkt_switch() -> KASSERT(gd->gd_spinlocks==0) -> panic. */
static void
spam(const char *path, unsigned base)
{
	unsigned uid = base;
	for (;;) {
		(void)set_limit_uid(path, uid, 7ULL);
		uid += 64;	/* stride by one chunk (ACCT_CHUNK_BITS=5) */
	}
}

int
main(int argc, char **argv)
{
	const char *path = argc >= 2 ? argv[1] : "/tmp";
	int nhog = argc >= 3 ? atoi(argv[2]) : 8;
	int nspam = argc >= 4 ? atoi(argv[3]) : 6;
	int i;
	pid_t pid;

	setvbuf(stdout, NULL, _IONBF, 0);
	printf("aggressive: %d hogs + %d spammers on %s (pid %d)\n",
	    nhog, nspam, path, getpid());
	for (i = 0; i < nhog; i++) {
		pid = fork();
		if (pid == 0) { hog(); _exit(0); }
	}
	/* let hogs build memory+swap pressure first */
	sleep(6);
	printf("aggressive: hogs ramped, launching %d new-chunk spammers\n",
	    nspam);
	for (i = 0; i < nspam; i++) {
		pid = fork();
		/* disjoint, ever-growing uid ranges so every call is a NEW chunk */
		if (pid == 0) { spam(path, 400000u + i * 2000000u); _exit(0); }
	}
	/* parent: run for up to ~60s, then kill kids */
	sleep(55);
	printf("aggressive: window elapsed, killing kids\n");
	kill(0, SIGTERM);
	while (waitpid(-1, NULL, 0) > 0) ;
	return 0;
}