/*
 *  kmonlog  by Davide Libenzi (linux kernel kmon logger)
 *  Version 0.17 - Copyright (C) 2007  Davide Libenzi
 *
 *  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 2 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, write to the Free Software
 *  Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
 *
 *  Davide Libenzi <davidel@xmailserver.org>
 *
 *  Assuming you have your kmon-patched kernel tree in $LNXPATH, to
 *  build  kmonlog  simply run:
 *
 *  gcc -I $LNXPATH/include/ -o kmonlog kmonlog.c -lrt
 *
 */

#include <sys/types.h>
#include <sys/stat.h>
#include <sys/time.h>
#include <fcntl.h>
#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <poll.h>
#include <signal.h>
#include <string.h>
#include <time.h>
#include <errno.h>
#include <sys/ioctl.h>
#include <linux/kmon.h>
#include <linux/kmon-stats.h>

#define READ_BUFSIZ	(1024 * 128)
#define MAX_PIDS	(64 * 1024)
#define KMON_HDRSIZE	11

static int stop_test;
static char *pids_cmdline[MAX_PIDS];

static char *pid_get_cmdline(unsigned int pid) {

	return pids_cmdline[pid & (MAX_PIDS - 1)];
}

static void pid_set_cmdline(unsigned int pid, char *cmdln) {

	pids_cmdline[pid & (MAX_PIDS - 1)] = cmdln;
}

static void pid_foreach(void (*fn)(void *, unsigned int, char *), void *priv) {
	unsigned int i;

	for (i = 0; i < MAX_PIDS; i++)
		if (pids_cmdline[i] != NULL)
			(*fn)(priv, i, pids_cmdline[i]);
}

static char *pid_load_cmdline(unsigned int pid) {
	static char buf[1024 * 16];
	int fd;
	ssize_t i, rdb;
	char *cmdln, *nsta, *nend;
	char fn[512];

	snprintf(fn, sizeof(fn) - 1, "/proc/%u/cmdline", pid);
	if ((fd = open(fn, O_RDONLY)) == -1)
		return strdup("");
	rdb = read(fd, buf, sizeof(buf));
	close(fd);
	if (rdb == 0) {
		/*
		 * Kernel threads do not have the command line.
		 */
		snprintf(fn, sizeof(fn) - 1, "/proc/%u/stat", pid);
		if ((fd = open(fn, O_RDONLY)) == -1)
			return strdup("");
		rdb = read(fd, buf, sizeof(buf));
		close(fd);
		buf[rdb] = 0;
		if ((nsta = strchr(buf, '(')) == NULL ||
		    (nend = strchr(buf, ')')) == NULL)
			return strdup("");
		buf[0] = '[';
		memmove(buf + 1, nsta + 1, nend - nsta - 1);
		buf[1 + nend - nsta - 1] = ']';
		buf[2 + nend - nsta - 1] = 0;
		rdb = 2 + nend - nsta - 1;
	}
	for (i = 0; i < rdb; i++)
		if (buf[i] == 0)
			buf[i] = ';';
	if ((cmdln = (char *) malloc(rdb + 1)) == NULL) {
		perror("allocating command line");
		exit(1);
	}
	memcpy(cmdln, buf, rdb);
	cmdln[rdb] = 0;

	return cmdln;
}

static void pid_check(unsigned int pid) {
	char *cmdln;

	if ((cmdln = pid_get_cmdline(pid)) == NULL) {
		cmdln = pid_load_cmdline(pid);
		pid_set_cmdline(pid, cmdln);
	}
}

static u_int16_t get_16(const unsigned char *p) {

	return ((u_int16_t) p[0]) | (((u_int16_t) p[1]) << 8);
}

static u_int32_t get_32(const unsigned char *p) {

	return ((u_int32_t) p[0]) | (((u_int32_t) p[1]) << 8) |
		(((u_int32_t) p[2]) << 16) | (((u_int32_t) p[3]) << 24);
}

static u_int64_t get_64(const unsigned char *p) {

	return ((u_int64_t) p[0]) | (((u_int64_t) p[1]) << 8) |
		(((u_int64_t) p[2]) << 16) | (((u_int64_t) p[3]) << 24) |
		(((u_int64_t) p[4]) << 32) | (((u_int64_t) p[5]) << 40) |
		(((u_int64_t) p[6]) << 48) | (((u_int64_t) p[7]) << 56);
}

static void sig_int(int sig) {

	++stop_test;
	signal(sig, sig_int);
}

static unsigned long long getnstime(void) {
	struct timespec tm;

	clock_gettime(CLOCK_REALTIME, &tm);
	return 1000000000ULL * tm.tv_sec + tm.tv_nsec;
}

static unsigned long long getmstime(void) {

	return getnstime() / 1000000;
}

static void record_parse(const unsigned char *buf, size_t count) {
	u_int16_t pid, pid2;

	/*
	 * This is a minimalistic parse, in order to build the process
	 * information (command line) table, that will be used together
	 * with the kmon log data to display meaningful results to the
	 * user. Since the  kmonlog  processing time should be kept at
	 * the minimum, we do not fully translate the logged data here.
	 * This can be done at later time w/out impacting the system
	 * during the sampling time. Process information must be grabbed
	 * immediately though, because of their volatility.
	 */
	switch (buf[0]) {
	case KMON_DT_TSKQUEUE:
		pid = get_16(buf + 1);
		pid_check(pid);
		break;
	case KMON_DT_TSKDEQUEUE:
		pid = get_16(buf + 1);
		pid_check(pid);
		break;
	case KMON_DT_TSKSWITCH:
		pid = get_16(buf + 1);
		pid2 = get_16(buf + 3);
		pid_check(pid);
		pid_check(pid2);
		break;
	case KMON_DT_TSKWAKEUP:
		pid = get_16(buf + 1);
		pid2 = get_16(buf + 3);
		if (pid != 0)
			pid_check(pid);
		pid_check(pid2);
		break;
	default:
		fprintf(stderr, "unknown record %u\n", (unsigned int) buf[0]);
	}
}

static void buffer_process(const unsigned char *buf, size_t count,
			   size_t *rem) {
	size_t i;
	u_int16_t cpu, rsize;
	u_int64_t nstime;

	for (i = 0; i < count;) {
		if (count - i < KMON_HDRSIZE)
			break;
		cpu = buf[i];
		nstime = get_64(buf + i + 1);
		rsize = get_16(buf + i + 9);
		if (count - i < KMON_HDRSIZE + rsize)
			break;
		record_parse(buf + i + KMON_HDRSIZE, rsize);
		i += rsize + KMON_HDRSIZE;
	}
	*rem = count - i;
}

static void pid_list_cmdline(void *priv, unsigned int pid, char *cmdln) {
	FILE *file = (FILE *) priv;

	fprintf(file, "%u\t%s\n", pid, cmdln);
}

static void usage(const char *prg) {

	fprintf(stderr,
		"use: %s [-d KDEV] [-b BUFSIZ] [-o OUTFIL] [-t MSTEST]\n"
		"\t[-m BMAX] [-p PIDFIL] [-h]\n", prg);
	exit(99);
}

int main(int ac, char **av) {
	int i, kfd, ofd = 1;
	ssize_t rdb;
	size_t buffsize = 1024 * 1024, size, rem;
	unsigned long long ts, maxtime = 0, maxdata = 0, totdata;
	const char *kdevfile = "/dev/kmon", *pidfile = "pids.dat", *outfile = NULL;
	void *rdbuf;
	FILE *pfile;
	struct pollfd pfd;

	while ((i = getopt(ac, av, "d:b:o:p:t:m:h")) != -1) {
		switch (i) {
		case 'd':
			kdevfile = optarg;
			break;
		case 'b':
			buffsize = atoi(optarg);
			break;
		case 'o':
			outfile = optarg;
			break;
		case 'p':
			pidfile = optarg;
			break;
		case 't':
			maxtime = atoll(optarg);
			break;
		case 'm':
			maxdata = atoll(optarg);
			break;
		case 'h':
			usage(av[0]);
		}
	}
	signal(SIGINT, sig_int);
	if ((rdbuf = malloc(READ_BUFSIZ)) == NULL) {
		perror("allocating read buffer");
		return 2;
	}
	if (outfile != NULL &&
	    (ofd = open(outfile, O_CREAT | O_WRONLY | O_TRUNC)) == -1) {
		perror(outfile);
		return 3;
	}
	pid_set_cmdline(0, strdup("[IDLE]"));
	if ((kfd = open(kdevfile, O_RDWR)) == -1) {
		perror(kdevfile);
		return 4;
	}
	if (ioctl(kfd, KMON_BUFFSIZE, &buffsize)) {
		perror("ioctl(KMON_BUFFSIZE)");
		return 5;
	}
	if (ioctl(kfd, KMON_START, NULL)) {
		perror("ioctl(KMON_START)");
		return 6;
	}
	ts = getmstime();
	for (totdata = 0, size = READ_BUFSIZ, rem = 0; !stop_test;) {
		pfd.fd = kfd;
		pfd.events = POLLIN;
		pfd.revents = 0;
		poll(&pfd, 1, -1);
		if (pfd.revents & POLLERR) {
			fprintf(stderr, "kmon device data buffer overrun\n");
			break;
		}
		if (pfd.revents & POLLIN) {
			size = READ_BUFSIZ - rem;
			rdb = read(kfd, rdbuf + rem, size);
			if (rdb == -1) {
				perror("reading kmon device");
				break;
			}
			size = rem + rdb;
			buffer_process(rdbuf, size, &rem);
			if (write(ofd, rdbuf, size - rem) != (size - rem)) {
				perror("writing log file");
				break;
			}
			totdata += size;
			if (rem)
				memmove(rdbuf, rdbuf + size - rem, rem);
			if (maxtime && (getmstime() - ts) >= maxtime)
				break;
			if (maxdata && totdata >= maxdata)
				break;
		}
	}
	if (ioctl(kfd, KMON_STOP, NULL))
		perror("ioctl(KMON_STOP)");
	close(kfd);
	if (outfile != NULL)
		close(ofd);
	if ((pfile = fopen(pidfile, "w")) == NULL) {
		perror(pidfile);
		return 7;
	}
	pid_foreach(pid_list_cmdline, pfile);
	fclose(pfile);

	return 0;
}