/*
	This program is designed to use up computer cycles in such a way
	as to stress the system.  Specifically, it does the following:
		- uses CPU cycles at a given duty cycle, loop time, 
		  & priority for a specified amount of time
		- optionally accesses additional memory to force paging
		- keeps track of time used and reports actual vs.
		  expected

	Runstring:
*/
#define runstring "stressor  [-l loop time (secs)] \\ \n \
	[-r randomness (0-100)] [-m memory (pages)] \\ \n \
	[-t time to run (seconds)] [-v (verbose)] \\ \n \
	[-V (very verbose)] [-d duty-cycle (1-100%) \\ \n \
	[-i (use integer math)] [-p priority] [-P RR|FIFO|OTHER ] \n"

/*
	In addition, the signal SIGUSR1 will toggle verbose mode.
*/

#include <stdio.h>
#include <unistd.h>
#include <sys/resource.h>
#include <stdlib.h>
#include <math.h>
#include <signal.h>
#include <sys/time.h>
#include <sys/times.h>
#include <sched.h>

#define PIDHEADER printf("stressor(%d): ",mypid)
#define FPIDHEADER fprintf(stderr,"stressor(%d): ",mypid)
#define POLICY(x)  x ? x-1 ? "RR" : "FIFO" : "OTHER"

int verbose=0;	/* 0=none, 1=verbose, 2+=Very Verbose */
int integer=0;	/* 0=floating point, 1=integer math in busy loop */
int running;	/* for signaling from SIGALRM handler */

int bebusy(float time);
int besleepy(float time);
int mypid;		/* this processes PID */
int pages=0;		/* number of memory pages to access for paging */
char *memptr;		/* pointer variable for paging access */
void usrhandler(int sig);	/* handler for SIGUSR1 */
void alarmhandler(int sig);	/* handler for SIGALRM */


main(int argc, char *argv[])
{
  int   c;		/* generic single character */
  int   optprobs=0; 	/* any problems with options? */

  int  	loop=1, 	/* loop time in seconds */
	prio=0, 	/* priority to run at */
	policy=-1,	/* scheduling policy */
	passedpolicy,
	pmin,		/* priority min and max */
	pmax;
  char  policystring[10]; /* temp holding for the passed string */
  struct sched_param 
  	  sched_parms;	/* scheduling parameters */
  float duty=50.,	/* duty cycle 1-100% */
	random=0.;	/* randomness 0-100 */
  float	time=0;		/* time to run in seconds */
  float looptime;	/* temp variable */
  int   secs,		/* for scheduling calls */
  	nsecs;		
  int	timeleft=1;	/* logical flag */

  struct sigaction 
        usraction,	/* sig handler structures */
  	alrmaction,	
	oldaction;
  
  struct tms timebuf;	/* used to hold the time usage info */
  int   i=0;		/* status keeping stuff */
  float t=0.;
  float min=1000,max=0;
  float minloop;

  /* "standard" option parsing... */

  mypid=(int)getpid();

  while ( (c=getopt(argc, argv, "d:il:m:p:P:r:t:?vV")) != EOF)
  {
    switch (c) {
      case 'd':
        sscanf(optarg,"%f",&duty); 
	if (duty < 1 || duty > 100 ) {
	  FPIDHEADER;
	  fprintf(stderr,
	  	"stressor: duty out of range(1-100%): %f.\n",duty);
	  optprobs=1;
	} 
        break;
      case 'i':
        integer=1;
        break;
      case 'l':
        sscanf(optarg,"%d",&loop); 
	if (loop < 1) {
	  FPIDHEADER;
	  fprintf(stderr,
	   	"stressor: loop time cannot be less than 1: %d\n",loop);
	  optprobs=1;
	}
        break;
      case 'm':
        sscanf(optarg,"%d",&pages); 
	if (pages < 0) {
	  FPIDHEADER;
	  fprintf(stderr,
	   	"stressor: memory pages cannot be negative: %d\n",pages);
	  optprobs=1;
	}
        break;
      case 'p':
        sscanf(optarg,"%d",&prio); 
	/* Value is checked later - can depend on the sched policy */
	break;
      case 'P':
        sscanf(optarg,"%s",policystring);
	for ( i=0 ; i<3 ; i++ ) {
	  if ( strcmp(policystring, POLICY(i)) == 0 ) {
	    policy=i;
	    break;
	  }
	}
	if (policy < 0 ) {
	  FPIDHEADER;
	  fprintf(stderr,"stressor: illegal value for policy: %s.\n",
	  	policystring);
	  fprintf(stderr,"  policy string must be one of:\n");
	  for ( i=0 ; i<3 ; i++) 
	    fprintf(stderr,"     %s\n",POLICY(i));

	  optprobs=1;
	}
        break;
      case 'r':
        sscanf(optarg,"%f",&random); 
	if (random < 0 || random > 100 ) {
	  fprintf(stderr,
	   	"stressor: Random must be between 0-100.\n");
	  optprobs=1;
	} 
        break;

      case 't':
        sscanf(optarg,"%f",&time); 
	if (loop < 1) {
	  fprintf(stderr,
	   	"stressor: time to run must be at least one second\n");
	  optprobs=1;
	}
        break;
      case 'v':
        verbose+=1;	/* account for -V also being set */
        break;
      case 'V':
        verbose=2;
        break;
      case '?':
        printf("Runstring:\n%s",runstring);
	exit(0);
      default:
	optprobs=1;	/* we'll deal with this problem later */
        break;
    }
  }
  /*
		Sanity time - find and fix conflicts....

  */

  if (optprobs) {
    fprintf(stderr,"Runstring:\n");
    fprintf(stderr,runstring);
    exit(1);
  }

  PIDHEADER;
  printf("Starting...\n");

  if (time == 0) {
    time=10.;	/* default to ten seconds */
    if (verbose) {
      PIDHEADER;
      printf("using default time of %f seconds.\n",time);
    }
  }
  /*  
     if policy was not set, default it.....    */
  if (policy < 0 ) {
    /* Policy was not set - default it appropriately...
        prio=0  policy=0 (SCHED_OTHER)
	prio>0  policy=1 (SCHED_RR)
    */ 
    if (prio > 0 )
      policy=SCHED_RR;
    else
      policy=SCHED_OTHER;

    PIDHEADER;
    printf("defaulting scheduling policy to SCHED_%s.\n",
      POLICY(policy));
  }
  
  /* one last check, in case an illegal policy was set... */
  passedpolicy=policy;
  if (prio == 0 ) {
    if ( policy != SCHED_OTHER )
      policy=SCHED_OTHER;
  } else 
    if ( ( policy != SCHED_FIFO ) && ( policy != SCHED_RR ) )
      policy=SCHED_RR;

  if ( policy != passedpolicy ) { /* musta been a problem - report it */
    FPIDHEADER;
    fprintf(stderr,"Policy SCHED_%s is illegal for prio %d.\n",
        POLICY(passedpolicy), prio );
    fprintf(stderr,"  (%d): changing policy to SCHED_%s.\n",
       mypid,POLICY(policy) );
  }

  if (verbose) {
    PIDHEADER;
    printf("options: \n");
    printf("  (%d): loop is %d seconds\n",mypid,loop);
    printf("  (%d): priority is %d\n",mypid,prio);
    printf("  (%d): scheduling policy is SCHED_%s\n",mypid,
        POLICY(policy) );
      
    printf("  (%d): paging memory %d pages\n",mypid,pages);
    if (time == 0) 
      printf("  (%d): time to run is indefinitely...\n",mypid);
    else
      printf("  (%d): time to run is %f seconds\n",mypid,time);
    printf("  (%d): %s math will be used in the busy loop\n",
      mypid, integer ? "Integer" : "Floating Point" );
    printf("  (%d): verbose mode: %s\n",mypid, verbose-1 ? "Very" : "On");
    printf("  (%d): duty cycle: %.0f%%.\n",mypid,duty);
    printf("  (%d): randomness is %.0f%%.\n",mypid,random);
  }

  /* convert duty cycle and randomness from percentages to factors */
  duty/=100.;
  random/=100.;

  /*
     set up the signal handlers...
  */
  if (verbose) {
    PIDHEADER;
    printf("Setting up handler for signal %d.\n",SIGUSR1);
  }
  if (signal(SIGUSR1,&usrhandler)) {
    fprintf(stderr,"Error from signal...\n");
    perror("stressor");
    exit(1);
  }

  if (verbose) {
    PIDHEADER;
    printf("Setting up handler for signal %d.\n",SIGALRM);
  }
  if (signal(SIGALRM,&alarmhandler)) {
    fprintf(stderr,"Error from signal...\n");
    perror("stressor");
    exit(1);
  }


  /* for some reason, this doesn't seem to work... */
#ifdef not
  usraction.sa_handler=usrhandler;
  usraction.sa_sigaction=0;
  if (sigaction(SIGUSR1,&usraction,&oldaction)) {
    fprintf(stderr,"Error from sigaction...\n");
    perror("help!:");
    exit(1);
  }
  printf("Setting up handler for signal %d.\n",SIGALRM);
  usraction.sa_handler=alarmhandler;
  usraction.sa_sigaction=0;
  if  (sigaction(SIGALRM, &usraction, &oldaction)) {
    fprintf(stderr,"Error from sigaction...\n");
    perror("help!:");
    exit(1);
  }
#endif
  /*
  	If memory pages have been asked for, attempt to allocate them
	and make them available.
  */
  if (pages > 0) {
    if ((memptr=malloc(pages*4096)) == 0) {
      fprintf(stderr,"(%d): could not allocate %d pages.\n",mypid,pages);
      fprintf(stderr,"(%d): paging is disabled.\n",mypid);
      pages=0;
    }
    else
      if (verbose) 
        printf("stressor(%d): Allocated %d bytes of paging memory.\n",
	   mypid, pages*4096);

  }
  /*
  	Almost there.  Now it's time to set our priority and/or
	scheduling policy.  Since the priority range can depend
	on the policy the range is determined from the policy, 
	then both policy and priority are set if all is right.
  */
  pmin=sched_get_priority_min(policy);
  pmax=sched_get_priority_max(policy);
  if ( (prio  < pmin) || (prio > pmax) ) {
    fprintf(stderr,"(%d): priority is out of range: %d.\n",mypid,prio);
    fprintf(stderr,"(%d): must be between %d and %d.\n",mypid,pmin,pmax);
    fprintf(stderr,"(%d): for scheduling policy: SCHED_%s.\n",mypid,POLICY(policy));

    exit(1);
  }

  sched_parms.sched_priority=prio;
  if (sched_setscheduler(0, policy, &sched_parms)) {
    fprintf(stderr,"(%d): Failuring setting priority to %d.\n",mypid,prio);
    perror("stressor");
    exit(1);
  }

  /*
  	After all that - we're ready to get started.  The basic drill
	is we set our priority, then run in a loop.  The length is
	"loop" seconds and we are busy for "duty" percent, sleep for
	100-"duty".  Both numbers are randomized by multiplying times 
	1-"random".  During the busy cycle, if "pages" is non-zero, the
	the memory pages are accesssed randomly.
  */

  srand48((long)24011956);	/* seed the random numbers */
  minloop = loop * duty * (1. - random*.5);
  while(timeleft) {
    /* Calculate the busy time */
    if (verbose > 1) {
      PIDHEADER;
      printf("New loop, %.2f seconds left.\n",time);
    }

    /*   this might look a bit odd.  The point is that if randomness
         is 0, looptime should be loop times the duty cycle percentage.
	 If randomness is not zero, the result should be randomly
	 distributed around the loop time duty cycle value - such that
	 the average works out to loop - that's why randomness is
	 multiplied times .5 in the final term - keeps the value from
	 dropping to zero - not good.  
    */
    looptime = loop * duty * ((1. - random*drand48()) + random*.5);

    if (looptime > time)
      looptime = time;
    if (looptime < minloop)	/* let's don't get too small */
      looptime = minloop;
    
    t+=looptime;		/* keep some stats */
    i++;
    if (looptime < min) min=looptime;
    if (looptime > max) max=looptime;
     
    if ( (time-=looptime) <= 0. ) {
      time=0;
      timeleft=0;	/* Time for this one to go home */
    }

    bebusy(looptime);	/* go be busy */
    
    /* and again - this time for the sleepy time - note we use the other
       portion of the duty cycle, thus the '1-duty' */
    looptime = loop * (1.0 - duty) * 
    	((1. - random*drand48()) + random*.5);

    if (looptime > time)
      looptime = time;

    if (looptime != 0 )
      besleepy(looptime);
      
    if ( (time-=looptime) <= 0. ) {
      time=0;
      timeleft=0;	/* Time for this one to go home */
    }
  }

  PIDHEADER;
  printf("Done - %d loops completed.\n",i);
		
  times(&timebuf);	/* how much time did we use? */
  if (verbose>1) {
    PIDHEADER;
    printf("Total time used is %.2f seconds in %d loops.\n",t,i);
    printf("stressor(%d):  Average is: %f \n",mypid,t/i);
    printf("stressor(%d):  Minimum loop was %f \n",mypid,min);
    printf("stressor(%d):  Maximum loop was %f \n",mypid,max);
    printf("stressor(%d):  minloop is %f \n",mypid,minloop);

    printf("stressor(%d):  Actual usage from times() in seconds: \n",mypid);
    printf("stressor(%d):   User time: %f.\n",
	mypid, (float)timebuf.tms_utime/100.);
    printf("stressor(%d):   System time: %f.\n",
	mypid, (float)timebuf.tms_stime/100.);
  }
  /* print out expected vs. actual... */
  printf("stressor(%d):  Actual/Expected = %.2f / %.2f = %.2f%%.\n",
     mypid, (float)timebuf.tms_utime/100.,t,(float)timebuf.tms_utime/t);




}

int bebusy(float time)
{
  int i,j,max;
  int index;
  struct itimerval timerval;
  float a, b, c, pi;
  int m,n;

  /*
    use setitimer to set up a siglarm, then loop around like an
    idiot, until we get the sigalarm, or we run out of cycles - running
    out of cycles is a problem - means we missed a sigalarm. (Assuming
    the loops/per second value is set correctly.
  */

  if ((verbose-1)>0) {
    PIDHEADER;
    printf("Busy time: %f seconds.\n",time);
  }


  i=0;
  running=1;	/* needs to be set *before* the setitimer call */

  timerval.it_value.tv_sec=(int)time;	/* whole number part */
  timerval.it_value.tv_usec=(time-timerval.it_value.tv_sec)*1000000; /*frac*/
  timerval.it_interval.tv_sec=0;
  timerval.it_interval.tv_usec=0;

  if (setitimer(ITIMER_REAL,&timerval,0) ) {
    fprintf(stderr,"stressor(%d): Failure setting timer!\n");
    perror("stressor:");
    exit(1);
  }
  

  max=time*1000000;  /* Upper limit of how long to run */
  if ((verbose-1)>0) {
    PIDHEADER;
    printf("Max loops: %d.\n",max);
  }


  if ( integer )
  {	/* integer busy loop */
    m=32666;	/* arbitrary number */
    while ( running && (i<max) ) {
      /* The number of loops is adjusted to so that this loop is
         approximately the same timing as the floating point loop */
      for ( j=0 ; j<50 ; j++ ) {    
        if (m < 41234 ) {	    /* another arbitrary number */
          n=m*3;	/* numbers get bigger  */
          m=n/2;
        }
        else {
          n=m*5;  /* number get smaller */
          m=n/7;
        }
        if (pages) {	/* are we forcing memory accesses? */
          /*index=drand48()*4096*pages; */
	  index=(lrand48() & 0xFFF ) * pages; /* mask to 0-4095 */
          memptr[index]+=i;
        }
      }
      i++;
    }
  } else
  {

    /* just some goofiness in a loop... */
    a=3.1415926539;
    while ( running && (i<max) ) {
       if (pages) {	/* are we forcing memory accesses? */
         index=drand48()*4096*pages;
         b=index/6000.;
         c=exp(b);
         b=log(c+a);
         memptr[index]+=1;
       } else {
         b=a*drand48();
         c=exp(b);
         b=log(c+a);
       }
  
       i++;
    }
  }
  if (i>=max) {	   /* this means we did NOT get a SIGALRM! */
    PIDHEADER;
    printf("did not receive a SIGALRM within the expected time!\n");
  }
  if (verbose > 1) {
    PIDHEADER;
    printf("loop done: i is %d, max is %d.\n",i,max);
    PIDHEADER;
    printf("loops per second is about %.0f.\n",i/time);
  }
  return 0;
}

int besleepy(float time)
{
  int secs;
  struct timespec timerval;
  
  if ((verbose-1)>0) {
    PIDHEADER;
    printf("Sleepy time: %f seconds.\n",time);
  }
  timerval.tv_sec=(int)time;	/* whole number part */
  timerval.tv_nsec=(time-timerval.tv_sec)*1000000000; /*frac*/
  if (nanosleep(&timerval,NULL) ) {
    fprintf(stderr,"stressor(%d): Failure in nanosleep!\n",mypid);
    perror("stressor:");
  }
  return 0;
}

void usrhandler(int sig)
{

  PIDHEADER;
  printf("Received SIGUSR1!!\n");

  if (verbose)
    verbose=0;
  else
    verbose=2;
  PIDHEADER;
  printf("Verbose flag is %s.\n", verbose ? "ON" : "OFF");

}

void alarmhandler(int sig)
{
  if (verbose>1) {
    PIDHEADER;
    printf("Received SIGALRM!! - parameter is %d.\n",sig);
  }
  running=0;
}