/*
 * DF-5442 - Cross-CPU MTRR store/refetch race in amd64_mem.c
 *
 * amd64_mrstore() queues amd64_mrstoreone() to remote CPUs with the
 * asynchronous lwkt_send_ipiq_mask() (no completion barrier), then
 * amd64_mrset() immediately calls amd64_mrfetch() which REWRITES the
 * shared sc->mr_desc[] array that the remote CPUs are still reading.
 * A remote CPU that executes its storeone while the local CPU is inside
 * mrfetch() reads a torn descriptor generation and programs a torn MTRR
 * pair.  Because that storeone is the last one queued, nothing repairs
 * it: the desync persists on that CPU.
 *
 * This program:
 *   phase 1 (control): SET/GET/REMOVE roundtrip on a variable MTRR,
 *                      proving the ioctl path works and showing state.
 *   phase 2 (race):    hammer alternating MEMRANGE_SET UPDATE calls
 *                      (A/B), then one final distinctive SET F, then
 *                      read the per-CPU MTRR snapshot sysctl exposed by
 *                      the mtrrdump.ko module.  The module's cpusync
 *                      barrier drains all pending storeone IPIs before
 *                      sampling, so ANY cross-CPU divergence in the
 *                      post-drain snapshot is a persistent torn state.
 *
 * Must run as root (/dev/mem is root:kmem 0640 + caps gate at open).
 */
#include <sys/types.h>
#include <sys/ioctl.h>
#include <sys/memrange.h>
#include <sys/sysctl.h>
#include <sys/syscall.h>
#include <sys/usched.h>
#include <fcntl.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <errno.h>
#include <pthread.h>

#define NDESC_MAX 512

static int fd;
static struct mem_range_desc descs[NDESC_MAX];
static int ndesc;
static volatile int hammer_stop;
static volatile int busy_stop;
static volatile long hammer_count;

/* the three competing configurations: distinct base AND len so that any
 * torn (base,len) pairing is not equal to any legitimate generation */
#define A_BASE 0x400000000ULL
#define A_LEN  0x1000ULL
#define B_BASE 0x500000000ULL
#define B_LEN  0x4000ULL
#define F_BASE 0x600000000ULL
#define F_LEN  0x2000ULL
/* the final generation changes every cycle so the torn read is visible */
static u_int64_t f_base_now = F_BASE;

static int
mr_get(void)
{
	struct mem_range_op mo;

	memset(&mo, 0, sizeof(mo));
	mo.mo_arg[0] = NDESC_MAX;
	mo.mo_desc = descs;
	if (ioctl(fd, MEMRANGE_GET, &mo) < 0)
		return (-1);
	ndesc = mo.mo_arg[0];
	return (0);
}

static int
mr_set(u_int64_t base, u_int64_t len, int flags, const char *owner, int arg)
{
	struct mem_range_op mo;
	struct mem_range_desc md;

	memset(&md, 0, sizeof(md));
	md.mr_base = base;
	md.mr_len = len;
	md.mr_flags = flags;
	strncpy(md.mr_owner, owner, sizeof(md.mr_owner) - 1);
	memset(&mo, 0, sizeof(mo));
	mo.mo_desc = &md;
	mo.mo_arg[0] = arg;
	return (ioctl(fd, MEMRANGE_SET, &mo));
}

static void
dump_variable(const char *tag)
{
	int i;

	printf("---- %s: ndesc=%d\n", tag, ndesc);
	for (i = 0; i < ndesc; i++) {
		if (descs[i].mr_flags & MDF_ACTIVE)
			printf("[%3d] base=0x%016llx len=0x%016llx flags=0x%08x owner=\"%s\"\n",
			    i, (unsigned long long)descs[i].mr_base,
			    (unsigned long long)descs[i].mr_len,
			    descs[i].mr_flags, descs[i].mr_owner);
	}
}

static int
find_active(u_int64_t base, u_int64_t len)
{
	int i;

	for (i = 0; i < ndesc; i++) {
		if (descs[i].mr_base == base && descs[i].mr_len == len &&
		    (descs[i].mr_flags & MDF_ACTIVE))
			return (i);
	}
	return (-1);
}

static void
snapshot(const char *tag)
{
	char buf[16384];
	size_t len = sizeof(buf) - 1;

	memset(buf, 0, sizeof(buf));
	if (sysctlbyname("debug.mtrr_snap", buf, &len, NULL, 0) < 0) {
		printf("snapshot(%s): sysctl failed: %s\n", tag, strerror(errno));
		return;
	}
	printf("==== SNAPSHOT %s ====\n%s", tag, buf);
}

static int i_dummy;

static void
cleanup(void)
{
	int i;

	printf("cleanup: removing all our ranges ...\n");
	if (mr_get() == 0) {
		for (i = 0; i < ndesc; i++) {
			if ((descs[i].mr_flags & MDF_ACTIVE) &&
			    (strcmp(descs[i].mr_owner, "pocA") == 0 ||
			     strcmp(descs[i].mr_owner, "pocB") == 0 ||
			     strcmp(descs[i].mr_owner, "finl") == 0 ||
			     strcmp(descs[i].mr_owner, "mrtrace") == 0)) {
				printf("  remove [%d] base=0x%016llx owner=%s: ",
				    i, (unsigned long long)descs[i].mr_base,
				    descs[i].mr_owner);
				if (mr_set(descs[i].mr_base, descs[i].mr_len, 0,
				    "", MEMRANGE_SET_REMOVE) < 0)
					printf("%s\n", strerror(errno));
				else
					printf("ok\n");
			}
		}
	}
}

static void *
hammer_thread(void *arg)
{
	/*
	 * Concurrent MEMRANGE_SET stream.  mem_lock (kern_memio.c:493)
	 * serializes the ioctl critical sections against each other, but
	 * the remote-CPU amd64_mrstoreone() IPIs queue up behind whatever
	 * each remote CPU is still executing, so per-CPU IPI backlog (and
	 * thus remote storeone arrival time relative to the setter's
	 * amd64_mrfetch() rewrite window) spreads across the torn band.
	 */
	while (!hammer_stop) {
		if (mr_set(A_BASE, A_LEN, MDF_UNCACHEABLE, "pocA",
		    MEMRANGE_SET_UPDATE) < 0) {
			if (errno != EBUSY && errno != ENOSPC)
				break;
		}
		if (mr_set(B_BASE, B_LEN, MDF_UNCACHEABLE, "pocB",
		    MEMRANGE_SET_UPDATE) < 0) {
			if (errno != EBUSY && errno != ENOSPC)
				break;
		}
		hammer_count += 2;
	}
	return (NULL);
}

static void *
busy_thread(void *arg)
{
	/*
	 * DF-5442 load injector: continuously fire 60-120us busy-IPIs at
	 * all CPUs via the mtrrdump.ko instrumentation hook.  This is the
	 * realistic kernel-IPI contention (TLB shootdown / cpusync style)
	 * that delays a remote cpu's amd64_mrstoreone() execution into the
	 * setter's amd64_mrfetch() rewrite window.
	 */
	int usec;
	size_t sz;

	while (!busy_stop) {
		usec = 10 + random() % 71;
		sz = sizeof(usec);
		if (sysctlbyname("debug.mtrr_busy", NULL, NULL, &usec, sz) < 0)
			break;
	}
	return (NULL);
}

int
main(int argc, char **argv)
{
	int cycles = (argc > 1) ? atoi(argv[1]) : 50;
	int inner = (argc > 2) ? atoi(argv[2]) : 200;
	int nthreads = (argc > 3) ? atoi(argv[3]) : 3;
	int c, k, idx, errors = 0, divergent_hits = 0;
	pthread_t th[9];
	char snapbuf[16384];
	size_t len;

	fd = open("/dev/mem", O_RDWR);
	if (fd < 0) {
		perror("open /dev/mem (must be root)");
		exit(1);
	}
	if (mr_get() < 0) {
		perror("MEMRANGE_GET");
		exit(1);
	}
	printf("---- initial: ndesc=%d (variable descs: %d)\n", ndesc,
	    ndesc - 88 > 0 ? ndesc - 88 : 0);
	for (i_dummy = 88; i_dummy < ndesc; i_dummy++) {
		if (descs[i_dummy].mr_flags & MDF_ACTIVE)
			printf("[%3d] base=0x%016llx len=0x%016llx flags=0x%08x owner=\"%s\"\n",
			    i_dummy, (unsigned long long)descs[i_dummy].mr_base,
			    (unsigned long long)descs[i_dummy].mr_len,
			    descs[i_dummy].mr_flags, descs[i_dummy].mr_owner);
	}
	snapshot("baseline");

	/* ---- phase 1: control roundtrip ---- */
	printf("---- control: SET A(0x%llx,0x%llx,UC) ...\n",
	    (unsigned long long)A_BASE, (unsigned long long)A_LEN);
	if (mr_set(A_BASE, A_LEN, MDF_UNCACHEABLE, "mrtrace",
	    MEMRANGE_SET_UPDATE) < 0) {
		perror("control SET A");
		exit(1);
	}
	if (mr_get() < 0)
		perror("GET after SET A");
	idx = find_active(A_BASE, A_LEN);
	printf("control: A active at desc %d (expect >= 0)\n", idx);
	if (idx < 0) {
		printf("CONTROL FAILED: A not active\n");
		exit(1);
	}
	if (mr_set(A_BASE, A_LEN, 0, "", MEMRANGE_SET_REMOVE) < 0)
		perror("control REMOVE A");
	if (mr_get() == 0)
		printf("control: after REMOVE, A active desc = %d (expect -1)\n",
		    find_active(A_BASE, A_LEN));
	printf("---- control OK\n");
	snapshot("post-control");

	/* ---- phase 2: race hammer ---- */
	if (nthreads > 8)
		nthreads = 8;
	for (k = 0; k < nthreads; k++)
		pthread_create(&th[k], NULL, hammer_thread, NULL);
	for (c = 0; c < cycles; c++) {
		for (k = 0; k < inner; k++) {
			if (mr_set(A_BASE, A_LEN, MDF_UNCACHEABLE, "pocA",
			    MEMRANGE_SET_UPDATE) < 0 &&
			    errno != EBUSY && errno != ENOSPC) {
				errors++;
				perror("SET A");
			}
			if (mr_set(B_BASE, B_LEN, MDF_UNCACHEABLE, "pocB",
			    MEMRANGE_SET_UPDATE) < 0 &&
			    errno != EBUSY && errno != ENOSPC) {
				errors++;
				perror("SET B");
			}
		}
		/* quiesce the SET hammer only; keep the busy-IPI injectors
		 * running THROUGH the final SET so the remote storeone
		 * arrival time spreads across the torn band */
		hammer_stop = 1;
		for (k = 0; k < nthreads; k++)
			pthread_join(th[k], NULL);
		hammer_stop = 0;

		/*
		 * Arrival-time SCAN: fire a single busy-IPI of swept
		 * duration U just before the final SET.  The remote
		 * storeone IPIs queue behind it, so their arrival shifts
		 * by ~U; scanning U across 0..150usec walks the remote's
		 * variable-slot reads through the caller's
		 * amd64_mrfetch() rewrite window deterministically.
		 * The final setter is pinned to cpu0 so lwkt_send_ipiq3_
		 * mask() runs the self-storeone inline first and queues
		 * every remote behind the busy item.
		 */
		{
			int U = (c % 20) * 50;
			size_t sz = sizeof(U);
			int cpu0 = 0;

			if (sysctlbyname("debug.mtrr_busy", NULL, NULL,
			    &U, sz) < 0)
				printf("cycle %d: busy(%d) failed: %s\n", c, U,
				    strerror(errno));
			if (syscall(SYS_usched_set, 0, USCHED_SET_CPU,
			    &cpu0, sizeof(cpu0)) < 0)
				printf("usched_set(cpu0): %s (continuing "
				    "unpinned)\n", strerror(errno));
		}
		usleep(2000);

		/* single final distinctive generation - the last mutation */
		f_base_now = F_BASE + (u_int64_t)(c % 32) * 0x10000ULL;
		if (mr_set(f_base_now, F_LEN, MDF_UNCACHEABLE, "finl",
		    MEMRANGE_SET_UPDATE) < 0) {
			errors++;
			perror("SET F");
			cleanup();
			exit(1);
		}
		/* let every remote CPU drain its storeone backlog; a torn
		 * read of the final generation is never repaired */
		usleep(200000);
		len = sizeof(snapbuf) - 1;
		memset(snapbuf, 0, sizeof(snapbuf));
		if (sysctlbyname("debug.mtrr_snap", snapbuf, &len, NULL, 0) == 0) {
			char *hit = strstr(snapbuf, "divergent_slots=");

			printf("cycle %d: %s", c, snapbuf);
			if (hit && atoi(hit + 17) > 0) {
				divergent_hits++;
				printf("*** CYCLE %d: POST-DRAIN CROSS-CPU "
				    "MTRR DIVERGENCE (torn state) ***\n", c);
			}
		} else {
			printf("cycle %d: snapshot failed: %s\n", c,
			    strerror(errno));
		}
		/* free the slot so next cycle's F lands in the same slot
		 * with a DIFFERENT value (torn reads stay visible) */
		if (mr_set(f_base_now, F_LEN, 0, "", MEMRANGE_SET_REMOVE) < 0)
			printf("cycle %d: remove F: %s\n", c, strerror(errno));
		/* restart hammer threads for the next cycle */
		for (k = 0; k < nthreads; k++)
			pthread_create(&th[k], NULL, hammer_thread, NULL);
	}
	printf("hammer done: cycles=%d inner=%d threads=%d "
	    "extra_sets=%ld ioctl_errors=%d divergent_hits=%d\n",
	    cycles, inner, nthreads, hammer_count, errors, divergent_hits);

	cleanup();
	if (mr_get() == 0)
		dump_variable("final-after-cleanup");
	snapshot("post-cleanup");
	return (divergent_hits ? 2 : 0);
}