#include <sys/mman.h>
#include <sys/wait.h>
#include <string.h>
#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>

void usage(char *prog)
{
  printf("Usage: %s <bytes>[K|M] <command>\n", prog);
  printf("\tMost kernels limit the amount of memory to be mlock()ed by a\n");
  printf("\tsingle process to 1/2 of the available physical memory.\n");
  printf("\tWorkaround if you need to steal more (e.g., 130MB on a 144MB machine):\n");
  printf("\t\t%s 70M %s 60M <command>\n\n", prog, prog);
  printf("\tIf you don't want to run your program as a child of %s,\n", prog);
  printf("\tuse e.g., 'fullsleep <time>' as the command.\n\n");
  printf("\tHacked up 2000 by Marcel Waldvogel\n");
  exit(1);
}

int main(int argc, char **argv)
{
  void *mem;
  size_t bytes;
  int status;
  pid_t pid;

  if (argc < 3) usage(argv[0]);
  bytes = atoi(argv[1]);
  if (strchr(argv[1], 'k') || strchr(argv[1], 'K')) {
    bytes *= 1024;
  }
  if (strchr(argv[1], 'm') || strchr(argv[1], 'M')) {
    bytes *= 1024*1024;
  }
  if (bytes == 0) usage(argv[0]);

  mem = malloc(bytes);
  if (mem == NULL) {
    perror("malloc");
    exit(1);
  }
  memset(mem, 0, bytes);

  if (mlock(mem, bytes)) { /* May get one page too many */
    perror("mlock");
    exit(1);
  }

  switch (pid = fork()) {
  case -1:	perror("fork");
		exit(1);
  case  0:	execvp(argv[2], &argv[2]);
		/* Should not get here if everything's all right */
		perror("exec");
		exit(1);
  default:
#ifdef MLOCK_WAIT_DEBUG
		do{
#endif
		  pid = wait(&status);
		  if (pid < 0) {
		    perror("wait");
		    exit(1);
		  }
#ifdef MLOCK_WAIT_DEBUG
		  if (WIFEXITED(status)) {
		    printf("%s: %s exited with return code %d\n",
			argv[0], argv[2], WEXITSTATUS(status));
		    exit(0);
		  }
		  if (WIFSIGNALED(status)) {
		    printf("%s: %s exited due to signal %d\n",
			argv[0], argv[2], WTERMSIG(status));
		    exit(0);
		  }
		  if (WIFSTOPPED(status)) {
		    printf("%s: %s was stopped due to signal %d\n",
			argv[0], argv[2], WSTOPSIG(status));
		  }
		} while (1);
#endif
  }

  return 0;
}