/*
 *  irman by Davide Libenzi ( irman load generator )
 *  Copyright (C) 2003  Davide Libenzi
 *
 *  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
 *
 *  Davide Libenzi <davidel@xmailserver.org>
 *
 */

#include <sys/types.h>
#include <sys/stat.h>
#include <sys/time.h>
#include <sys/socket.h>
#include <sys/signal.h>
#include <sys/resource.h>
#include <sys/socket.h>
#include <fcntl.h>
#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <unistd.h>
#include <errno.h>


#define BUFSIZE (1024 * 32)


static int *pipes, *child;
static int num_pipes, num_active, num_child, use_socket;
static unsigned long burn_ms;
static char buf1[BUFSIZE], buf2[BUFSIZE];
static volatile sig_atomic_t run = 1;
pid_t parent;

static void signal_all(int signum) {
	if (getpid() == parent) {
		while (num_child >= 0) {
			kill(child[num_child], SIGKILL);
			num_child--;
		}
		exit(0);
	} else if (signum == SIGKILL || getppid() == 1)
		run = 0;
}

unsigned long long getustime(void) {
	struct timeval tm;

	gettimeofday(&tm, NULL);
	return (unsigned long long) tm.tv_sec * 1000ULL + (unsigned long long) tm.tv_usec / 1000ULL;
}


int burn_ms_cpu(unsigned long ms) {
	int i, cmp = 0;
	unsigned long long ts;

	ts = getustime();
	do {
		for (i = 0; i < 4; i++)
			cmp += memcmp(buf1, buf2, BUFSIZE);
	} while (ts + ms > getustime());
	return cmp;
}


pid_t hog_process(void) {
	pid_t pid;

	if (!(pid = fork())) {
		while (run) {
			printf("HOG running %u\n", time(NULL));
			burn_ms_cpu(burn_ms);
		}
		exit(0);
	}
	return pid;
}


pid_t irman_process(int n) {
	int nn;
	pid_t pid;
	u_char ch;

	if (!(pid = fork())) {
		if ((nn = n + num_active) >= num_pipes)
			nn -= num_pipes;
		while (run) {
			printf("reading %u\n", n);
			read(pipes[2 * n], &ch, 1);
			burn_ms_cpu(burn_ms);
			printf("writing %u\n", nn);
			write(pipes[2 * nn + 1], "s", 1);
		}
		exit(0);
	}
	return pid;
}

int main (int argc, char **argv) {
	struct rlimit rl;
	int i, c;
	long work;
	int *cp, run_secs = 0;
	extern char *optarg;
	struct sigaction action;

	parent = getpid();
	num_pipes = 40;
	num_active = 1;
	burn_ms = 300;
	use_socket = 0;
	while ((c = getopt(argc, argv, "n:b:a:s:S:")) != -1) {
		switch (c) {
		case 'n':
			num_pipes = atoi(optarg);
			break;
		case 'b':
			burn_ms = atoi(optarg);
			break;
		case 'a':
			num_active = atoi(optarg);
			break;
		case 's':
			run_secs = 1 + atoi(optarg);
			break;
		case 'S':
			use_socket = 1;
			break;
		default:
			fprintf(stderr, "Illegal argument \"%c\"\n", c);
			exit(1);
		}
	}

	rl.rlim_cur = rl.rlim_max = num_pipes * 2 + 50;
	if (setrlimit(RLIMIT_NOFILE, &rl) == -1) {
		perror("setrlimit"); 
		exit(1);
	}

	pipes = calloc(num_pipes * 2, sizeof(int));
	if (pipes == NULL) {
		perror("malloc");
		exit(1);
	}

	child = calloc(num_pipes, sizeof(int));
	if (child == NULL) {
		perror("malloc");
		exit(1);
	}

	for (cp = pipes, i = 0; i < num_pipes; i++, cp += 2) {
		if (!use_socket) {
			if(pipe(cp) == -1) {
				perror("pipe");
				exit(1);
			}
		} else if (socketpair(AF_UNIX, SOCK_STREAM, 0, cp) == -1) {
			perror("socketpair");
			exit(1);
		}
	}

	memset(buf1, 'f', sizeof(buf1));
	memset(buf2, 'f', sizeof(buf2));

	sigemptyset(&action.sa_mask);
	/* establish termination handler */
	action.sa_handler = signal_all;
	action.sa_flags = SA_NODEFER;
	if (sigaction(SIGTERM, &action, NULL) == -1) {
		perror("Could not install signal handler");
		exit(1);
	}

	for (i = 0; i < num_pipes; i++)
		child[i] = irman_process(i);

	child[i] = hog_process();
	num_child = i;

	for (i = 0; i < num_active; i++)
		write(pipes[2 * i + 1], "s", 1);

	while (--run_secs)
		sleep(1);
	signal_all(SIGKILL);
	exit(0);
}