/*
  $Id: rlcat.c,v 2.9 2008/02/23 00:42:32 reid Exp $

    Rate-limited cat command
	rlcat -n 13 -d sec < file
    will cat the file at 13 lines per second.
	rlcat -n 13 -d sec -m xyzzy < file
    wile cat the file at 13 chunks per second, where each chunk is terminated
	by the string "xyzzy".

    Options:
	-d	denominator, one of "sec", "min", or "hr"
	-m	marker (delimits a chunk))
	-n	numerator, an integer
	-x	discard the marker (do not output it)
	-v	verbose--diagnostic output to stderr

	Brian Reid, ISC, September 2007

 */


     #include <ctype.h>
     #include <stdio.h>
     #include <stdlib.h>
     #include <unistd.h>
     #include <sys/stat.h>
     #include <errno.h>
     #include <time.h>
     #include <sys/time.h>
     #include <math.h>
     #include <signal.h>

/* ------------------------------------------------------------------ */

#define	TRUE		1
#define	FALSE		0

#define MAXCHUNK	20000
#define MAXFILENAME	1024
#define MAXMARKER	200		/* max length of -m argument */
#define MAXTICKS	300		/* max ticks between clock probes */
#define MSFUDGE		0.004		/* see init_timing */
#define SIGNALFUDGE	0.0001		/* see interval_sleep */
#define	TC_SMOOTH	50.0		/* see init_timing */
#define	PRECISION	0.0001		/* a ratio */
#define	EPSILON		0.000000001	/* smallest useful wait time */

#define MINIMUM_NANOSLEEP 0.000000001

/* Globals */

     char *numerator = NULL, *denominator = "sec";
     char marker[MAXMARKER+1] = "\n";
     char current_file[MAXFILENAME];

     int sendmarker = 1,
	verbose = 0;

     long double fnum=1.0,
	fdenom=1.0;

     long double 
	recent_send_rate,
     	desired_send_rate,
	max_send_rate,
	cumulative_send_rate,
	smoothed_average_send_rate,
	idealized_wait_interval,
	computed_wait_interval,
	alpha;

     double total_elapsed_seconds,
        interval_elapsed_seconds;

     long
	interval_chunks_needed,
        interval_chunks_sent,
        total_chunks_sent,
	chunk_goal_at_interval_start,
	chunk_goal_by_interval_end,
	chunk_goal_next_interval_end,
	chunk_count_at_interval_start,
	chunk_count_at_interval_end;

     int timer_correction_needed;

     struct timeval starting_time,
	last_clock_probe,
	previous_clock_probe;

     struct itimerval
	one_second_interval = { {1,0}, {1,0}};
     
/* ------------------------------------------------------------------ */

     usage () {
	   fprintf (stderr, "\n\
rlcat, a rate-limited cat(1)\n\
Usage:   rlcat -n <num> [-d <denom>] [-m <marker>] [-xv] [file ...] \n\
	<num> must be an integer\n\
	<denom> must be \"sec\", \"min\", or \"hr\" (default \"sec\")\n\
	<marker> is string that delimits chunks (default \"\\n\")\n\
	If -x is set, the markers are not output.\n\
	-v gives some diagnostic output to stderr\n\
\n\
	rlcat sends <num> delimited chunks of data every <denom> \n\
	and sends the delimiter too unless -x is set.\n\
\n\
	If -m is not set, the marker is a newline, which is to say rlcat\n\
	sends lines at a limited rate. \n\n\
\n\
	Examples:\n\
	 rlcat -n 24 -d min 		sends 24 lines per minute to stdout\n\
	 rlcat -n 300 -d sec -m \"\\n\\n\"  sends 300 blankline-delimited\n\
					chunks per second to stdout\n\n\
");
	    return (0);
       }

/* ------------------------------------------------------------------ */
/*
   This signal handler is called once per second from a timer. It takes
   some notes, tweaks the separation interval if needed, and then resumes
   execution. 
 */
      void
      alarm_tick_handler(int signum) {
	struct timeval
		elapsed,
		interval_elapsed;
	long double clock_correction, rate_correction;

	if (signum != SIGALRM) return;

   /* Take and process a timestamp */
	previous_clock_probe = last_clock_probe;
	gettimeofday(&last_clock_probe, NULL);
	timersub(&last_clock_probe, &starting_time, &elapsed);
	timersub(&last_clock_probe, &previous_clock_probe, &interval_elapsed);
	interval_elapsed_seconds = interval_elapsed.tv_sec +
		((double)interval_elapsed.tv_usec) / 1000000.00;
	total_elapsed_seconds = elapsed.tv_sec +
		((double)elapsed.tv_usec) / 1000000.00;

	chunk_count_at_interval_start = chunk_count_at_interval_end;
	chunk_count_at_interval_end = total_chunks_sent;

	chunk_goal_at_interval_start = chunk_goal_by_interval_end;
	chunk_goal_by_interval_end = chunk_goal_next_interval_end;
	chunk_goal_next_interval_end = (int) 1.0 + (desired_send_rate * (1.0+total_elapsed_seconds));
	recent_send_rate = chunk_count_at_interval_end - chunk_count_at_interval_start;
        smoothed_average_send_rate = alpha*smoothed_average_send_rate
	   +(1.0-alpha)*recent_send_rate;

	/* This computation doesn't work if we're sending slowly */
	if (desired_send_rate > 1.0) {
	    cumulative_send_rate = (double)total_chunks_sent/(1.0+total_elapsed_seconds);
	}

	if (cumulative_send_rate > 0.0) {
	    rate_correction = cumulative_send_rate / desired_send_rate;
	    clock_correction = 1.0/rate_correction;
	    if (abs(1.0-rate_correction) > PRECISION) {
		computed_wait_interval *= sqrt(clock_correction);
	    } 
	}

	if (verbose) {	
	    fprintf(stderr,"TICK: %6.2lf  Burst: %Lg Cum: %lg Avg: %lg (+%lg) Corr: %lG %Lg, %Lg\n",
	    total_elapsed_seconds, recent_send_rate, (double) cumulative_send_rate,
	    (double) smoothed_average_send_rate,
	    (double) (chunk_goal_next_interval_end - total_chunks_sent),
	    (double) rate_correction,
	    idealized_wait_interval,computed_wait_interval);
	}

 	signal(SIGALRM, alarm_tick_handler);
      }
/* ------------------------------------------------------------------ */

/* Initialize the timing apparatus */

      init_timing () {
	total_chunks_sent = 0;
	total_elapsed_seconds=0;
	gettimeofday(&starting_time, NULL);
	gettimeofday(&last_clock_probe, NULL);

	desired_send_rate = fnum / fdenom;
	smoothed_average_send_rate = desired_send_rate;

 /* 
    The idealized wait interval is the timing between chunks if the 
    computer were infininitely fast. MSFUDGE is a guess at the length
    of time it will take to execute the code loop, which is then 
    subtracted out of the ideal time. If we guess wrong, the timing
    refinement algorithm will fix it, iteratively, but if we start off
    more accurately it will converge faster.
  */

	idealized_wait_interval = (1.0 / desired_send_rate);
	computed_wait_interval = idealized_wait_interval - MSFUDGE;
	if (computed_wait_interval < 0.0) 
		computed_wait_interval = 0;

	max_send_rate = 0.0;

	recent_send_rate = -1;
	alpha = exp(-desired_send_rate/TC_SMOOTH);

	chunk_count_at_interval_start = 0;
	chunk_count_at_interval_end = 0;

	chunk_goal_at_interval_start = 0;
	chunk_goal_by_interval_end = 0;

	chunk_goal_next_interval_end = (int) desired_send_rate;

	signal(SIGALRM, alarm_tick_handler);

	if (setitimer(ITIMER_REAL, &one_second_interval, 0)) {
	    perror("setitimer");
	}
      }
    
/* ------------------------------------------------------------------ */

/* Program initialization and option processing */

     int init(int argc, char *argv[]) {

       char errbuf[512];
       int printusage = 0;
       struct stat statbuf;

       int index, in, out, j;
       char *jptr;
       int c;
     
       opterr = 0;
     
       while ((c = getopt (argc, argv, "d:hm:n:vx")) != -1)
         switch (c)
           {

/* -d option: specify denominator of the limit rate */
           case 'd':
             denominator = optarg;
	     fdenom = 0.0;
	     if (!strcmp(denominator,"sec")) fdenom=1.0;
	     else if (!strcmp(denominator,"min")) fdenom=60.0;
	     else if (!strcmp(denominator,"hr")) fdenom=3600.0;
	     if (fdenom == 0.0) {
		    fprintf(stderr, "Unknown denominator unit %s; must be hr, min, or sec.\n",denominator);
		    printusage++;
	     }
             break;

/* -h option: print help message and exit */
           case 'h':
             printusage++;
             break;

/* -m option: note the inter-record marker string */
	   case 'm':
	   /* BKR: I wasted entirely too much time looking for a 
		package or function or library function to do the
		escape processing here. Foo. I was tempted to do
		popen("perl") :-)
            */
	     out=0;
	     for (in=0; in<strlen(optarg); in++) {
	        switch (optarg[in]) {
		  case '\\':
		     jptr = &optarg[in];
		     j = optarg[++in];
		     switch (j) {
			case '\\': marker[out++] = '\\'; break;
			case 'n': marker[out++] = '\n'; break;
			case 'r': marker[out++] = '\r'; break;
			case 't': marker[out++] = '\t'; break;
			case 'f': marker[out++] = '\f'; break;
			default:
			   if (isdigit(j)) {
			       int n;
			       n = (int) j - (int) '0';
  /* We have to look ahead one character to see where the integer ends */
			       while (isdigit(optarg[++in])) {
			           j = optarg[in];
			           n = (n*8) + ((int)j - (int)'0');
			       }
			       marker[out++] = (char) n;
  /* Here we put back the lookahead character. We are counting on the
     optarg string being null-terminated to avoid an out-of-bounds 
     array reference (worst case, the lookahead stops on the null) */
			       in--;
			   } else {
			       fprintf(stderr,"\
Unrecognized escape sequence \\%c.\n\
Use \\nnn (base 8) for arbitrary characters.\n", j);
			       printusage++;
			   }
			   break;
		     }
		  break;
		  case (char) 0:
		  break;
		  default:
		    marker[out++] = optarg[in];
		}
	     }
	     marker[out++] = (char) 0;
	     break;

/* -n option: specify numerator of the limit rate */
             case 'n':
             numerator = optarg;
	     sscanf(numerator,"%Lg",&fnum);
	     if (fnum <= 0.0) {
	         fprintf(stderr, "Numerator must be positive definite. You specified %f\n\n", fnum);
		 printusage++;
	     }
             break;

/* -v option: run in verbose mode*/
	     case 'v': verbose++;
	     break;

/* -x option: set flag not to output the marker */
	    case 'x':
	      sendmarker = 0;
	      break;

/* unknown option */
           case '?':
	     fprintf(stderr, "Unknown option -%c\n\n",optopt);
             printusage++;
	     break;
           default:
             abort ();
           }
     
	for (index = optind; index < argc; index++) {
	    if (stat(argv[index], &statbuf) == -1) {
		snprintf(errbuf, sizeof(errbuf), "File %s",argv[index]);
		perror(errbuf);
		printusage++;
	    }
	}
	return (printusage);
     }


/* ------------------------------------------------------------------ */

 /* 
   Wait between chunk transmissions. This is complex because the hardware
   clock does not have enough precision, and in fact has unknown precision.
   We have to deal with whatever precision the hardware clock is able
   to deliver, and we can't even know what it is.
  */

	void
	interval_sleep (int reason_for_sleeping) {

#define SENDING	-1
#define IDLING	0
#define WAITING 1

	    long double 
	    	unslept_amount, slept_amount,
		sleep_time_remaining,
		theoretical_correction;
	    int sleep_seconds, interval_chunks_unsent;
	    struct timespec 
		nanosleep_time,
		unslept_time;

	    sleep_time_remaining = computed_wait_interval;
	    while (sleep_time_remaining > 0) {
		sleep_seconds = (int) sleep_time_remaining;
		nanosleep_time.tv_sec = sleep_seconds;
		nanosleep_time.tv_nsec = (int)
		  (1e9*(sleep_time_remaining-(long double) sleep_seconds));
		if (nanosleep(&nanosleep_time, &unslept_time) < 0) {
		    if (errno == EINTR) {
			if (reason_for_sleeping == IDLING) 
			   sleep_time_remaining = 0.0;
			else if (reason_for_sleeping == WAITING)
			   sleep_time_remaining = 0.0;
			else {
			    sleep_time_remaining =
			        (long double)unslept_time.tv_sec
			        + ( (long double) unslept_time.tv_nsec) / 1e9;
			    sleep_time_remaining -= SIGNALFUDGE;
			}
		    }
		} else 
		    sleep_time_remaining = 0;
	    }
     };

/* ------------------------------------------------------------------ */

/*	Send one file at a regulated speed 	*/

     void send_file (FILE *input_file) {

        char buffer[2*MAXCHUNK+MAXMARKER+1];
	int bufnext=0, buflast = 0;
	char *markerloc, *chunkloc;
	int chunksize, sendsize;

       /* this loop structure means that chunks cannot be split between
	  files. To do it otherwise would require a coroutine structure,
	  which is too hard to do in C */

	while (TRUE) {
	    size_t bytes_read;
	    if (buflast < MAXCHUNK) {
	       bytes_read = fread(&buffer[buflast], 1, MAXCHUNK, input_file);

	       if (bytes_read <= 0) {
	           if (feof(input_file) && bufnext >= buflast) return;
		   if (ferror(input_file)) {
		       perror(current_file);
		   }
	       }
	       buflast += (int) bytes_read;
	       buffer[buflast] = (char) 0;
	    }

	/* Find the marker at the end of the chunk. If it's not there,
	   we have to read in another bufferload. */

	    markerloc = strstr(&buffer[bufnext], marker);
	    if (markerloc == NULL) {
	        int fragmentsize;
		int i;

	        fragmentsize = buflast-bufnext+1;

	/* slide the fragment up to the top of the buffer to make
	   room for another block to be read just after it. */

		strncpy(buffer, &buffer[bufnext], fragmentsize);
		bufnext = 0;
		buflast = fragmentsize-1;
	    } else {
		
     /* the next section is the entire reason for the existence
        of this program. We wait a certain amount of time as
        determined by the rate limit, and then we send the
	chunk to standard output. */

		chunkloc = &buffer[bufnext];
		chunksize = markerloc - chunkloc;

		sendsize = chunksize;
		if (sendmarker) sendsize += strlen(marker);
		chunksize += strlen(marker);

		bufnext += chunksize;

		fwrite(chunkloc, 1, sendsize, stdout); 
		total_chunks_sent++;

		interval_sleep(SENDING);

		if (total_elapsed_seconds > 0) {
		    cumulative_send_rate = (double)total_chunks_sent/(1.0+total_elapsed_seconds);
		}

     /* Make sure we have quota to send; if not, wait for it. Our goal
        is to have the interval timer signal arrive in the first part
	of the IDLING sleep interval. 
      */

		if (total_chunks_sent >= chunk_goal_next_interval_end) {
		    interval_sleep(IDLING);
		}

		while (total_chunks_sent >= chunk_goal_next_interval_end) {
		    interval_sleep(WAITING);
		}
	    }
	}
	
     }

/* ------------------------------------------------------------------ */
/* ------------------------------------------------------------------ */

     int  main (int argc, char **argv)
     {
       int firstarg = 0, lastarg = 0;
       int index;
       char *unitsuffix="";
       FILE *file_to_send;

/* initialize */

       if (init(argc, argv)) { usage(); exit(0);}
       firstarg = optind;
       lastarg = argc;

/* say what we are going to do */

	if (fnum != 1.0) unitsuffix="s";	
	if (!strcmp(marker,"\n")) {
	    fprintf(stderr, "will cat %Lg line%s per %s\n",fnum,unitsuffix,denominator);
	    if (!sendmarker) fprintf(stderr, "Newline characters will not be sent\n");
        } else {
	    fprintf(stderr, "will cat %Lg block%s per second\n",fnum,unitsuffix);
	    fprintf(stderr, "block marker is %s. ", marker);
	    if (sendmarker) fprintf(stderr, "Marker will be transmitted.\n");
	    else fprintf(stderr, "Marker will not be transmitted.\n");
	}

/* Do it */
	init_timing();

	   /* get the file names out of argv, make sure they are
	      still there, and then send them. */

	for (index = firstarg; index < lastarg; index++) {
	    if (argv[index] == "-") {
  	        strncpy(current_file, "standard input", sizeof(current_file));
	        send_file(stdin);
	    } else {
		file_to_send = fopen(argv[index], "r");
		if (file_to_send == NULL) {
		    perror(argv[index]);
		    exit(1);
		}
	        strncpy(current_file, argv[index], sizeof(current_file));
 	        send_file(file_to_send);
	        fclose(file_to_send);
	    }
	}

	   /* If there were no files specified to send, then send stdin */

	if (firstarg >= lastarg) {
	    strncpy(current_file, "standard input", sizeof(current_file));
	    send_file(stdin);
	}

/* Say we did it */

	    fprintf(stderr,"Sent %d blocks in %2.2lf secs; mean rate=%Lg/sec\n",
		total_chunks_sent, (double) total_elapsed_seconds, cumulative_send_rate);


	exit(0);
     }
