/*
 * steamroller.h: common wrappers and calibration functions.
 * 
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA 02111-1307, USA.
 *
 * Copyright (C) IBM Corporation, 2005
 *
 * Author: Paul E. McKenney <paulmck@us.ibm.com>
 */

#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <sys/types.h>
#include <sys/mman.h>
#include <sys/stat.h>
#include <sys/wait.h>
#include <errno.h>
#include <fcntl.h>
#include <limits.h>
#include <sched.h>
#include <signal.h>

#define STEAMROLLER_EARLY 0
#define STEAMROLLER_RACED 1
#define STEAMROLLER_LATE  2

/*
 * Map the specified amount of memory.  If specified, map from the
 * file, otherwise, if "fd" is -1, anonymously.
 */

char zerobuf[256] = { 0 };

void *mapmem(size_t length, int fd)
{
	int i;
	void *p;
	size_t map_length;
	size_t pgsz = getpagesize();

	map_length = ((length + pgsz - 1) / pgsz) * pgsz;
	if (fd != -1) {
		lseek(fd, (off_t)0, SEEK_SET);
		i = 0;
		while (i < map_length) {
			write(fd, zerobuf, sizeof(zerobuf));
			i += sizeof(zerobuf);
		}
	}
	p = mmap(NULL, map_length,
		 PROT_READ|PROT_WRITE,
		 (fd != -1 ? 0 : MAP_ANONYMOUS) | MAP_SHARED,
		 (fd != -1 ? fd : 0), (off_t)0);
	if (p == (void *)-1) {
		perror("mmap");
		exit(-1);
	}
}

/*
 * Return time as a floating-point number rather than struct timeval.
 */

double d_gettimeofday(void)
{
	int retval;
	struct timeval tv;

	retval = gettimeofday(&tv, (struct timezone *)NULL);
	if (retval != 0) {
		perror("gettimeofday");
		exit(-1);
	}
	return (tv.tv_sec + ((double)tv.tv_usec) / 1000000.);
}

/*
 * Data for spinloop.
 */

#define SPINDELAY_CALIBRATE_N 1000000
double spindelay_multiplier;

/*
 * Calibrate spinloop for delay times.
 */

void calibrate_spindelay(void)
{
	int i;
	double et;
	double st;

	st = d_gettimeofday();
	i = SPINDELAY_CALIBRATE_N;
	while (--i > 0) {
		continue;
	}
	et = d_gettimeofday();
	spindelay_multiplier = ((double)SPINDELAY_CALIBRATE_N) /
			       ((et - st) * 1000000.);
}

/*
 * Spin for the specified number of loops.
 */

void spindelay(int nloops)
{
	int i = nloops;
	
	while (--i > 0) {
		continue;
	}
}

/*
 * Spin for the specified number of microseconds.
 */

void spindelayus(int us)
{
	int i = (int)(us * spindelay_multiplier);
	
	while (--i > 0) {
		continue;
	}
}

/*
 * Convert from microseconds to spindelay() units.
 * Useful when attempting to carpet-bomb a suspected race window.
 */

int us2spindelay(int us)
{
	return (int)(us * spindelay_multiplier);
}

/*
 * Get a rough measurement of the weight of process creation and deletion.
 * Multiply by a safety factor in order to feel better about it.
 * Returns the measure in microseconds.
 */

int process_creation_deletion_us(void)
{
	double et;
	double st;
	int status;

	st = d_gettimeofday();
	if (fork() == 0) {
		exit(0);
	}
	wait(&status);
	et = d_gettimeofday();
	return (int)((et - st) * 1000000);
}

long steamroller_remaining_cpus;

/*
 * Get affinity mask for the system.  Affinity the caller onto the
 * first CPU (skipping as many as specified first), and return the
 * mask for the second CPU.  A child process would normally
 * affinity itself to this returned mask.
 */

long getaffinity(int skipcpus)
{
	long childcpuset;
	int i;
	long oldcpuset = 0;
	long parentcpuset;
	int skipped = 0;

	i = sched_getaffinity(0, sizeof(oldcpuset), &oldcpuset);
	for (i = 0; i < sizeof(oldcpuset) * 8; i++) {
		if (oldcpuset & (1 << i)) {
			if (++skipped >= skipcpus) {
				break;
			}
		}
	}
	for (i = 0; i < sizeof(oldcpuset) * 8; i++) {
		if (oldcpuset & (1 << i)) {
			parentcpuset = 1 << i;
			break;
		}
	}
	for (i++; i < sizeof(oldcpuset) * 8; i++) {
		if (oldcpuset & (1 << i)) {
			childcpuset = 1 << i;
			break;
		}
	}
	if (i >= sizeof(oldcpuset) * 8) {
		fprintf(stderr, "need at least two CPUs!\n");
		exit(-1);
	}
	sched_setaffinity(0, sizeof(parentcpuset), &parentcpuset);
	steamroller_remaining_cpus = oldcpuset & ~((childcpuset << 1) - 1);
	return childcpuset;
}

/*
 * Called by a child process that needs to run on all available
 * CPUs.  Example uses are race cases involving fork() storms.
 */

int child_affinity_rest(long childcpuset)
{
	long newcpuset = childcpuset | steamroller_remaining_cpus;

	return (sched_setaffinity(0, sizeof(newcpuset), &newcpuset));
}

struct steamroller_args {
	long verbosity;
} steamroller_args;

/*
 * Print usage message.  Keep synched with steamroller_init().
 */
void steamroller_usage(char *name)
{
	fprintf(stderr, "Usage: %s [args]\n", name);
	fprintf(stderr, "\t--verbose [level]\n");
	exit(-1);
}

/*
 * Initialize the steamroller.
 */

long steamroller_init(int argc, char *argv[])
{
	int i = 1;

	while (i < argc) {
		steamroller_args.verbosity = 0;
		if (strcmp(argv[i], "--verbose") == 0) {
			steamroller_args.verbosity = 1;
			i++;
			if ((i < argc) && (argv[i][0] != '-')) {
				steamroller_args.verbosity =
					strtol(argv[i], NULL, 0);
				i++;
			}
			continue;
		}
		steamroller_usage(argv[0]);
	}

	calibrate_spindelay();
	return getaffinity(0);
}

/*
 * Invoke the specified function "nevals" times, accumulating the
 * early/raced/late counts in "resulttab".
 */

void raceeval(int (*f)(void *, int, long), void *p, long childcpuset,
	      int childdelay, int nevals, int resulttab[3])
{
	int j;
	int result;

	resulttab[0] = resulttab[1] = resulttab[2] = 0;
	for (j = 0; j < nevals; j++) {
		result = f(p, childdelay, childcpuset);
		if ((result != STEAMROLLER_EARLY) &&
		    (result != STEAMROLLER_RACED) &&
		    (result != STEAMROLLER_LATE))
		{
			fprintf(stderr,
				"invalid test return code: %d\n", result);
			exit(-1);
		}
		resulttab[result]++;
	}
	if (steamroller_args.verbosity >= 100) {
		printf("childdelay = %d  %d:%d:%d\n",
		       childdelay, resulttab[0], resulttab[1], resulttab[2]);
	}
}

/*
 * Scan the specified range in a multiplicative search, printing
 * early/raced/late distribution for each point selected.
 */

void racescan(int (*f)(void *, int, long), void *p, long childcpuset,
	      int start, int mult, int div, int lim)
{
	int i;
	int j;
	int result;
	int resulttab[3];
		
	for (i = start; i < lim; i = i * mult / div + 1) {
		raceeval(f, p, childcpuset, i, 9, resulttab);
		printf("%d:%d:%d -- child # %d\n",
		       resulttab[0], resulttab[1], resulttab[2], i);
	}
}

/*
 * Run the race test function "neval" times, returning TRUE if
 * the "desiredresult" occurs at least "nvotes" times.  This is
 * useful in binary searches, since it rejects occasional errors
 * due to daemons, interrupts, etc.  The "desiredresult" value
 * may be STEAMROLLER_EARLY, STEAMROLLER_RACED, or STEAMROLLER_LATE.
 */

int racevote(int (*f)(void *, int, long), void *p, long childcpuset,
	     int childdelay, int neval, int desiredresult, int nvotes)
{
	int i;
	int ndesired = 0;
	int result;
	int resulttab[3];

	raceeval(f, p, childcpuset, childdelay, neval, resulttab);
	return resulttab[desiredresult] >= nvotes;
}

/*
 * Use power searches ("mult" == 2 and "div" == 1 for binary search) to
 * locate both sides of a race window for the specified race function.
 * The "eps" parameter specified a desired sloppiness -- values in the
 * 100-1000 range seem to work well.
 */

int racepowersearch(int (*f)(void *, int, long), void *p, long childcpuset,
		    int start, int mult, int div, int lim, int eps,
		    int *before, int *after)
{
	int cur;
	int early;
	int foundafter = 0;
	int foundbefore = 0;
	int i;
	int late;
	int race;
	int result;
	int resulttab[3];
		
	/* Do power-search to straddle the race window. */

	for (i = start; i < lim; i = i * mult / div + 1) {
		raceeval(f, p, childcpuset, i, 9, resulttab);
		if (resulttab[STEAMROLLER_EARLY] >= 7) {
			early = i;
			foundbefore = 1;
		} else if (resulttab[STEAMROLLER_LATE] >= 7) {
			if (foundbefore) {
				late = i;
				foundafter = 1;
				break;
			} else {
				return 0;
			}
		}
	}
	if (!foundafter) {
		return 0;
	}

	/* Do binary search to locate beginning of race window. */

	race = late;
	do {
		cur = (early + race) / 2;
		if (racevote(f, p, childcpuset,
			     cur, 10, STEAMROLLER_EARLY, 9)) {
			early = cur;
		} else {
			race = cur;
		}
	} while (race - early > eps);
	*before = early;

	/* Do binary search to locate end of race window. */

	race = early;
	do {
		cur = (race + late) / 2;
		if (racevote(f, p, childcpuset,
			     cur, 10, STEAMROLLER_LATE, 9)) {
			late = cur;
		} else {
			race = cur;
		}
	} while (late - race > eps);
	*after = late;
	
	return 1;
}

/*
 * Invoke the race function for every possible value within the specified
 * race window, accumulating the early/raced/late results in "resulttab".
 */

void steamroller(int (*f)(void *, int, long), void *p, long childcpuset,
	         int start, int lim, int resulttab[3])
{
	int i;
	int result;
	double lastverbose;
	double curverbose;

	resulttab[0] = resulttab[1] = resulttab[2] = 0;
	if (steamroller_args.verbosity >= 2) {
		lastverbose = d_gettimeofday();
		printf("steamroller: %d spindelay units\n", start);
	}
	for (i = start; i < lim; i++) {
		result = f(p, i, childcpuset);
		if ((result != STEAMROLLER_EARLY) &&
		    (result != STEAMROLLER_RACED) &&
		    (result != STEAMROLLER_LATE))
		{
			fprintf(stderr,
				"invalid test return code: %d\n", result);
			exit(-1);
		}
		resulttab[result]++;
		if (steamroller_args.verbosity >= 2) {
			curverbose = d_gettimeofday();
			if (curverbose - lastverbose > 5.0) {
				printf("steamroller: %d spindelay units\n", i);
				lastverbose = curverbose;
			}
		}
	}
}

/*
 * Given race function "f" that returns STEAMROLLER_EARLY,
 * STEAMROLLER_RACED, or STEAMROLLER_LATE, search for the race window
 * and "steamroller" over it.
 */

void search_and_steamroller(int (*f)(void *, int, long),
			    void *p, long childcpuset)
{
	int before;
	int after;
	int resulttab[3];
	
	if (!racepowersearch(f, p, childcpuset, 0, 3, 2, INT_MAX, 10,
			     &before, &after)) {
		fprintf(stderr, "Failed to bracket race.\n");
	}
	if (steamroller_args.verbosity) {
		printf("Race range: %d:%d spindelay units\n", before, after);
	}
	steamroller(f, p, childcpuset, before, after, resulttab);
	printf("steamroller distribution: %d:%d:%d\n",
	       resulttab[0], resulttab[1], resulttab[2]);
}
