path: root/coreupdown.c

/* See LICENSE file for copyright and license details. */
#include <ctype.h>
#include <errno.h>
#include <fcntl.h>
#include <limits.h>
#include <signal.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>
#include <unistd.h>

static const char *argv0;
static volatile sig_atomic_t caught_sighup = 0;

#define coreup() coreupdown(ULONG_MAX)
#define coredown() coreupdown(1)

static int
coreupdown(unsigned long int count)
{
	char path[sizeof("/sys/devices/system/cpu/cpu/online") + 3 * sizeof(count)];
	unsigned long int i;
	int fd, ret = 0;

	for (i = 1; i; i++) {
		sprintf(path, "/sys/devices/system/cpu/cpu%zu/online", i);
		fd = open(path, O_WRONLY);
		if (fd < 0) {
			if (errno == ENOENT)
				break;
			fprintf(stderr, "%s: open %s O_WRONLY\n", argv0, path);
			return -1;
		}
		if (write(fd, i < count ? "1\n" : "0\n", 2) < 0) {
			fprintf(stderr, "%s: write %s\n", argv0, path);
			ret = -1;
		}
		close(fd);
	}
	return ret;
}

static int
oneshot_main(int argc, char **argv, unsigned long int count)
{
	if (argc && !strcmp(argv[0], "--")) {
		argv++;
		argc--;
	}

	if (argc > 1) {
	usage:
		fprintf(stderr, "usage: %s [count]\n", argv0);
		return 1;
	}

	if (argc) {
		char *end;
		errno = 0;
		if (!isdigit(**argv))
			goto usage;
		count = strtoul(*argv, &end, 10);
		if (!count || errno || *end)
			goto usage;
		argv++;
		argc--;
	}

	return -coreupdown(count);
}

static int
getusage(uintmax_t *total_out, uintmax_t *idle_out, size_t *cpus_out)
{
#define IDLE_FIELD 3

	static uintmax_t last_total = 0, last_idle = 0;
	uintmax_t cur_total, cur_idle;
	char buf[512];
	int fd;
	size_t off, len, field;
	ssize_t r;
	int line_incomplete = 0, beginning_of_line = 1;
	uintmax_t amount[] = {0, 0}, value;

	*cpus_out = 0;

	fd = open("/proc/stat", O_RDONLY);
	if (fd < 0) {
		fprintf(stderr, "%s: open /proc/stat O_RDONLY\n", argv0);
		return -1;
	}

	len = 0;
	for (;;) {
		r = read(fd, &buf[len], sizeof(buf) - len);
		if (!r) {
			fprintf(stderr, "%s: /proc/stat did not contain \"cpu\" line\n", argv0);
			close(fd);
			return -1;
		} else if (r < 0) {
			if (errno == EINTR)
				continue;
			fprintf(stderr, "%s: read /proc/stat: %s\n", argv0, strerror(errno));
			close(fd);
			return -1;
		}
		len += (size_t)r;

		if (line_incomplete)
			goto skip_line;

		for (off = 0; off < len;) {
			if (len - off < 4) {
				memmove(&buf[0], &buf[off], len -= off);
				break;
			} else if (!strncmp(&buf[off], "cpu", 3)) {
				if (isspace(buf[off + 3]))
					goto cpu_line_found;
				else if (isdigit(buf[off + 3]))
					*cpus_out += 1;
			}

		skip_line:
			while (off < len && buf[off] != '\n')
				off++;
			if (off == len) {
				line_incomplete = 0;
				break;
			} else {
				off++;
				line_incomplete = 1;
			}
		}
	}
cpu_line_found:

	off += 4;
	field = 0;
	value = 0;
	for (;;) {
		if (off == len) {
			off = len = 0;
			r = read(fd, buf, sizeof(buf));
			if (!r) {
				goto line_done;
			} else if (r < 0) {
				if (errno == EINTR)
					continue;
				fprintf(stderr, "%s: read /proc/stat: %s\n", argv0, strerror(errno));
				close(fd);
				return -1;
			}
			len = (size_t)r;
		}
		while (off < len && isspace(buf[off]))
			off++;
		if (off == len)
			continue;
		while (off < len) {
			if (buf[off] == '\n') {
				amount[field++ == IDLE_FIELD] += value;
				off++;
				goto line_done;
			} else if (isspace(buf[off])) {
				amount[field++ == IDLE_FIELD] += value;
				value = 0;
				off++;
				break;
			} else {
				value *= 10;
				value += buf[off] & 15;
				off++;
			}
		}
	}
line_done:

	for (;;) {
		if (off == len) {
			off = len = 0;
		read_more:
			r = read(fd, &buf[len], sizeof(buf) - len);
			if (!r) {
				break;
			} else if (r < 0) {
				if (errno == EINTR)
					continue;
				fprintf(stderr, "%s: read /proc/stat: %s\n", argv0, strerror(errno));
				close(fd);
				return -1;
			}
			len = (size_t)r;
		}
		for (; off < len; off++) {
			if (beginning_of_line) {
				if (len - off < 5) {
					memmove(&buf[0], &buf[off], len -= off);
					goto read_more;
				}
				if (!strncmp(&buf[off], "cpu", 3) && isdigit(buf[off + 3]))
					*cpus_out += 1;
				beginning_of_line = 0;
			}
			if (buf[off] == '\n')
				beginning_of_line = 1;
		}
	}

	close(fd);

	if (field < IDLE_FIELD) {
		fprintf(stderr, "%s: /proc/stat did not contain CPU idle time\n", argv0);
		close(fd);
		return -1;
	}

	cur_total = amount[0] + amount[1];
	cur_idle = amount[1];
	*total_out = cur_total - last_total;
	*idle_out = cur_idle - last_idle;
	last_total = cur_total;
	last_idle = cur_idle;
	return 0;

#undef IDLE_FIELD
}

static void
sighup_handler(int signo)
{
	(void) signo;
	caught_sighup = 1;
}

static int
daemon_main(int argc, char **argv, int orig_argc, char **orig_argv, int reexeced)
{
	uintmax_t total, idle;
	size_t cpus;
	int up_time = 0, down_time = 0, cooldown = 0, err;
	struct timespec interval, rem;
	uintmax_t coreup_threshold_cpu, coredown_threshold_cpu;
	int coreup_threshold_time, coredown_threshold_time;
	int coreup_cooldown, coredown_cooldown;
	time_t daemon_interval_sec;
	long int daemon_interval_nsec;

restart:
	coreup_threshold_cpu = COREUP_THRESHOLD_CPU;
	coreup_threshold_time = COREUP_THRESHOLD_TIME;
	coreup_cooldown = COREUP_COOLDOWN;
	coredown_threshold_cpu = COREDOWN_THRESHOLD_CPU;
	coredown_threshold_time = COREDOWN_THRESHOLD_TIME;
	coredown_cooldown = COREDOWN_COOLDOWN;
	daemon_interval_sec = DAEMON_INTERVAL_SEC;
	daemon_interval_nsec = DAEMON_INTERVAL_NSEC;

	if (reexeced)
		argv[0] = "coreupdownd";

	if (argc && !strcmp(argv[0], "--")) {
		argv++;
		argc--;
	}

	if (argc > 0) {
		fprintf(stderr, "usage: %s\n", argv0);
		return 1;
	}

	if (!reexeced) {
		/* TODO add proper daemonisation */
	}

	/* TODO configurations should be configurable */

	if (signal(SIGHUP, sighup_handler) == SIG_ERR)
		fprintf(stderr, "%s: signal SIGHUP <function>: %s\n", argv0, strerror(errno));

	interval.tv_sec = daemon_interval_sec;
	interval.tv_nsec = daemon_interval_nsec;

	if (getusage(&total, &idle, &cpus))
		return 1;

	close(STDIN_FILENO);
	close(STDOUT_FILENO);

	for (;;) {
		if (caught_sighup) {
			caught_sighup = 0;
			orig_argv[0] = "coreupdownd ";
			execv(DAEMON_PATH, orig_argv);
			fprintf(stderr, "%s: execv %s: %s\n", argv0, DAEMON_PATH, strerror(errno));
			reexeced = 1;
			argv = &orig_argv[1];
			argc = orig_argc - 1;
			goto restart;
		}

		if ((err = clock_nanosleep(CLOCK_REALTIME /* = monotonic */, 0, &interval, &rem))) {
			do {
				if (err != EINTR) {
					fprintf(stderr, "%s: clock_nanosleep CLOCK_REALTIME: %s\n",
					        argv0, strerror(errno));
					return 1;
				}
			} while ((err = clock_nanosleep(CLOCK_REALTIME, 0, &rem, &rem)));
		}
		if (getusage(&total, &idle, &cpus))
			return 1;
		if (cooldown) {
			cooldown--;
			continue;
		}

		if (cpus == 1 && 100 * (total - idle) > coreup_threshold_cpu * total) {
			down_time = 0;
			if (up_time++ >= coreup_threshold_time) {
				cooldown = coreup_cooldown;
				coreup();
				up_time = 0;
			}
		} else if (cpus > 1 && 100 * (total - idle) < coredown_threshold_cpu * total * cpus) {
			up_time = 0;
			if (down_time++ >= coredown_threshold_time) {
				cooldown = coredown_cooldown;
				coredown();
				down_time = 0;
			}
		} else {
			up_time = 0;
			down_time = 0;
		}
	}

	return 0;
}

int
main(int argc, char **argv)
{
	const char *command;

	if (argc--) {
		command = strrchr(argv0 = *argv++, '/');
		command = command ? &command[1] : argv0;
		if (!strcmp(command, "coreup"))
			return oneshot_main(argc, argv, ULONG_MAX);
		if (!strcmp(command, "coredown"))
			return oneshot_main(argc, argv, 1);
		if (!strcmp(command, "coreupdownd"))
			return daemon_main(argc, argv, argc + 1, &argv[-1], 0);
		if (!strcmp(command, "coreupdownd "))
			return daemon_main(argc, argv, argc + 1, &argv[-1], 1);
	}

	fprintf(stderr, "usage: coreup [count]\n"
	                "usage: coredown [count]\n"
	                "usage: coreupdownd\n");

	return 1;
}