#include <sys/param.h> /* XXX include file bugs in stock system */
/*
 * mmap - mmap(2) stress-tester
 *
 * Usage: mmap [-s] [-{wire|lock}[all]] [-unit U] [-v] [-r] [-w] [-z]
 *		[-{spin|loop}] [-fork F] [-fr] [-fw] [-fz] N
 *
 * U and N are byte counts, but can have suffixes k, m, or g
 *  (case-independent) to indicate multiplication by 1024, 1048576, or
 *  1073741824, respectively.  The number may have a fractional part,
 *  but the actual value must be a multiple of the page size.
 *  (Fractions are useful for things like "1.25G" or "2.5M".)
 *
 * U must be non-negative; if omitted, the default is 1M.
 *
 * If N is greater than zero, mmap tries to allocate that much space;
 *  if it succeeds, it then sleeps forever (in a loop with sigpause).
 *
 * If N is less than zero, mmap allocates as much space as it can, then
 *  frees as much as the negative of N; if all this succeeds, it then
 *  sleeeps forever.
 *
 * If N is equal to zero, mmap exits immediately without doing
 *  anything.
 *
 * mmap is normally silent unless an error occurs.  -v makes it
 *  somewhat chatty.
 *
 * mmap normally does not do anything with the memory it allocates.  -r
 *  causes it to read the first byte of each page; -w is similar, but
 *  causes writes.  -z causes it to bzero the whole block.
 *
 * -fork causes mmap to finish its allocation, as above, and then fork
 *  F children before sleeping.  Each child also sleeps forever.  If
 *  -fr is given, the children after forking, and the parent after all
 *  forks are done, does the equivalent of -r, read-touching each page
 *  of memory allocated; -fw and -fz are similar, but akin to -w and -z
 *  as -fr is to -r.  If -fork is not given, or if F is zero, the
 *  -f[rwz] flags have no effect.  If -fork is given with nonzero F,
 *  and any of -f[rwz] are given, mmap uses malloc() to allocate memory
 *  to remember what pages to touch after forking.
 *
 * Normally the mapping is done MAP_PRIVATE.  -s makes it use
 *  MAP_SHARED instead.
 *
 * -wire (which also can be spelled -lock) causes mmap to mlock() the
 *  mmapped memory in-core as soon as it's allocated.  Using -wireall
 *  (or -lockall) uses mlockall(MCL_CURRENT|MCL_FUTURE) instead.
 *
 * -spin (which can also be spelled -loop) causes the "sleep forever"
 *  to not really be a sleep; instead, the process sets its nice value
 *  to +20 and goes into an infinite loop.
 *
 * Note that children "sleep forever" in a different way from the
 *  parent, in order to ensure that they die if the parent does.
 */

#include <poll.h>
#include <stdio.h>
#include <errno.h>
#include <stdlib.h>
#include <stdarg.h>
#include <signal.h>
#include <unistd.h>
#include <strings.h>
#include <sys/mman.h>
#include <sys/sysctl.h>
#include <sys/resource.h>

extern const char *__progname;

static unsigned int U = 1048576;
static unsigned int N;
#define M_ZERO  0 /* command-line arg = 0 */
#define M_ALLOC 1 /* command-line arg > 0 */
#define M_ALBUT 2 /* command-line arg < 0 */
static int mode = M_ZERO;
static int vflag = 0;
static int rflag = 0;
static int wflag = 0;
static int zflag = 0;
static unsigned int F = 0;
static int frflag = 0;
static int fwflag = 0;
static int fzflag = 0;
static int mapflag = MAP_PRIVATE;
static int wireflag = 0;
static int wireallflag = 0;
static int spinflag = 0;

typedef struct blk BLK;

struct blk {
  BLK *link;
  char *base;
  int npgs;
  char *end;
  } ;

static int remember;
static BLK *mem = 0;
static BLK *lastmem = 0;
static pid_t mypid;

static int signed_number(const char *s)
{
 char *ep;

 if (index(s,'.'))
  { double v;
    int m;
    v = strtod(s,&ep);
    if (s == ep)
     { fprintf(stderr,"%s: %s: invalid number\n",__progname,s);
       exit(1);
     }
    switch (*ep++)
     { case 'k': case 'K':
	  m = 1 << 10;
	  break;
       case 'm': case 'M':
	  m = 1 << 20;
	  break;
       case 'g': case 'G':
	  m = 1 << 30;
	  break;
       default:
	  m = 1;
	  ep --;
	  break;
     }
    if (*ep)
     { fprintf(stderr,"%s: %s: junk after number\n",__progname,s);
       exit(1);
     }
    return(v*m);
  }
 else
  { long int v;
    int m;
    v = strtol(s,&ep,0);
    if (s == ep)
     { fprintf(stderr,"%s: %s: invalid number\n",__progname,s);
       exit(1);
     }
    switch (*ep++)
     { case 'k': case 'K':
	  m = 1 << 10;
	  break;
       case 'm': case 'M':
	  m = 1 << 20;
	  break;
       case 'g': case 'G':
	  m = 1 << 30;
	  break;
       default:
	  m = 1;
	  ep --;
	  break;
     }
    if (*ep)
     { fprintf(stderr,"%s: %s: junk after number\n",__progname,s);
       exit(1);
     }
    return(v*m);
  }
}

static unsigned int unsigned_number(const char *s)
{
 char *ep;

 if (index(s,'.'))
  { double v;
    int m;
    v = strtod(s,&ep);
    if ((s == ep) || (v < 0))
     { fprintf(stderr,"%s: %s: invalid number\n",__progname,s);
       exit(1);
     }
    switch (*ep++)
     { case 'k': case 'K':
	  m = 1 << 10;
	  break;
       case 'm': case 'M':
	  m = 1 << 20;
	  break;
       case 'g': case 'G':
	  m = 1 << 30;
	  break;
       default:
	  m = 1;
	  ep --;
	  break;
     }
    if (*ep)
     { fprintf(stderr,"%s: %s: junk after number\n",__progname,s);
       exit(1);
     }
    return(v*m);
  }
 else
  { unsigned long int v;
    int m;
    v = strtoul(s,&ep,0);
    if (s == ep)
     { fprintf(stderr,"%s: %s: invalid number\n",__progname,s);
       exit(1);
     }
    switch (*ep++)
     { case 'k': case 'K':
	  m = 1 << 10;
	  break;
       case 'm': case 'M':
	  m = 1 << 20;
	  break;
       case 'g': case 'G':
	  m = 1 << 30;
	  break;
       default:
	  m = 1;
	  ep --;
	  break;
     }
    if (*ep)
     { fprintf(stderr,"%s: %s: junk after number\n",__progname,s);
       exit(1);
     }
    return(v*m);
  }
}

static unsigned int pgsz(void) __attribute__((__const__));
static unsigned int pgsz(void)
{
 static unsigned int ps = 0;

 if (ps == 0)
  { int mib[2];
    int v;
    size_t vsz;
    mib[0] = CTL_HW;
    mib[1] = HW_PAGESIZE;
    vsz = sizeof(v);
    sysctl(&mib[0],2,&v,&vsz,0,0);
    ps = v;
    if (ps & (ps-1))
     { fprintf(stderr,"%s: pagesize (%d) is not a power of two\n",__progname,v);
       exit(1);
     }
  }
 return(ps);
}

static unsigned int pgshf(void) __attribute__((__const__));
static unsigned int pgshf(void)
{
 static unsigned int psh = 0;

 if (psh == 0)
  { unsigned int sz;
    sz = pgsz();
    while (sz > 1)
     { psh ++;
       sz >>= 1;
     }
  }
 return(psh);
}

static unsigned int pageroundup(unsigned int n)
{
 return((n+pgsz()-1)&~(pgsz()-1));
}

static unsigned int pagecount(unsigned int n)
{
 return((n+pgsz()-1)>>pgshf());
}

static void handleargs(int ac, char **av)
{
 int skip;
 int errs;
 int argno;

 skip = 0;
 errs = 0;
 argno = 0;
 for (ac--,av++;ac;ac--,av++)
  { if (skip > 0)
     { skip --;
       continue;
     }
  if (**av != '-')
     { switch (argno++)
	{ case 0:
	      { signed int a;
		a = signed_number(*av);
		if (a < 0)
		 { N = pageroundup(-a);
		   mode = M_ALBUT;
		 }
		else if (a > 0)
		 { N = pageroundup(a);
		   mode = M_ALLOC;
		 }
	      }
	     break;
	  default:
	     fprintf(stderr,"%s: extra argument `%s'\n",__progname,*av);
	     errs ++;
	     break;
	}
       continue;
     }
    if (0)
     {
needarg:;
       fprintf(stderr,"%s: %s needs a following argument\n",__progname,*av);
       errs ++;
       continue;
     }
#define WANTARG() do { if (++skip >= ac) goto needarg; } while (0)
    if (!strcmp(*av,"-v"))
     { vflag ++;
       continue;
     }
    if (!strcmp(*av,"-r"))
     { rflag ++;
       continue;
     }
    if (!strcmp(*av,"-w"))
     { wflag ++;
       continue;
     }
    if (!strcmp(*av,"-z"))
     { zflag ++;
       continue;
     }
    if (!strcmp(*av,"-fr"))
     { frflag ++;
       continue;
     }
    if (!strcmp(*av,"-fw"))
     { fwflag ++;
       continue;
     }
    if (!strcmp(*av,"-fz"))
     { fzflag ++;
       continue;
     }
    if (!strcmp(*av,"-s"))
     { mapflag = MAP_SHARED;
       continue;
     }
    if (!strcmp(*av,"-wire") || !strcmp(*av,"-lock"))
     { wireflag = 1;
       continue;
     }
    if (!strcmp(*av,"-wireall") || !strcmp(*av,"-lockall"))
     { wireallflag = 1;
       continue;
     }
    if (!strcmp(*av,"-spin") || !strcmp(*av,"-loop"))
     { spinflag = 1;
       continue;
     }
    if (!strcmp(*av,"-unit"))
     { WANTARG();
       U = pageroundup(unsigned_number(av[skip]));
       continue;
     }
    if (!strcmp(*av,"-fork"))
     { WANTARG();
       F = unsigned_number(av[skip]);
       continue;
     }
#undef WANTARG
    fprintf(stderr,"%s: unrecognized option `%s'\n",__progname,*av);
    errs ++;
  }
 if (argno < 1)
  { fprintf(stderr,"%s: need memory size argument\n",__progname);
    errs ++;
  }
 else if (mode == M_ZERO)
  { fprintf(stderr,"%s: memory size argument must be nonzero\n",__progname);
    errs ++;
  }
 if (errs)
  { fprintf(stderr,"Usage: %s [-unit U] [-v] [-r] [-w] [-z] N\n",__progname);
    exit(1);
  }
}

static void vpr(const char *, ...)
	__attribute__((__format__(__printf__,1,2)));
static void vpr(const char *fmt, ...)
{
 va_list ap;

 if (! vflag) return;
 va_start(ap,fmt);
 vprintf(fmt,ap);
 va_end(ap);
 fflush(stdout);
}

static void sleep_forever(void) __attribute__((__noreturn__));
static void sleep_forever(void)
{
 if (spinflag)
  { int n;
    setpriority(PRIO_PROCESS,0,20);
    n = 0;
    while (1)
     { volatile int i;
       vpr("spinnning: %d\n",n);
       for (i=1<<20;i>0;i--) ;
       n ++;
     }
  }
 else
  { sigset_t s;
    sigemptyset(&s);
    vpr("sleeping...\n");
    while (1) sigsuspend(&s);
  }
}

static void record_mem(void *vp, int n)
{
 BLK *b;
 BLK *l;
 BLK *p;

 if (! remember) return;
 n = pagecount(n);
 if (! mem)
  { mem = malloc(sizeof(BLK));
    mem->link = 0;
    mem->base = vp;
    mem->npgs = n;
    mem->end = mem->base + (mem->npgs << pgshf());
    lastmem = mem;
    return;
  }
 if ((char *)vp == lastmem->end)
  { lastmem->npgs += n;
    return;
  }
 if ((n<<pgshf())+(char *)vp == mem->base)
  { mem->base = vp;
    mem->npgs += n;
    return;
  }
 b = malloc(sizeof(BLK));
 b->base = vp;
 b->npgs = n;
 b->end = b->base + (n << pgshf());
 if (b->base > lastmem->end)
  { b->link = 0;
    lastmem->link = b;
    lastmem = b;
    return;
  }
 if (b->end < mem->base)
  { b->link = mem;
    mem = b;
    return;
  }
 p = 0;
 for (l=mem;l;p=l,l=l->link)
  { if (b->base < l->base)
     { if (b->end > l->base)
	{ fprintf(stderr,"%s: overlapping blocks!\n",__progname);
	  fprintf(stderr,"  have %d@%p, now got %d@%p\n",l->npgs,(void *)l->base,n,vp);
	  exit(1);
	}
       if (! p) abort();
       if (b->end == l->base)
	{ l->base = b->base;
	  l->npgs += b->npgs;
	  free(b);
	  if (p->end == l->base)
	   { p->npgs += l->npgs;
	     p->link = l->link;
	     free(l);
	   }
	  return;
	}
       if (p->end == b->base)
	{ p->npgs += b->npgs;
	  free(b);
	  return;
	}
       p->link = b;
       b->link = l;
       return;
     }
  }
}

static void allocmem(void)
{
 unsigned int left;
 unsigned int done;
 int i;
 void *mmrv;

 left = N;
 done = 0;
 while (left > 0)
  { unsigned int a;
    a = U;
    if (a > left) a = left;
    mmrv = mmap(0,a,PROT_READ|PROT_WRITE,MAP_ANON|mapflag,-1,0);
    if (mmrv == MAP_FAILED)
     { fprintf(stderr,"%s: mmap %u failed (after %u done): %s\n",__progname,a,done,strerror(errno));
       exit(1);
     }
    vpr("mmap %u ok...",a);
    record_mem(mmrv,a);
    if (wireflag)
     { vpr("lock...");
       if (mlock(mmrv,a) < 0)
	{ fprintf(stderr,"%s: mlock failed: %s\n",__progname,strerror(errno));
	}
     }
    if (rflag)
     { vpr("read...");
       for (i=0;i<a;i+=pgsz()) ((volatile char *)mmrv)[i];
     }
    if (wflag)
     { vpr("write...");
       for (i=0;i<a;i+=pgsz()) ((volatile char *)mmrv)[i] = 0;
     }
    if (zflag)
     { vpr("bzero...");
       bzero(mmrv,a);
     }
    vpr("\n");
    done += a;
    left -= a;
  }
}

static void allocbut(void)
{
 unsigned int left;
 unsigned int done;
 int i;
 void *mmrv;
 int nch;

 nch = (N+U-1) / U;
  { void *chunks[nch];
    int chh;
    chh = 0;
    done = 0;
    while (1)
     { mmrv = mmap(0,U,PROT_READ|PROT_WRITE,MAP_ANON|MAP_PRIVATE,-1,0);
       if (mmrv == MAP_FAILED) break;
       vpr("mmap %u ok...",U);
       record_mem(mmrv,U);
       if (wireflag)
	{ vpr("lock...");
	  if (mlock(mmrv,U) < 0)
	   { fprintf(stderr,"%s: mlock failed: %s\n",__progname,strerror(errno));
	   }
	}
       if (rflag)
	{ vpr("read...");
	  for (i=0;i<U;i+=pgsz()) ((volatile char *)mmrv)[i];
	}
       if (wflag)
	{ vpr("write...");
	  for (i=0;i<U;i+=pgsz()) ((volatile char *)mmrv)[i] = 0;
	}
       if (zflag)
	{ vpr("bzero...");
	  bzero(mmrv,U);
	}
       vpr("\n");
       if (--chh < 0) chh = nch - 1;
       chunks[chh] = mmrv;
       done += U;
     }
    if (done < N)
     { fprintf(stderr,"%s: allocated only %u, can't free %u\n",__progname,done,N);
       exit(1);
     }
    left = N;
    while (left > 0)
     { if (left > U)
	{ munmap(chunks[chh],U);
	  left -= U;
	}
       else
	{ munmap(chunks[chh],left);
	  left = 0;
	}
       if (++chh >= nch) chh = 0;
     }
  }
}

static void map_recorded(void (*fn)(void *, int))
{
 BLK *b;

 for (b=mem;b;b=b->link) (*fn)(b->base,b->npgs);
}

static void r_pages(void *base, int npgs)
{
 int i;
 char *bp;

 for (i=0,bp=base;i<npgs;i++,bp+=pgsz()) *((volatile char *)bp);
}

static void w_pages(void *base, int npgs)
{
 int i;
 char *bp;

 for (i=0,bp=base;i<npgs;i++,bp+=pgsz()) *((volatile char *)bp) = 0;
}

static void z_pages(void *base, int npgs)
{
 bzero(base,npgs<<pgshf());
}

static void forktouch(void)
{
 mypid = getpid();
 if (frflag)
  { vpr("read[%d]...\n",(int)mypid);
    map_recorded(r_pages);
  }
 if (fwflag)
  { vpr("write[%d]...\n",(int)mypid);
    map_recorded(w_pages);
  }
 if (fzflag)
  { vpr("bzero[%d]...\n",(int)mypid);
    map_recorded(z_pages);
  }
}

static void kidread(int fd)
{
 char junk;

 if (read(fd,&junk,1) >= 0) exit(0);
 switch (errno)
  { case EWOULDBLOCK:
    case EINTR:
       return;
       break;
  }
 fprintf(stderr,"%s: child read: %s\n",__progname,strerror(errno));
 exit(1);
}

static void kidsleep(int fd)
{
 if (spinflag)
  { struct pollfd pfd;
    int n;
    int p;
    setpriority(PRIO_PROCESS,0,20);
    n = 0;
    while (1)
     { int i;
       vpr("spinnning: %d\n",n);
       for (i=1<<10;i>0;i--)
	{ pfd.fd = fd;
	  pfd.events = POLLIN | POLLRDNORM;
	  p = poll(&pfd,1,0);
	  if (p < 0)
	   { switch (errno)
	      { case EWOULDBLOCK:
		case EINTR:
		   continue;
		   break;
	      }
	     fprintf(stderr,"%s: child poll: %s\n",__progname,strerror(errno));
	     exit(1);
	   }
	  if (p == 0) continue;
	  kidread(fd);
	}
       n ++;
     }
  }
 else
  { vpr("sleeping[%d]...\n",mypid);
    while (1) kidread(fd);
  }
}

static void forkstuff(void)
{
 unsigned int i;
 pid_t kid;
 int kidpipe[2];

 if (pipe(kidpipe) < 0)
  { fprintf(stderr,"%s: pipe: %s\n",__progname,strerror(errno));
    exit(1);
  }
 for (i=0;i<F;i++)
  { kid = fork();
    if (kid < 0)
     { fprintf(stderr,"%s: fork: %s\n",__progname,strerror(errno));
       exit(1);
     }
    if (kid == 0)
     { close(kidpipe[1]);
       forktouch();
       kidsleep(kidpipe[0]);
     }
    vpr("forked %d...\n",(int)kid);
  }
 close(kidpipe[0]);
 forktouch();
}

int main(int, char **);
int main(int ac, char **av)
{
 handleargs(ac,av);
 remember = (F > 0) && (frflag || fwflag || fzflag);
 if (wireallflag)
  { if (mlockall(MCL_CURRENT|MCL_FUTURE) < 0)
     { fprintf(stderr,"%s: mlockall: %s\n",__progname,strerror(errno));
       exit(1);
     }
  }
 switch (mode)
  { case M_ALLOC:
       allocmem();
       break;
    case M_ALBUT:
       allocbut();
       break;
    default:
       abort();
       break;
  }
 if (F) forkstuff();
 sleep_forever();
}