repro-invlpgb.c

// SPDX-License-Identifier: MIT
#define _GNU_SOURCE

/*
 * Build:
 *   gcc -O2 -Wall -Wextra -pthread -static -o repro-server-tlb repro-server-tlb.c
 *
 * Example run:
 *   ./repro-server-tlb --batch --rounds 20 --jobs 32 -d 5 -w 8 -m 2 -s 512 -q
 *
 * Expected result on a good kernel/CPU:
 *   all rounds complete without output containing "CORRUPTION" or
 *   "UNEXPECTED SIGSEGV".
 *
 * Failure mode:
 *   CORRUPTION slot=... gen=... off=... got={...} want={...}
 *
 * The program protects each slot with a pthread rwlock. Mutators hold the
 * write lock while unmapping/remapping the slot and filling it with a marker;
 * readers hold the read lock while checking that marker. A mismatch means a
 * load from one virtual offset returned bytes that do not match the marker
 * written for that offset/generation. In the observed failure, the returned
 * bytes usually decode as a marker written for a different offset, which is
 * consistent with a stale or wrong translation after virtual-address reuse.
 */

#include <errno.h>
#include <inttypes.h>
#include <pthread.h>
#include <sched.h>
#include <signal.h>
#include <stdatomic.h>
#include <stdbool.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/mman.h>
#include <sys/random.h>
#include <sys/wait.h>
#include <time.h>
#include <ucontext.h>
#include <unistd.h>

#define SLOT_PAGES 16UL
#define GUARD_PAGES 1UL
#define CHECK_STRIDE 4096UL
#define FILL_STRIDE 64UL
#define CHECKSUM_CONST 0x9e3779b97f4a7c15ULL

struct opts {
	int workers;
	int mutators;
	int duration_s;
	int slots;
	int rounds;
	int jobs;
	bool quiet;
	bool batch;
};

struct marker {
	uint64_t cookie;
	uint64_t slot;
	uint64_t gen;
	uint64_t off;
	uint64_t checksum;
};

struct slot {
	unsigned char *base;
	unsigned char *data;
	size_t len;
	atomic_uint gen;
	pthread_rwlock_t lock;
};

struct arg {
	int id;
	int cpu;
};

static struct opts opts = {
	.duration_s = 300,
	.slots = 4096,
	.rounds = 1,
	.jobs = 1,
};
static long page_size;
static struct slot *slots;
static int *cpus;
static int nr_cpus;
static uint64_t cookie;
static atomic_bool running = ATOMIC_VAR_INIT(true);
static atomic_ulong ops = ATOMIC_VAR_INIT(0);
static atomic_ulong corruptions = ATOMIC_VAR_INIT(0);

static uint64_t nsec_now(void)
{
	struct timespec ts;

	clock_gettime(CLOCK_MONOTONIC, &ts);
	return (uint64_t)ts.tv_sec * 1000000000ULL + ts.tv_nsec;
}

static void die(const char *what)
{
	fprintf(stderr, "%s: %s\n", what, strerror(errno));
	exit(1);
}

static uint64_t rnd64(uint64_t *s)
{
	uint64_t x = *s;

	x ^= x >> 12;
	x ^= x << 25;
	x ^= x >> 27;
	*s = x;
	return x * 0x2545F4914F6CDD1DULL;
}

static uint64_t marker_checksum(uint64_t c, uint64_t slot, uint64_t gen, uint64_t off)
{
	return c ^ (slot * 0xbf58476d1ce4e5b9ULL) ^
	       (gen * 0x94d049bb133111ebULL) ^
	       (off * 0xd6e8feb86659fd93ULL) ^ CHECKSUM_CONST;
}

static struct marker make_marker(int slot, unsigned gen, size_t off)
{
	struct marker m = {
		.cookie = cookie,
		.slot = (uint64_t)slot,
		.gen = gen,
		.off = off,
	};

	m.checksum = marker_checksum(m.cookie, m.slot, m.gen, m.off);
	return m;
}

static bool marker_equal(const struct marker *a, const struct marker *b)
{
	return a->cookie == b->cookie && a->slot == b->slot &&
	       a->gen == b->gen && a->off == b->off &&
	       a->checksum == b->checksum;
}

static bool marker_checksum_valid(const struct marker *m)
{
	return m->checksum == marker_checksum(m->cookie, m->slot, m->gen, m->off);
}

static void print_marker(FILE *f, const struct marker *m)
{
	fprintf(f, "{cookie=0x%016" PRIx64 " slot=%" PRIu64
		   " gen=%" PRIu64 " off=%" PRIu64 " checksum=0x%016" PRIx64
		   " valid=%d}",
		m->cookie, m->slot, m->gen, m->off, m->checksum,
		marker_checksum_valid(m));
}

static void init_cookie(void)
{
	uint64_t seed;

	if (getrandom(&seed, sizeof(seed), 0) != sizeof(seed))
		seed = nsec_now() ^ ((uint64_t)getpid() << 32) ^ (uintptr_t)&seed;

	/* Avoid zero, which makes uninitialised pages look slightly less noisy. */
	cookie = seed ? seed : 0xa5a5a5a55a5a5a5aULL;
}

static void pin_cpu(int cpu)
{
	cpu_set_t set;

	CPU_ZERO(&set);
	CPU_SET(cpu, &set);
	if (sched_setaffinity(0, sizeof(set), &set))
		fprintf(stderr, "warning: sched_setaffinity cpu %d: %s\n",
			cpu, strerror(errno));
}

static void load_cpus(void)
{
	cpu_set_t set;
	int max = sysconf(_SC_NPROCESSORS_CONF);

	if (max <= 0)
		max = 1024;
	cpus = calloc(max, sizeof(*cpus));
	if (!cpus)
		die("calloc cpus");
	CPU_ZERO(&set);
	if (sched_getaffinity(0, sizeof(set), &set))
		die("sched_getaffinity");
	for (int i = 0; i < max; i++)
		if (CPU_ISSET(i, &set))
			cpus[nr_cpus++] = i;
	if (!nr_cpus) {
		fprintf(stderr, "no CPUs in affinity mask\n");
		exit(1);
	}
	if (!opts.workers)
		opts.workers = nr_cpus < 256 ? nr_cpus : 256;
	if (!opts.mutators)
		opts.mutators = nr_cpus / 8 ? nr_cpus / 8 : 1;
	if (opts.mutators > 32)
		opts.mutators = 32;
}

static void segv_handler(int sig, siginfo_t *si, void *ctx)
{
	(void)sig;
	ucontext_t *uc = ctx;
	unsigned long ip = 0, err = 0;

#if defined(__x86_64__)
	ip = uc->uc_mcontext.gregs[REG_RIP];
	err = uc->uc_mcontext.gregs[REG_ERR];
#endif
	dprintf(2, "UNEXPECTED SIGSEGV addr=%p ip=0x%lx err=0x%lx ops=%lu corruptions=%lu\n",
		si->si_addr, ip, err, atomic_load(&ops), atomic_load(&corruptions));
	_exit(139);
}

static void install_segv(void)
{
	struct sigaction sa = {0};

	sa.sa_sigaction = segv_handler;
	sa.sa_flags = SA_SIGINFO | SA_RESETHAND;
	sigemptyset(&sa.sa_mask);
	if (sigaction(SIGSEGV, &sa, NULL))
		die("sigaction");
}

static void cpuid_invlpgb(void)
{
#if defined(__x86_64__) || defined(__i386__)
	uint32_t a, b, c, d, maxext;

	__asm__ volatile("cpuid" : "=a"(maxext), "=b"(b), "=c"(c), "=d"(d) :
			 "0"(0x80000000U), "2"(0));
	if (maxext >= 0x80000008U) {
		__asm__ volatile("cpuid" : "=a"(a), "=b"(b), "=c"(c), "=d"(d) :
				 "0"(0x80000008U), "2"(0));
		printf("cpuid: invlpgb=%u ebx=0x%08x edx=0x%08x count_max=%u\n",
		       (b >> 3) & 1, b, d, (d & 0xffff) + 1);
	}
#endif
}

static void fill_slot(int idx, unsigned gen)
{
	struct slot *s = &slots[idx];

	for (size_t off = 0; off + sizeof(struct marker) <= s->len; off += FILL_STRIDE) {
		struct marker m = make_marker(idx, gen, off);

		memcpy(s->data + off, &m, sizeof(m));
	}
}

static void read_marker(unsigned char *addr, struct marker *m)
{
	volatile uint64_t *p = (volatile uint64_t *)addr;

	m->cookie = p[0];
	m->slot = p[1];
	m->gen = p[2];
	m->off = p[3];
	m->checksum = p[4];
}

static void check_slot(int idx, unsigned gen)
{
	struct slot *s = &slots[idx];

	for (size_t off = 0; off + sizeof(struct marker) <= s->len; off += CHECK_STRIDE) {
		struct marker got;
		struct marker want = make_marker(idx, gen, off);

		read_marker(s->data + off, &got);
		if (!marker_equal(&got, &want)) {
			atomic_fetch_add(&corruptions, 1);
			fprintf(stderr, "CORRUPTION slot=%d gen=%u off=%zu got=", idx, gen, off);
			print_marker(stderr, &got);
			fprintf(stderr, " want=");
			print_marker(stderr, &want);
			fprintf(stderr, " ops=%lu\n", atomic_load(&ops));
			_exit(140);
		}
	}
}

static void init_slots(void)
{
	size_t data_len = SLOT_PAGES * (size_t)page_size;
	size_t total = data_len + 2 * GUARD_PAGES * (size_t)page_size;

	slots = calloc(opts.slots, sizeof(*slots));
	if (!slots)
		die("calloc slots");
	for (int i = 0; i < opts.slots; i++) {
		unsigned char *p = mmap(NULL, total, PROT_NONE,
				       MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
		if (p == MAP_FAILED)
			die("mmap slot");
		slots[i].base = p;
		slots[i].data = p + GUARD_PAGES * (size_t)page_size;
		slots[i].len = data_len;
		if (mprotect(slots[i].data, data_len, PROT_READ | PROT_WRITE))
			die("mprotect init");
		atomic_store(&slots[i].gen, 1);
		if (pthread_rwlock_init(&slots[i].lock, NULL)) {
			errno = EINVAL;
			die("pthread_rwlock_init");
		}
		fill_slot(i, 1);
	}
}

static void *worker(void *argp)
{
	struct arg *a = argp;
	uint64_t seed = 0x9e3779b97f4a7c15ULL ^ (uint64_t)a->id;

	pin_cpu(a->cpu);
	while (atomic_load_explicit(&running, memory_order_acquire)) {
		int idx = rnd64(&seed) % opts.slots;
		struct slot *s = &slots[idx];
		unsigned gen;

		if (pthread_rwlock_rdlock(&s->lock))
			die("pthread_rwlock_rdlock");
		gen = atomic_load_explicit(&s->gen, memory_order_acquire);
		check_slot(idx, gen);
		if (pthread_rwlock_unlock(&s->lock))
			die("pthread_rwlock_unlock rd");
		atomic_fetch_add_explicit(&ops, 1, memory_order_relaxed);
	}
	return NULL;
}

static void mutate_slot(int idx)
{
	struct slot *s = &slots[idx];
	unsigned gen;
	void *p;

	if (pthread_rwlock_wrlock(&s->lock))
		die("pthread_rwlock_wrlock");

	if (munmap(s->data, s->len))
		die("munmap data");
	p = mmap(s->data, s->len, PROT_READ | PROT_WRITE,
		 MAP_PRIVATE | MAP_ANONYMOUS | MAP_FIXED, -1, 0);
	if (p == MAP_FAILED)
		die("mmap fixed data");
	if (p != s->data) {
		fprintf(stderr, "mmap fixed returned %p, wanted %p\n", p, s->data);
		exit(1);
	}

	gen = atomic_fetch_add_explicit(&s->gen, 1, memory_order_acq_rel) + 1;
	fill_slot(idx, gen);
	if (pthread_rwlock_unlock(&s->lock))
		die("pthread_rwlock_unlock wr");
	atomic_fetch_add_explicit(&ops, 1, memory_order_relaxed);
}

static void *mutator(void *argp)
{
	struct arg *a = argp;
	uint64_t seed = 0xd1b54a32d192ed03ULL ^ (uint64_t)a->id;

	pin_cpu(a->cpu);
	while (atomic_load_explicit(&running, memory_order_acquire))
		mutate_slot(rnd64(&seed) % opts.slots);
	return NULL;
}

static int run_instance(void)
{
	pthread_t *wt, *mt;
	struct arg *wa, *ma;
	uint64_t end, last;

	page_size = sysconf(_SC_PAGESIZE);
	if (page_size <= 0)
		die("page size");
	install_segv();
	init_cookie();
	load_cpus();
	init_slots();

	if (!opts.quiet) {
		printf("mmap/munmap TLB stress\n");
		printf("pid=%d workers=%d mutators=%d cpus=%d slots=%d footprint=%zuMiB duration=%ds cookie=0x%016" PRIx64 "\n",
		       getpid(), opts.workers, opts.mutators, nr_cpus, opts.slots,
		       (opts.slots * SLOT_PAGES * (size_t)page_size) >> 20,
		       opts.duration_s, cookie);
		cpuid_invlpgb();
		fflush(stdout);
	}

	wt = calloc(opts.workers, sizeof(*wt));
	mt = calloc(opts.mutators, sizeof(*mt));
	wa = calloc(opts.workers, sizeof(*wa));
	ma = calloc(opts.mutators, sizeof(*ma));
	if (!wt || !mt || !wa || !ma)
		die("calloc threads");

	for (int i = 0; i < opts.workers; i++) {
		wa[i].id = i;
		wa[i].cpu = cpus[i % nr_cpus];
		if (pthread_create(&wt[i], NULL, worker, &wa[i]))
			die("pthread_create worker");
	}
	for (int i = 0; i < opts.mutators; i++) {
		ma[i].id = i;
		ma[i].cpu = cpus[(opts.workers + i) % nr_cpus];
		if (pthread_create(&mt[i], NULL, mutator, &ma[i]))
			die("pthread_create mutator");
	}

	end = nsec_now() + (uint64_t)opts.duration_s * 1000000000ULL;
	last = nsec_now();
	while (nsec_now() < end) {
		sleep(1);
		if (!opts.quiet && nsec_now() - last >= 1000000000ULL) {
			printf("ops=%lu corruptions=%lu\n",
			       atomic_load(&ops), atomic_load(&corruptions));
			fflush(stdout);
			last = nsec_now();
		}
	}
	atomic_store_explicit(&running, false, memory_order_release);
	for (int i = 0; i < opts.workers; i++)
		pthread_join(wt[i], NULL);
	for (int i = 0; i < opts.mutators; i++)
		pthread_join(mt[i], NULL);
	if (!opts.quiet)
		printf("done: ops=%lu corruptions=%lu no SIGSEGV\n",
		       atomic_load(&ops), atomic_load(&corruptions));
	return 0;
}

static int run_batch(void)
{
	printf("batch: rounds=%d jobs=%d duration=%d workers=%d mutators=%d slots=%d\n",
	       opts.rounds, opts.jobs, opts.duration_s, opts.workers,
	       opts.mutators, opts.slots);
	fflush(stdout);

	for (int round = 1; round <= opts.rounds; round++) {
		pid_t *pids = calloc(opts.jobs, sizeof(*pids));
		bool failed = false;

		if (!pids)
			die("calloc pids");
		printf("round=%d\n", round);
		fflush(stdout);

		for (int job = 0; job < opts.jobs; job++) {
			pids[job] = fork();
			if (pids[job] < 0)
				die("fork");
			if (!pids[job])
				_exit(run_instance());
		}

		for (int job = 0; job < opts.jobs; job++) {
			int status;

			if (waitpid(pids[job], &status, 0) < 0)
				die("waitpid");
			if (!WIFEXITED(status) || WEXITSTATUS(status) != 0) {
				failed = true;
				fprintf(stderr, "job %d pid %d failed status=0x%x\n",
					job + 1, pids[job], status);
			}
		}
		free(pids);
		if (failed)
			return 139;
	}

	printf("completed all rounds without reproducing\n");
	return 0;
}

static void usage(const char *p)
{
	fprintf(stderr,
		"Usage: %s [--batch] [--rounds N] [--jobs N] [-d seconds] "
		"[-w workers] [-m mutators] [-s slots] [-q]\n",
		p);
	exit(2);
}

static int env_int(const char *name, int fallback)
{
	const char *v = getenv(name);

	return v && *v ? atoi(v) : fallback;
}

static void parse_env(void)
{
	opts.rounds = env_int("ROUNDS", opts.rounds);
	opts.jobs = env_int("JOBS", opts.jobs);
	opts.duration_s = env_int("DURATION", opts.duration_s);
	opts.workers = env_int("WORKERS", opts.workers);
	opts.mutators = env_int("MUTATORS", opts.mutators);
	opts.slots = env_int("SLOTS", opts.slots);
}

static void parse(int argc, char **argv)
{
	parse_env();
	for (int i = 1; i < argc; i++) {
		if (!strcmp(argv[i], "--batch"))
			opts.batch = true;
		else if (!strcmp(argv[i], "--rounds") && i + 1 < argc)
			opts.rounds = atoi(argv[++i]);
		else if (!strcmp(argv[i], "--jobs") && i + 1 < argc)
			opts.jobs = atoi(argv[++i]);
		else if ((!strcmp(argv[i], "-d") || !strcmp(argv[i], "--duration")) && i + 1 < argc)
			opts.duration_s = atoi(argv[++i]);
		else if ((!strcmp(argv[i], "-w") || !strcmp(argv[i], "--workers")) && i + 1 < argc)
			opts.workers = atoi(argv[++i]);
		else if ((!strcmp(argv[i], "-m") || !strcmp(argv[i], "--mutators")) && i + 1 < argc)
			opts.mutators = atoi(argv[++i]);
		else if ((!strcmp(argv[i], "-s") || !strcmp(argv[i], "--slots")) && i + 1 < argc)
			opts.slots = atoi(argv[++i]);
		else if (!strcmp(argv[i], "-q") || !strcmp(argv[i], "--quiet"))
			opts.quiet = true;
		else
			usage(argv[0]);
	}
	if (opts.duration_s <= 0 || opts.slots <= 0 || opts.workers < 0 ||
	    opts.mutators < 0 || opts.rounds <= 0 || opts.jobs <= 0)
		usage(argv[0]);
}

int main(int argc, char **argv)
{
	parse(argc, argv);
	if (opts.batch)
		return run_batch();
	return run_instance();
}