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 <stdlib.h>
#include <string.h>
#include <sys/resource.h>
#include <sys/stat.h>
#include <sys/time.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 */
int 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 */
unsigned long statsSlowdowns = 0;
unsigned long 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,"    -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\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;

#ifdef __FreeBSD__
	long total;
	struct loadavg sysload;
	size_t size = sizeof(sysload);
	static int mib[] = {CTL_VM, VM_LOADAVG};


	if (sysctl(mib, 2, &sysload, &size, NULL, 0) < 0) {
		(void) fprintf(stderr, "top: sysctl failed: %s\n", strerror(errno));
		bzero(&total, sizeof(total));
	}

	loadavgn = ((double) sysload.ldavg[0]) / sysload.fscale;
#else
	char buf[1024];
	char *pos;
	int fd;
	int 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);
		}
	}

#endif
	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 an entire process groups CPU time */
static unsigned long getProcessGroupCPUTime(pid_t pgrp)
{
	/* Handles & structures */
	DIR *proc;
	struct dirent *entry = NULL;
	struct cputool_stat pstat;
	/* Combined total of CPU time consumed */
	unsigned long 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 unsigned long timediff_us(const struct timeval *tv1,const struct timeval *tv2)
{
	/* Calculate the total time difference by adding up secs + usecs */
	return (tv1->tv_sec - tv2->tv_sec) * 1000000 + (tv1->tv_usec - tv2->tv_usec);
}


/* 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) {

	/* If we have a child we can handle it here */
	if(child_pgid > 0) {
		/* Resume so the child can handle the signal */
		killpg(child_pgid, SIGCONT);

		/* Only kill if its ours */
		if (!child_external) {
			killpg(child_pgid, signum);
			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 timeval now;
	struct timeval lastUpdate;

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

	/* Misc */
	int i;

	/* Our long options */
	struct option long_options[] = {
		{"pid",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:c:l:vh",long_options,&option_index);

		if (c == -1)
			break;

		/* Check... */
		switch (c) {
			case 'p':
				pid = atoi(optarg);
				break;
			case 'c':
				cpuLimit = atoi(optarg);
				break;
			case 'l':
				loadLimit = atof(optarg);
				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 && optind == argc) {
		fprintf(stderr,"%s: Nothing to manage, either --pid/-p or a command must be given\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 && 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 {
		pgid = getpgid(pid);
	}

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

	/* Last update is right now */
	gettimeofday(&lastUpdate,NULL);

	if (verbose > 1)
			logmsg("Child PID/PGID => %lu/%lu\n",pid,pgid);

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


		/* Check if we have dead children */
		if (haveChild) {
			if (waitpid(pid, &waitStatus, WNOHANG) < 0)
				break;

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

		/* We need to grab "now" so we can calculate below */
		gettimeofday(&now,NULL);
 		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 */
			unsigned int ticks_delta;
			/* Number of ticks allowed */
			double ticks_allowed;


			/* Grab current CPU time for entire process group */
			cpuNow = getProcessGroupCPUTime(pgid);

			/* Change in ticks for period */
			ticks_delta = cpuNow - cpuLast;
			/* Number of tickes we can eat */
			ticks_allowed = ( (float) elapsed_us / (float) CLOCK_PRECISION) * HZ * ( (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 > HZ) {
				tickBucket = HZ;
			/* Check if we don't have an insane negative value eithr */
			} 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 */
		gettimeofday(&lastUpdate,NULL);

		/* 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) {
			if (verbose > 2)
				logmsg("KILLPG: SIGSTOP sent to %lu\n",pgid);

			killpg(pgid,SIGSTOP);
			isRunning = 0;

			statsSlowdowns++;

		/* If we not running and should be then continue */
		} else if (!(exceededLoad || exceededCPU) && !wasRunning) {
			if (verbose > 2)
				logmsg("KILLPG: SIGCONT sent to %lu\n",pgid);

			killpg(pgid,SIGCONT);
			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 */