cputool.c

/*
 *  cputool.c - CPU & load managmenet tool
 *  Copyright (C) 2012-2013, AllWorldIT
 *  Copyright (C) 2012, Nigel Kukard <nkukard@lbsd.net>
 *
 *  This program is free software: you can redistribute it and/or modify
 *  it under the terms of the GNU General Public License as published by
 *  the Free Software Foundation, either version 3 of the License, or
 *  (at your option) any later version.
 *
 *  This program is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
*/

#include "config.h"

#include "cputool.h"

#include <dirent.h>
#include <errno.h>
#include <fcntl.h>
#include <getopt.h>
#include <signal.h>
#include <stdarg.h>
#include <stdio.h>
#include <stdint.h>
#include <stdlib.h>
#include <string.h>
#include <sys/resource.h>
#include <sys/stat.h>
#include <sys/types.h>
#include <sys/wait.h>
#include <time.h>
#include <unistd.h>



/* How accurate is the clock, per second? */
#define CLOCK_PRECISION 1000000
/* How many times do we plan to sleep per second? */
#define DEFAULT_SLEEP CLOCK_PRECISION / 10


/* Be verbose? */
int verbose = 0;

/* These two child variable are used for the signalling function */
pid_t child_pid = 0;
pid_t child_pgid = 0;
int child_external = 0;
/* Use with the SIGUSR1 and SIGUSR2 to stop/resume suspending/unsuspending */
int canStopProcesses = 1;
/* Continue loopin flag */
int loop = 1;

/* Counters */
uint64_t statsSlowdowns = 0;
uint64_t statsTimeDelayed = 0;



/* Print out our usage */
static void printUsage(char **argv) {
	fprintf(stderr,"Usage: %s [options] [command]\n",argv[0]);
	fprintf(stderr,"\n");
	fprintf(stderr,"Options:\n");
	fprintf(stderr,"    -p,  --pid=<PID>              Manage the CPU usage of a specific PID\n");
	fprintf(stderr,"    -P,  --pid-pgrp=<PID>         Manage the CPU usage of a specific PID's entire\n");
	fprintf(stderr,"                                  process group.\n");
	fprintf(stderr,"    -c,  --cpu-limit=<PCNT>       Percentage of CPU to limit process to\n");
	fprintf(stderr,"    -l,  --load-limit=<LOAD>      Load to limit process to. Decimals allowed\n");
	fprintf(stderr,"    -v,  --verbose                Be verbose in what we do, -vvv being maximum.\n");
	fprintf(stderr,"    -h,  --help                   Display this page\n\n");
	fprintf(stderr,"\n");
}


/* Log messsage */
static void logmsg(const char* format, ...)
{
	/* Grab time */
	time_t t = time(NULL);
	struct tm tm = *localtime(&t);
	va_list argptr;

	/* Print out some fancy info */
	fprintf(stderr,"%d-%02d-%02d %02d:%02d:%02d - ", tm.tm_year + 1900, tm.tm_mon + 1, tm.tm_mday, tm.tm_hour, tm.tm_min, tm.tm_sec);

	/* And the rest of the mssage */
	va_start(argptr, format);
	vfprintf(stderr, format, argptr);
	va_end(argptr);
}


/* Function to grab the load average */
static float getload() {
	float loadavgn = -1.00;
	char buf[1024];
	char *pos;
	int fd;
	ssize_t len;


	if ((fd = open("/proc/loadavg", O_RDONLY) > 0)) {
		len = read(fd, buf, sizeof(buf) - 1);
		close(fd);
		if (len > -1) {
			buf[len] = '\0';
			loadavgn = strtod(buf, &pos);
		}
	}

	return loadavgn;
}


/* Grab the stat info for a process */
static int getProcessStat(pid_t pid, struct cputool_stat *pstat)
{
	char statfile[32];
	FILE *fd;
	int i;


	/* Create filename */
	sprintf(statfile, "/proc/%d/stat", pid);

	/* Open stat file */
	if (!(fd = fopen(statfile,"r"))) {
//			logmsg("ERROR: Failed to open '%s': %s\n",statfile,strerror(errno));
			return -1;
	}

	/* Scan in stat */
	i = fscanf(fd,CPUTOOL_STAT_FORMAT,
		&pstat->pid,pstat->comm,&pstat->state,&pstat->ppid,&pstat->pgrp,
		&pstat->session,&pstat->tty_nr,&pstat->tpgid,
		&pstat->flags,
		&pstat->minflt,&pstat->cminflt,&pstat->majflt,&pstat->cmajflt,
		&pstat->utime,&pstat->stime,&pstat->cutime,&pstat->cstime,
		&pstat->priority, &pstat->nice,
		&pstat->num_threads,
		&pstat->itrealvalue,
		&pstat->starttime,
		&pstat->vsize,&pstat->rss,&pstat->rlim,
		&pstat->startcode,&pstat->endcode,&pstat->startstack,
		&pstat->kstkesp,&pstat->kstkeip,
		&pstat->signal,&pstat->blocked,
		&pstat->sigignore,&pstat->sigcatch,
		&pstat->wchan,
		&pstat->nswap,&pstat->cnswap,
		&pstat->exit_signal,
		&pstat->processor,&pstat->rt_priority,
		&pstat->policy,
		&pstat->delayacct_blkio_ticks
	);
	fclose(fd);

	/* Check result */
	if (i < 42) {
		return -1;
	}

	return 0;
}


/* Grab single processes CPU time */
static uint64_t getProcessCPUTime(pid_t pid)
{
	/* Handles & structures */
	struct cputool_stat pstat;
	/* Combined total of CPU time consumed */
	uint64_t cpuTime = 0;


	/* Grab process stat for this PID */
	if (!getProcessStat(pid,&pstat)) {
		cpuTime = pstat.utime + pstat.stime;
	}

	return cpuTime;
}


/* Grab an entire process groups CPU time */
static uint64_t getProcessGroupCPUTime(pid_t pgrp)
{
	/* Handles & structures */
	DIR *proc;
	struct dirent *entry = NULL;
	struct cputool_stat pstat;
	/* Combined total of CPU time consumed */
	uint64_t cpuTime = 0;


	/* Open /proc */
	if ((proc = opendir("/proc")) == NULL) {
		logmsg("ERROR: Failed to opendir() on '/proc': %s\n",strerror(errno));
		return -1;
	}

	/* Loop reading in directory entries */
	while ((entry = readdir(proc))) {
		if (strtok(entry->d_name,"0123456789"))
			continue;
		/* Grab process stat for this PID */
		if (!getProcessStat(atoi(entry->d_name),&pstat)) {
			/* If its in our group, then add up the CPU time */
			if (pstat.pgrp == pgrp) {
				cpuTime += pstat.utime + pstat.stime;
			}
		}
	}
	/* Close off our handle in /proc */
	closedir(proc);

	return cpuTime;
}


/* Return time difference between two timevals in ms */
static inline uint64_t timediff_us(const struct timespec *tv1, const struct timespec *tv2)
{
	/* Calculate the total time difference by adding up secs + usecs */
	return (tv1->tv_sec - tv2->tv_sec) * 1000000 + (tv1->tv_nsec - tv2->tv_nsec) / 1000;
}


/* Signal handling */
static void sigusr1 () {
	signal(SIGUSR1, sigusr1);
	canStopProcesses = 0;
}
static void sigusr2 () {
	signal(SIGUSR2, sigusr2);
	canStopProcesses = 1;
}
/* And the handler itself */
static void signal_handler(int signum) {

	/* Resume so the child can handle the signal */
	if (child_pgid) {
		killpg(child_pgid,SIGCONT);
	} else if (child_pid) {
		kill(child_pid,SIGCONT);
	}

	/* Make sure its not an external process */
	if (!child_external) {

		/* Kill it with the signal we got */
		if (child_pgid) {
			killpg(child_pgid,signum);
		} else if (child_pid) {
			kill(child_pid,signum);
		}

		/* And wait... */
		waitpid(-1, NULL, 0);
	}

	/* We should not continue looping */
	loop = 0;
}



/* Main program */
int main (int argc, char *argv[]) {

	/* If this variable is set, we spawned a child */
	int haveChild = 0;

	/* PID of process we're working on */
	pid_t pid = 0;
	/* Process GROUP */
	pid_t pgid = 0;
	/* Wait status */
	int waitStatus;

	/* Child is running */
	int isRunning = 0;
	int exceededLoad = 0;
	int exceededCPU = 0;

	/* Sleep timespec */
	struct timespec sleepTime;
	/* Time now and before */
	struct timespec now;
	struct timespec lastUpdate;

	/* CPU usage now and before, plus the limit */
	uint64_t cpuNow = 0;
	uint64_t cpuLast = 0 ;
	uint32_t cpuLimit = 0;
	/* Load limit */
	float loadLimit = 0.00;
	/* bucket holding how many ticks we can consume */
	float tickBucket = 0;
	float tickBucketMax = 0;
	/* Setup default sleep period */
	uint32_t sleep_us = DEFAULT_SLEEP;

	/* Misc */
	int i;

	/* Our long options */
	struct option long_options[] = {
		{"pid",0,0,'p'},
		{"pid-pgrp",0,0,'P'},
		{"cpu-limit",0,0,'c'},
		{"load-limit",0,0,'l'},
		{"verbose",0,0,'v'},
		{"help",0,0,'h'},
		{0,0,0,0}
	};


	/* Setup signals */
	signal(SIGUSR1, sigusr1);
	signal(SIGUSR2, sigusr2);
//	signal(SIGCHLD, sigchld);
	signal(SIGINT, signal_handler);
	signal(SIGHUP, signal_handler);
	signal(SIGQUIT, signal_handler);
	signal(SIGILL, signal_handler);
	signal(SIGKILL, signal_handler);
	signal(SIGABRT, signal_handler);
	signal(SIGTERM, signal_handler);
	signal(SIGPIPE, signal_handler);
	signal(SIGSEGV, signal_handler);

	/* Loop with options */
	while (1) {
		int option_index = 0;
		char c;

		/* Process */
		c = getopt_long(argc,argv,"p:P:c:l:vh",long_options,&option_index);

		if (c == -1) {
			break;
		}

		/* Check... */
		switch (c) {
			case 'p':
				pid = atoi(optarg);
				break;
			case 'P':
				pid = atoi(optarg);
				pgid = pid;
				break;
			case 'c':
				cpuLimit = atoi(optarg);
				/* Check the value range for cpuLimit */
				if (cpuLimit < 1) {
					fprintf(stderr,"%s: The value for -c/--cpu-limit must be in above 1\n",argv[0]);
					return 1;
				}
				break;
			case 'l':
				loadLimit = atof(optarg);
				/* Check the value range for loadLimit */
				if (loadLimit < 0.01) {
					fprintf(stderr,"%s: The value for -l/--load-limit must be above 0.00\n",argv[0]);
					return 1;
				}
				break;
			case 'v':
				verbose++;
				break;
			case 'h':
				printUsage(argv);
				return 0;
			default:
				fprintf(stderr,"Try --help for more info\n");
				return 1;
		}
	}

	/* If we don't have a PID we should have a command to run */
	if (!pid && !pgid && optind == argc) {
		fprintf(stderr,"%s: Nothing to manage. You must specify --pid/-p, --pid-pgid/-P or a command\n\n",argv[0]);
		fprintf(stderr,"Try --help for more info\n");
		return 1;
	}

	/* If we STILL have params left over, its bad */
	if (pid && !pgid && optind < argc) {
		while (optind < argc)
			fprintf(stderr,"%s: Invalid option -- %s\n",argv[0],argv[optind++]);
		fprintf(stderr,"Try --help for more info\n");
		return 1;
	}

	/* If we don't have a PID, its more than likely we must commandline it */
	if (!pid) {
		/* Loop with extra args and build our new execve environment */
		for (i = 0; i < (argc - optind); i++) {
			argv[i] = argv[i + optind];
		}
		/* End it of with NULL */
		argv[i] = NULL;

		haveChild = 1;
		pid = fork();
		if (pid < 0) {
			logmsg("ERROR: Failed to fork new process\n");
			return 1;

		/* Parent */
		} else if (pid > 1) {
			pgid = pid;

		/* Child */
		} else {
			pid = getpid();

			/* Reset parent */
			if (setsid () == -1) {
				logmsg("ERROR: Error resetting PGID\n");
			}

			if (verbose) {
//				logmsg("Child process %i PRIO set to 20\n",pid);
			}

			/* FIXME */
//			setpriority (PRIO_PROCESS, pid, 20);
			execvp(argv[0], argv);

			/* We shouldn't really get here */
			logmsg("ERROR: Failed to execute command '%s': %s\n",argv[0],strerror(errno));
			return 1;
		}

	/* Set process group if we were specified on the commandline */
	} else {
		child_external = 1;
		/* Setup the process group properly, we just set it to pid earlier */
		if (pgid) {
			pgid = getpgid(pid);
			/* Reset CPU counters */
			cpuLast = cpuNow = getProcessGroupCPUTime(pgid);
		} else {
			/* Reset CPU counters */
			cpuLast = cpuNow = getProcessCPUTime(pid);
		}
	}

	/* Setup child parent group for signalling */
	child_pid = pid;
	child_pgid = pgid;
	/* Continue process */
	killpg(pgid, SIGCONT);
	isRunning = 1;

	/* Last update is right now */
	clock_gettime(CLOCK_MONOTONIC,&lastUpdate);

	if (verbose > 1) {
		logmsg("Child PID/PGID => %lu/%lu\n",pid,pgid);
		/* Check what additional debug info we're gonig to display */
		if (cpuLimit) {
			logmsg(" CPU Limit : %u%%\n",cpuLimit);
		}
		if (loadLimit > 0.00) {
			logmsg(" LOAD Limit: %.2f\n",loadLimit);
		}
	}

	/* Set max tickBucket size & initialize tickBucket to that */
	tickBucket = tickBucketMax = (float) HZ * ( (float) cpuLimit / (float) 100);

	/* This is the main program loop */
	while (loop) {
		/* Period from last check (ms) */
		uint64_t elapsed_us;
		/* Were we running here? */
		int wasRunning = isRunning;


		/* Check if we have dead children */
		if (haveChild) {
			if (waitpid(pid, &waitStatus, WNOHANG) < 0) {
				logmsg("Dead child\n");
				break;
			}

		/* Check process is still alive */
		} else if (kill(pid,0) == -1 && errno == ESRCH) {
			logmsg("Process not alive\n");
			break;
		}

		/* We need to grab "now" so we can calculate below */
		clock_gettime(CLOCK_MONOTONIC,&now);
 		elapsed_us = timediff_us(&now,&lastUpdate);
		/* Statistics */
		if (!wasRunning) {
			statsTimeDelayed += elapsed_us;
		}

		/* Are we processing load limits? */
		if (loadLimit > 0.00) {
			double load = getload();

			/* Check if our current load is exceeding our limit, stop */
			if (load > loadLimit) {
				exceededLoad = 1;

			/* If we running and we below the threshold, resume */
			} else {
				exceededLoad = 0;
			}

			if (verbose > 1) {
					logmsg("LOAD LIMIT => load %.2f/%.2f\n",load,loadLimit);
			}
		}

		/* Are we processing cpu limits? */
		if (cpuLimit) {
			/* Change in ticks for period */
			uint32_t ticks_delta;
			/* Number of ticks allowed */
			double ticks_allowed;


			/* Grab current CPU time for entire process group */
			if (pgid) {
				cpuNow = getProcessGroupCPUTime(pgid);
			} else {
				cpuNow = getProcessCPUTime(pid);
			}

			/* Change in ticks for period */
			ticks_delta = cpuNow - cpuLast;
			/* Number of tickes we can eat */
			ticks_allowed = ( (float) elapsed_us / (float) CLOCK_PRECISION) * tickBucketMax * ( (float) cpuLimit / 100);

			/* Remove ticks we ate and add ones we allowed */
			tickBucket -= ticks_delta;
			tickBucket += ticks_allowed;

			/* Check we did not exceed 1s */
			if (tickBucket > tickBucketMax) {
				tickBucket = tickBucketMax;
			/* Check if we don't have an insane negative value either */
			} else if (tickBucket < - HZ) {
				tickBucket = - HZ;
			}

			/* If we running and our tick bucket is screwed, stop the process */
			if (tickBucket < 0) {
				exceededCPU = 1;
				sleep_us = DEFAULT_SLEEP;

			/* If we not running and we now have some ticks to consume, start the process */
			} else if (tickBucket > 0) {
				exceededCPU = 0;
				/* Set new sleep time */
				sleep_us = (tickBucket / HZ) * CLOCK_PRECISION;
			};

			/* Only use this if its a greater value */
			if (sleep_us < DEFAULT_SLEEP)
					sleep_us = DEFAULT_SLEEP;

			/* If our value is higher than the clock precision we MUST adjust it lower */
			/* or we will get a error returned */
			else if (sleep_us > CLOCK_PRECISION - DEFAULT_SLEEP) {
				sleep_us = CLOCK_PRECISION - DEFAULT_SLEEP;
			}

			/* Set last values */
			cpuLast = cpuNow;

			/* Print out info if we're running in verbose mode */
			if (verbose > 1) {
				logmsg("CPU LIMIT => tickBucket = %.2f (allowed += %.2f, consumed -= %lu), elapsed us = %lu, sleep_us = %lu\n", tickBucket, ticks_allowed, ticks_delta, elapsed_us, sleep_us);
			}
		}

		/* Set last time we were updated */
		clock_gettime(CLOCK_MONOTONIC,&lastUpdate);

		/* If load is high, override the sleep time */
		if (exceededLoad) {
			sleepTime.tv_sec = 5;
		} else {
			sleepTime.tv_sec = 0;
		}

		/* If we running and we should not be, then stop */
		if ((exceededLoad || exceededCPU) && wasRunning) {

			/* Check if we signalling the process group or process */
			if (pgid) {
				killpg(pgid,SIGSTOP);
				if (verbose > 2) {
					logmsg("KILLPG: SIGSTOP sent to process group %lu (%lu)\n",pgid,pid);
				}
			} else if (pid) {
				kill(pid,SIGSTOP);
				if (verbose > 2) {
					logmsg("KILLPG: SIGSTOP sent to process %lu\n",pid);
				}
			}

			isRunning = 0;

			statsSlowdowns++;

		/* If we not running and should be then continue */
		} else if (!(exceededLoad || exceededCPU) && !wasRunning) {

			/* Check if we signalling the process group or process */
			if (pgid) {
				killpg(pgid,SIGCONT);
				if (verbose > 2) {
					logmsg("KILLPG: SIGCONT sent to process group %lu (%lu)\n",pgid,pid);
				}
			} else if (pid) {
				kill(pid,SIGCONT);
				if (verbose > 2) {
					logmsg("KILLPG: SIGCONT sent to process %lu\n",pid);
				}
			}

			isRunning = 1;
		}

		/* Sleep here */
		sleepTime.tv_nsec = sleep_us * 1000;
		nanosleep(&sleepTime,NULL);
	}

	if (verbose) {
		logmsg("STATISTICS: Slowdowns = %lu, Total Time Delayed = %.2fs\n",statsSlowdowns,( (double) statsTimeDelayed / (double) CLOCK_PRECISION ));
	}

	return WEXITSTATUS(waitStatus);
}

/* vim: ts=4 */