cpulimit

/**
 * This program is licensed under the GNU General Public License,
 * version 2. A copy of the license can be found in the accompanying
 * LICENSE file.
 *
 **********************************************************************
 *
 * Simple program to limit the cpu usage of a process
 * If you modify this code, send me a copy please
 *
 * Author:  Angelo Marletta
 * Date:    26/06/2005
 * Version: 1.1
 *
 * Modifications and updates by: Jesse Smith
 * Date: May 4, 2011
 * Version 1.2
 *
 */


#include <getopt.h>
#include <stdio.h>
#include <stdlib.h>
#include <time.h>
#include <sys/time.h>
#include <unistd.h>
#include <sys/types.h>
#include <signal.h>
#include <sys/resource.h>
#include <string.h>
#include <dirent.h>
#include <errno.h>
#include <string.h>
#include <limits.h>    // for compatibility


//kernel time resolution (inverse of one jiffy interval) in Hertz
//i don't know how to detect it, then define to the default (not very clean!)
#define HZ 100

//some useful macro
#define min(a,b) (a<b?a:b)
#define max(a,b) (a>b?a:b)

// For platforms without PATH_MAX
#ifndef PATH_MAX
#define PATH_MAX 4096
#endif

//pid of the controlled process
int pid=0;
//executable file name
char *program_name;
//verbose mode
int verbose=0;
//lazy mode
int lazy=0;
// is higher priority nice possible?
int nice_lim;

//reverse byte search
void *memrchr(const void *s, int c, size_t n);

//return ta-tb in microseconds (no overflow checks!)
inline long timediff(const struct timespec *ta,const struct timespec *tb) {
    unsigned long us = (ta->tv_sec-tb->tv_sec)*1000000 + (ta->tv_nsec/1000 - tb->tv_nsec/1000);
    return us;
}

int waitforpid(int pid) {
	//switch to low priority
	// if (setpriority(PRIO_PROCESS,getpid(),19)!=0) {
        if ( (nice_lim < INT_MAX) && 
             (setpriority(PRIO_PROCESS, getpid(), 19) != 0) ) {
		printf("Warning: cannot renice\n");
	}

	int i=0;

	while(1) {

		DIR *dip;
		struct dirent *dit;

		//open a directory stream to /proc directory
		if ((dip = opendir("/proc")) == NULL) {
			perror("opendir");
			return -1;
		}

		//read in from /proc and seek for process dirs
		while ((dit = readdir(dip)) != NULL) {
			//get pid
			if (pid==atoi(dit->d_name)) {
				//pid detected
				if (kill(pid,SIGSTOP)==0 &&  kill(pid,SIGCONT)==0) {
					//process is ok!
                                        if (closedir(dip) == -1) {
                                           perror("closedir");
                                           return -1;
                                        }
					goto done;
				}
				else {
					fprintf(stderr,"Error: Process %d detected, but you don't have permission to control it\n",pid);
				}
			}
		}

		//close the dir stream and check for errors
		if (closedir(dip) == -1) {
			perror("closedir");
			return -1;
		}

		//no suitable target found
		if (i++==0) {
			if (lazy) {
				fprintf(stderr,"No process found\n");
				exit(2);
			}
			else {
				printf("Warning: no target process found. Waiting for it...\n");
			}
		}

		//sleep for a while
		sleep(2);
	}

done:
	printf("Process %d detected\n",pid);
	//now set high priority, if possible
	// if (setpriority(PRIO_PROCESS,getpid(),-20)!=0) {
        if ( (nice_lim < INT_MAX) &&
             (setpriority(PRIO_PROCESS, getpid(), nice_lim) != 0) ) {
		printf("Warning: cannot renice.\nTo work better you should run this program as root.\n");
	}
	return 0;

}

//this function periodically scans process list and looks for executable path names
//it should be executed in a low priority context, since precise timing does not matter
//if a process is found then its pid is returned
//process: the name of the wanted process, can be an absolute path name to the executable file
//         or simply its name
//return: pid of the found process
int getpidof(const char *process) {

	//set low priority
	// if (setpriority(PRIO_PROCESS,getpid(),19)!=0) {
        if ( (nice_lim < INT_MAX) &&
             (setpriority(PRIO_PROCESS, getpid(), 19) != 0) ) {
		printf("Warning: cannot renice\n");
	}

	char exelink[20];
	char exepath[PATH_MAX+1];
	int pid=0;
	int i=0;

	while(1) {

		DIR *dip;
		struct dirent *dit;

		//open a directory stream to /proc directory
		if ((dip = opendir("/proc")) == NULL) {
			perror("opendir");
			return -1;
		}

		//read in from /proc and seek for process dirs
		while ((dit = readdir(dip)) != NULL) {
			//get pid
			pid=atoi(dit->d_name);
			if (pid>0) {
				sprintf(exelink,"/proc/%d/exe",pid);
				int size=readlink(exelink,exepath,sizeof(exepath));
				if (size>0) {
					int found=0;
					if (process[0]=='/' && strncmp(exepath,process,size)==0 && size==strlen(process)) {
						//process starts with / then it's an absolute path
						found=1;
					}
					else {
						//process is the name of the executable file
						if (strncmp(exepath+size-strlen(process),process,strlen(process))==0) {
							found=1;
						}
					}
					if (found==1) {
						if (kill(pid,SIGSTOP)==0 &&  kill(pid,SIGCONT)==0) {
							//process is ok!
                                                        if (closedir(dip) == -1) {
                                                          perror("closedir");
                                                          return -1;
                                                        }
							goto done;
						}
						else {
							fprintf(stderr,"Error: Process %d detected, but you don't have permission to control it\n",pid);
						}
					}
				}
			}
		}

		//close the dir stream and check for errors
		if (closedir(dip) == -1) {
			perror("closedir");
			return -1;
		}

		//no suitable target found
		if (i++==0) {
			if (lazy) {
				fprintf(stderr,"No process found\n");
				exit(2);
			}
			else {
				printf("Warning: no target process found. Waiting for it...\n");
			}
		}

		//sleep for a while
		sleep(2);
	}

done:
	printf("Process %d detected\n",pid);
	//now set high priority, if possible
	// if (setpriority(PRIO_PROCESS,getpid(),-20)!=0) {
        if ( (nice_lim < INT_MAX) &&
             (setpriority(PRIO_PROCESS, getpid(), nice_lim) != 0) ) {
		printf("Warning: cannot renice.\nTo work better you should run this program as root.\n");
	}
	return pid;

}

//SIGINT and SIGTERM signal handler
void quit(int sig) {
	//let the process continue if it's stopped
	kill(pid,SIGCONT);
	printf("Exiting...\n");
	exit(0);
}

//get jiffies count from /proc filesystem
int getjiffies(int pid) {
	static char stat[20];
	static char buffer[1024];
        char *p;
	sprintf(stat,"/proc/%d/stat",pid);
	FILE *f=fopen(stat,"r");
	if (f==NULL) return -1;
	p = fgets(buffer,sizeof(buffer),f);
	fclose(f);
	// char *p=buffer;
        if (p)
        {
	  p=memchr(p+1,')',sizeof(buffer)-(p-buffer));
	  int sp=12;
	  while (sp--)
		p=memchr(p+1,' ',sizeof(buffer)-(p-buffer));
	  //user mode jiffies
	  int utime=atoi(p+1);
	  p=memchr(p+1,' ',sizeof(buffer)-(p-buffer));
	  //kernel mode jiffies
	  int ktime=atoi(p+1);
	  return utime+ktime;
        }
        // could not read info
        return -1;
}

//process instant photo
struct process_screenshot {
	struct timespec when;	//timestamp
	int jiffies;	//jiffies count of the process
	int cputime;	//microseconds of work from previous screenshot to current
};

//extracted process statistics
struct cpu_usage {
	float pcpu;
	float workingrate;
};

//this function is an autonomous dynamic system
//it works with static variables (state variables of the system), that keep memory of recent past
//its aim is to estimate the cpu usage of the process
//to work properly it should be called in a fixed periodic way
//perhaps i will put it in a separate thread...
int compute_cpu_usage(int pid,int last_working_quantum,struct cpu_usage *pusage) {
	#define MEM_ORDER 10
	//circular buffer containing last MEM_ORDER process screenshots
	static struct process_screenshot ps[MEM_ORDER];
	//the last screenshot recorded in the buffer
	static int front=-1;
	//the oldest screenshot recorded in the buffer
	static int tail=0;

	if (pusage==NULL) {
		//reinit static variables
		front=-1;
		tail=0;
		return 0;
	}

	//let's advance front index and save the screenshot
	front=(front+1)%MEM_ORDER;
	int j=getjiffies(pid);
	if (j>=0) ps[front].jiffies=j;
	else return -1;	//error: pid does not exist
	clock_gettime(CLOCK_REALTIME,&(ps[front].when));
	ps[front].cputime=last_working_quantum;

	//buffer actual size is: (front-tail+MEM_ORDER)%MEM_ORDER+1
	int size=(front-tail+MEM_ORDER)%MEM_ORDER+1;

	if (size==1) {
		//not enough samples taken (it's the first one!), return -1
		pusage->pcpu=-1;
		pusage->workingrate=1;
		return 0;
	}
	else {
		//now we can calculate cpu usage, interval dt and dtwork are expressed in microseconds
		long dt=timediff(&(ps[front].when),&(ps[tail].when));
		long dtwork=0;
		int i=(tail+1)%MEM_ORDER;
		int max=(front+1)%MEM_ORDER;
		do {
			dtwork+=ps[i].cputime;
			i=(i+1)%MEM_ORDER;
		} while (i!=max);
		int used=ps[front].jiffies-ps[tail].jiffies;
		float usage=(used*1000000.0/HZ)/dtwork;
		pusage->workingrate=1.0*dtwork/dt;
		pusage->pcpu=usage*pusage->workingrate;
		if (size==MEM_ORDER)
			tail=(tail+1)%MEM_ORDER;
		return 0;
	}
	#undef MEM_ORDER
}

void print_caption() {
	printf("\n%%CPU\twork quantum\tsleep quantum\tactive rate\n");
}

void print_usage(FILE *stream,int exit_code) {
	fprintf(stream, "Usage: %s TARGET [OPTIONS...]\n",program_name);
	fprintf(stream, "   TARGET must be exactly one of these:\n");
	fprintf(stream, "      -p, --pid=N        pid of the process\n");
	fprintf(stream, "      -e, --exe=FILE     name of the executable program file\n");
        fprintf(stream, "                         The -e option only works when\n");
        fprintf(stream, "                         cpulimit is run with admin rights.\n");
	fprintf(stream, "      -P, --path=PATH    absolute path name of the\n");
        fprintf(stream, "                         executable program file\n");
	fprintf(stream, "   OPTIONS\n");
        fprintf(stream, "      -b  --background   run in background\n");
	fprintf(stream, "      -l, --limit=N      percentage of cpu allowed from 1 up.\n");
        fprintf(stream, "                         Usually 1 - 100, but can be higher\n");
        fprintf(stream, "                         on multi-core CPUs (mandatory)\n");
	fprintf(stream, "      -v, --verbose      show control statistics\n");
	fprintf(stream, "      -z, --lazy         exit if there is no suitable target process,\n");
        fprintf(stream, "                         or if it dies\n");
	fprintf(stream, "      -h, --help         display this help and exit\n");
	exit(exit_code);
}

int main(int argc, char **argv) {

	//get program name
	char *p=(char*)memrchr(argv[0],(unsigned int)'/',strlen(argv[0]));
	program_name = p==NULL?argv[0]:(p+1);
        int run_in_background = 0;
	//parse arguments
	int next_option;
	/* A string listing valid short options letters. */
	const char* short_options="p:e:P:l:bvzh";
	/* An array describing valid long options. */
	const struct option long_options[] = {
		{ "pid", required_argument, NULL, 'p' },
		{ "exe", required_argument, NULL, 'e' },
		{ "path", required_argument, NULL, 'P' },
		{ "limit", required_argument, NULL, 'l' },
                { "background", no_argument, NULL, 'b' },
		{ "verbose", no_argument, NULL, 'v' },
		{ "lazy", no_argument, NULL, 'z' },
		{ "help", no_argument, NULL, 'h' },
		{ NULL, 0, NULL, 0 }
	};
	//argument variables
	const char *exe=NULL;
	const char *path=NULL;
	int perclimit=0;
	int pid_ok=0;
	int process_ok=0;
	int limit_ok=0;
        struct rlimit maxlimit;

	do {
		next_option = getopt_long (argc, argv, short_options,long_options, NULL);
		switch(next_option) {
                        case 'b':
                                run_in_background = 1;
                                break;
			case 'p':
				pid=atoi(optarg);
				pid_ok=1;
                                lazy = 1;
				break;
			case 'e':
				exe=optarg;
				process_ok=1;
				break;
			case 'P':
				path=optarg;
				process_ok=1;
				break;
			case 'l':
				perclimit=atoi(optarg);
				limit_ok=1;
				break;
			case 'v':
				verbose=1;
				break;
			case 'z':
				lazy=1;
				break;
			case 'h':
				print_usage (stdout, 1);
				break;
			case '?':
				print_usage (stderr, 1);
				break;
			case -1:
				break;
			default:
				abort();
		}
	} while(next_option != -1);

	if (!process_ok && !pid_ok) {
		fprintf(stderr,"Error: You must specify a target process\n");
		print_usage (stderr, 1);
		exit(1);
	}
	if ((exe!=NULL && path!=NULL) || (pid_ok && (exe!=NULL || path!=NULL))) {
		fprintf(stderr,"Error: You must specify exactly one target process\n");
		print_usage (stderr, 1);
		exit(1);
	}
	if (!limit_ok) {
		fprintf(stderr,"Error: You must specify a cpu limit\n");
		print_usage (stderr, 1);
		exit(1);
	}
	float limit=perclimit/100.0;
	if (limit <= 0.00) // || limit >1) {
        {
		fprintf(stderr,"Error: limit must be greater than 0\n");
		print_usage (stderr, 1);
		exit(1);
	}

        // check to see if we should fork
        if (run_in_background)
        {
             pid_t process_id;
             process_id = fork();
             if (! process_id)
                exit(0);
             else
             {
                setsid();
                process_id = fork();
                if (process_id)
                  exit(0);
             }
        }

	//parameters are all ok!
	signal(SIGINT,quit);
	signal(SIGTERM,quit);

	if (setpriority(PRIO_PROCESS,getpid(),-20)!=0) {
	//if that failed, check if we have a limit 
        // by how much we can raise the priority
#ifdef RLIMIT_NICE 
//check if non-root can even make changes 
// (ifdef because it's only available in linux >= 2.6.13)
		nice_lim=getpriority(PRIO_PROCESS,getpid());
		getrlimit(RLIMIT_NICE, &maxlimit);

//if we can do better then current
		if( (20 - (signed)maxlimit.rlim_cur) < nice_lim &&  
		    setpriority(PRIO_PROCESS,getpid(),
                    20 - (signed)maxlimit.rlim_cur)==0 //and it actually works
		  ) {

			//if we can do better, but not by much, warn about it
			if( (nice_lim - (20 - (signed)maxlimit.rlim_cur)) < 9) 
                        {
			printf("Warning, can only increase priority by %d.\n",                                nice_lim - (20 - (signed)maxlimit.rlim_cur));
			}
                        //our new limit
			nice_lim = 20 - (signed)maxlimit.rlim_cur; 

		} else 
// otherwise don't try to change priority. 
// The below will also run if it's not possible 
// for non-root to change priority
#endif
		{
			printf("Warning: cannot renice.\nTo work better you should run this program as root, or adjust RLIMIT_NICE.\nFor example in /etc/security/limits.conf add a line with: * - nice -10\n\n");
			nice_lim=INT_MAX;
		}
	} else {
		nice_lim=-20;
	}
	//don't bother putting setpriority back down, 
        // since getpidof and waitforpid twiddle it anyway



	//time quantum in microseconds. it's splitted in a working period and a sleeping one
	int period=100000;
	struct timespec twork,tsleep;   //working and sleeping intervals
	memset(&twork,0,sizeof(struct timespec));
	memset(&tsleep,0,sizeof(struct timespec));

wait_for_process:

	//look for the target process..or wait for it
	if (exe!=NULL)
		pid=getpidof(exe);
	else if (path!=NULL)
		pid=getpidof(path);
	else {
		waitforpid(pid);
	}
	//process detected...let's play

	//init compute_cpu_usage internal stuff
	compute_cpu_usage(0,0,NULL);
	//main loop counter
	int i=0;

	struct timespec startwork,endwork;
	long workingtime=0;		//last working time in microseconds

	if (verbose) print_caption();

	float pcpu_avg=0;

	//here we should already have high priority, for time precision
	while(1) {

		//estimate how much the controlled process is using the cpu in its working interval
		struct cpu_usage cu;
		if (compute_cpu_usage(pid,workingtime,&cu)==-1) {
			fprintf(stderr,"Process %d dead!\n",pid);
			if (lazy) exit(2);
			//wait until our process appears
			goto wait_for_process;		
		}

		//cpu actual usage of process (range 0-1)
		float pcpu=cu.pcpu;
		//rate at which we are keeping active the process (range 0-1)
		float workingrate=cu.workingrate;

		//adjust work and sleep time slices
		if (pcpu>0) {
			twork.tv_nsec=min(period*limit*1000/pcpu*workingrate,period*1000);
		}
		else if (pcpu==0) {
			twork.tv_nsec=period*1000;
		}
		else if (pcpu==-1) {
			//not yet a valid idea of cpu usage
			pcpu=limit;
			workingrate=limit;
			twork.tv_nsec=min(period*limit*1000,period*1000);
		}
		tsleep.tv_nsec=period*1000-twork.tv_nsec;

		//update average usage
		pcpu_avg=(pcpu_avg*i+pcpu)/(i+1);

		if (verbose && i%10==0 && i>0) {
			printf("%0.2f%%\t%6ld us\t%6ld us\t%0.2f%%\n",pcpu*100,twork.tv_nsec/1000,tsleep.tv_nsec/1000,workingrate*100);
		}

		if (limit<1 && limit>0) {
			//resume process
			if (kill(pid,SIGCONT)!=0) {
				fprintf(stderr,"Process %d dead!\n",pid);
				if (lazy) exit(2);
				//wait until our process appears
				goto wait_for_process;
			}
		}

		clock_gettime(CLOCK_REALTIME,&startwork);
		nanosleep(&twork,NULL);		//now process is working	
		clock_gettime(CLOCK_REALTIME,&endwork);
		workingtime=timediff(&endwork,&startwork);

		if (limit<1) {
			//stop process, it has worked enough
			if (kill(pid,SIGSTOP)!=0) {
				fprintf(stderr,"Process %d dead!\n",pid);
				if (lazy) exit(2);
				//wait until our process appears
				goto wait_for_process;
			}
			nanosleep(&tsleep,NULL);	//now process is sleeping
		}
		i++;
	}

}