#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <string.h>
#include <errno.h>
#include <signal.h>

#include <sys/ipc.h>
#include <sys/shm.h>

#include <sys/types.h>
#include <sys/wait.h>

#define HAVE_SMEM

#define	MB	(1024 * 1024)

#define SHMEM_KEY	908234

#define USAGE "Usage: %s [-12fhsw] [mb]\n\
	-1	allocate memory in one process\n\
	-2	allocate memory in two processes (parent, child)\n\
	-f	fork a child process\n\
	-h	help (or --help)\n\
	-s	use shared memory, rather than malloc()\n\
	-w	wait indefinitely after allocations\n\
	mb	allocate megabyte of memory (default=1)\n"

#define PS_STR	"ps -o pid,trs,drs,vsz,rss,cmd -C %s | sort -n"

int shmid = -1;
void rm_shmem(int);

int
main(int argc, char **argv)
{
	int mb = 1, do_alloc_in_proc = 0, do_shmem = 0,
		do_fork = 0, do_wait = 0, opt;
	char cmd[1024];
	pid_t child_pid = 0;
	
	while ((opt	= getopt(argc, argv, "12fhsw")) != -1)
	{
		switch (opt)
		{
			case '1':
				do_alloc_in_proc = 1;
				break;
			case '2':
				do_alloc_in_proc = 2;
				break;
			case 'f':
				do_fork = 1;
				break;
			case 's':
				do_shmem = 1;
				break;
			case 'w':
				do_wait = 1;
				break;
			case 'h':	/* fall through */
			default:
				fprintf(stderr, USAGE, argv[0]);
				exit(1);
		}

	}

	if (optind < argc - 1 ||
		(optind == argc - 1 && strcmp(argv[optind], "--help") == 0))
	{
		fprintf(stderr, USAGE, argv[0]);
		exit(1);
	}
	else if (optind == argc - 1)
		mb = atoi(argv[optind]);

	if (!do_fork && do_alloc_in_proc == 2)
	{
		fprintf(stderr, "Only fork mode allows allocation in two processes\n");
		exit(1);
	}
	
	if (mb != 0)
	{
		char *mem;
		
		if (!do_shmem)
			mem = malloc(mb * MB);
		else
		{
			atexit((void (*)(void)) rm_shmem);
			signal(SIGINT, rm_shmem);
			signal(SIGQUIT, rm_shmem);
			signal(SIGTERM, rm_shmem);

			if ((shmid = shmget(SHMEM_KEY, mb * MB, IPC_CREAT | IPC_EXCL | 0600)) == -1)
			{
				fprintf(stderr, "Failure to allocate shared memory: %s\n", strerror(errno));
				exit(1);
			}

			if ((mem = shmat(shmid, NULL, 0)) == (void *) -1)
			{
				fprintf(stderr, "Failure to attach to shared memory: %s\n", strerror(errno));
				exit(1);
			}
		}

		if (do_fork)
			child_pid = fork();	/* parent and child fall through */

		if ((!do_fork && do_alloc_in_proc) ||
			(do_fork && (do_alloc_in_proc == 2 ||
				(child_pid != 0 && do_alloc_in_proc == 1))))
		{
			int i;
	
			for (i = 0; i < mb * MB; i++)
				mem[i] = 1;
		}
	}

	/* only run in parent */
	if (!do_fork || child_pid != 0)
	{
		if (do_fork)
			sleep(1);	/* give the child time to access the memory */

		snprintf(cmd, sizeof(cmd), PS_STR,
			strrchr(argv[0], '/') ? strrchr(argv[0], '/') + 1 : argv[0]);
		system(cmd);
#ifdef HAVE_SMEM
		snprintf(cmd, sizeof(cmd), "smem | head -1");
		system(cmd);
		/* use [ ] to avoid matching our own process */
		snprintf(cmd, sizeof(cmd), "smem | grep '[ ]%s\\>' | sort -n", argv[0]);
		system(cmd);
		puts("");
#endif
		
	}

	if (do_fork)
	{
		int dummy;
		
		if (child_pid != 0)
			wait(&dummy);	/* wait for child to exit so later 'ps' is not confused */
		else
			sleep(3);	/* give parent time to run 'ps' */
	}

	if (do_wait)
		sleep(999999);
		
	return 0;
}

void rm_shmem(int sig)
{
	if (shmid != -1)
		shmctl(shmid, IPC_RMID, NULL);
}