File collect_logs.c of Package mmdebug-tools

#define _GNU_SOURCE
#include <sys/types.h>
#include <sys/stat.h>
#include <sys/time.h>
#include <sys/mman.h>
#include <errno.h>
#include <fcntl.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <limits.h>
#include <unistd.h>
#include <time.h>

/*
 * A simple log collector into a file. It tries hard to guarantee
 * that the content will get into the output file even under a strong memory
 * pressure.
 *
 * Usage
 * ./read_vmstat -o output_file
 *
 * -i log_file - file to read. Mandatory argument
 * -o output_file - file to collect data into. File must not exist. Stdout if not
 *                  present
 * -t timeout - time period between two snapshots. s - seconds, ms - miliseconds
 * 	        and m - minutes suffix is allowed.
 * 	        1s by default
 * -T time_period[ms/s/h/d] - how long to collect data. The output file size will
 *                  be estimated to fit all the snapshots.
 *                  -s is ignored if this one is present
 * -s output_size[k/M/G] - size of the outoutput file. 10MB by default
 * -w - wraparound mode (aka ring buffer). Overwrites oldest entries.
 *      WARNING - this can clobber output format and will make it harder to parse.
 *      Use on own risk
 *      Doesn't have any effect when -T is used
 * -p - precision timestamps. Use us granularity for =S and E= timestamps. By default
 *      second (unix time - number of seconds and microseconds since the Epoch) is only used
 *      =S marker.
 *
 * Each snapshot is enclosed by header and footer.
 * =S timestamp[.usec]
 * [...]
 * E= [timestamp.usec]
 *
 * Please note that your ulimit has to be sufficient to allow to mlock the code+
 * all the buffers.
 *
 * This comes under GPL v2
 * Copyright: Michal Hocko <mhocko@suse.cz> 2015 
 */

/* TODOs
 * - snaphosts with better time granularity
 * - adaptive read buffer size 
 * - handle mlock limit transparently
 * - align the buffer size to records so no overlaps are possible on the overflow
 * - mark overflows
 */
#define NS_PER_MS (1000*1000)
#define NS_PER_SEC (1000*NS_PER_MS)
#define ALIGN(x, a) (((x) + (a) - 1) & ~((a) - 1))
#define PAGE_SIZE 4096
#define PAGE_ALIGN(addr) ALIGN(addr, PAGE_SIZE)
#define BUFER_SIZE (1UL<<20)

void usage(void)
{
	fprintf(stderr, "collect_logs -i collect_file [-o output_file] [-t timeout] [-T time_period] [-s buffer_size] [-w]\n");
	fprintf(stderr, "\n");
	fprintf(stderr, "\t-i collect_file - file to snapshot. Mandatory argument\n");
	fprintf(stderr, "\t-o output_file - file to store snapshots. stdout if missing\n");
	fprintf(stderr, "\t-t timeout ms/s/m/h/d - timeout between snapshots. 1s by default\n");
	fprintf(stderr, "\t-T time_period ms/s/m/h/d - time period to collect snapshots\n");
	fprintf(stderr, "\t-s buffer_size k/M/G - size of the output file. 10M by default. Ignored if -T is specified\n");
	fprintf(stderr, "\t-w - wrap around when the buffer gets full (ring buffer mode). WARNING - use with care, output data processing is PITA\n");
	fprintf(stderr, "\t-p - use precision timestamps (us granularity) for both start and end markers\n");
}

int open_file(const char *str)
{
	int fd;

	fd = open(str, O_CREAT|O_EXCL|O_RDWR, 0755);
	if (fd == -1) {
		perror("open input");
		return 1;
	}

	return fd;
}

int read_timeout(const char *str, struct timespec *timeout)
{
	char *end;
	unsigned long val;

	val = strtoul(str, &end, 10);
	if (val == ULONG_MAX) {
		perror("Invalid number");
		return 1;
	}
	switch(*end) {
		case 's':
			timeout->tv_sec = val;
			break;
		case 'm':
			/* ms vs minute*/
			if (*(end+1) == 's') {
				timeout->tv_sec = (val * NS_PER_MS) / NS_PER_SEC;
				timeout->tv_nsec = (val * NS_PER_MS) % NS_PER_SEC;
			} else {
				timeout->tv_sec = val*60;
			}
			break;
		case 'h':
			timeout->tv_sec = 3600 * val;
			break;
		case 'd':
			timeout->tv_sec = 24 * 3600 * val;
			break;
		default:
			fprintf(stderr, "Uknown number %s\n", str);
			return 1;
	}

	return 0;
}

size_t read_size(const char *str)
{
	char *end;
	size_t val = strtoul(str, &end, 10);

	switch (*end) {
		case 'K':
		case 'k':
			val <<= 10;
			break;
		case 'M':
			val <<= 20;
			break;
		case 'G':
			val <<= 30;
			break;
	}

	return val;
}

size_t read_snapshot(int in_fd, char *buff, size_t buffsize, int use_us)
{
	size_t size;
	struct timeval now;

	gettimeofday(&now, NULL);
	size = snprintf(buff, buffsize, "=S %lu", now.tv_sec);
	if (use_us)
		size += snprintf(buff + size, buffsize - size, ".%lu", now.tv_usec);
	size += snprintf(buff + size, buffsize - size, "\n");
	lseek(in_fd, 0, SEEK_SET);
	size += read(in_fd, buff + size, buffsize - size);
	size += snprintf(buff + size, buffsize - size, "E=");
	if (use_us) {
		gettimeofday(&now, NULL);
		size += snprintf(buff + size, buffsize - size, " %lu.%lu", now.tv_sec, now.tv_usec);
	}
	size += snprintf(buff + size, buffsize - size, "\n");

	return size;
}

size_t snapshot_size_est(int in_fd, int use_us)
{
	char buff[BUFER_SIZE];

	return read_snapshot(in_fd, buff, sizeof(buff), use_us);
}

size_t dump_str(char *buffer, size_t buffer_size, size_t pos, const char *in, size_t size)
{
	size_t i;
	for (i = 0; i < size; i++) {
		buffer[pos] = in[i];
		pos = (pos + 1) % buffer_size;
	}

	return pos;
}

/* buffer == NULL -> stdout */
int __collect_logs(int in_fd, const struct timespec *timeout, char *buffer, size_t buffer_size, int use_wraparound, int use_us)
{
	char buff[BUFER_SIZE]; /* dump to the file automatically */
	time_t before, after;
	size_t out_pos = 0;
	size_t in_pos = 0;
	size_t size = 0;
	int estimate = 0;

	/* lock everything in */
	if (mlockall(MCL_CURRENT) == -1) {
		perror("mlockall. Continuing anyway");
	}

	while (1) {
		size_t record_size;
		before = time(NULL);

		record_size = read_snapshot(in_fd, buff, sizeof(buff), use_us);
		size += record_size;
		if (buffer && !estimate) {
			printf("Estimated %d entries fit to the buffer\n", buffer_size/size);
			estimate = 1;
		}
		if (!use_wraparound && (size > buffer_size))
			break;

		/* Dump to stdout */
		if (!buffer) {
			printf("%s", buff);
		} else {
			size_t pos;
			/* TODO align to the size of the record to prevent from partial overwrites */
			pos = dump_str(buffer, buffer_size, out_pos, buff, record_size);
			if (pos < out_pos)
				fprintf(stderr, "%lu: Buffer wrapped\n", before);
			out_pos = pos;
		}

		after = time(NULL);

		if (after - before > 2) {
			fprintf(stderr, "%d: Snapshoting took %d!!!\n", before, after-before);
		}
		if (nanosleep(timeout, NULL) == -1)
			if (errno == EINTR)
				return 0;
		/* kick in the flushing */
		if (buffer)
			msync(buffer, buffer_size, MS_ASYNC);
	}

	return 0;
}

int collect_logs(int in_fd, int out_fd, const struct timespec *timeout, size_t buffer_size, int use_wraparound, int use_us)
{
	unsigned char *buffer = NULL;

	if (out_fd != -1) {
		size_t buffer_psize = PAGE_ALIGN(buffer_size);

		if (ftruncate(out_fd, buffer_psize) == -1) {
			perror("ftruncate");
			return 1;
		}

		if (fallocate(out_fd, 0, 0, buffer_psize) && errno != EOPNOTSUPP) {
			perror("fallocate");
			return 1;
		}

		/* commit it to the disk */
		sync();

		buffer = mmap(NULL, buffer_psize, PROT_READ | PROT_WRITE,
				MAP_SHARED | MAP_POPULATE, out_fd, 0);
		if (buffer == MAP_FAILED) {
			perror("mmap");
			return 1;
		}
	}

	return __collect_logs(in_fd, timeout, buffer, buffer_size, use_wraparound, use_us);
}

int main(int argc, char **argv)
{
	struct timespec timeout = {.tv_sec = 1};
	const char *out_file;
	struct timespec time_period;
	int use_time_period = 0;
	int use_wraparound = 0;
	int use_us = 0;
	int in_fd = -1, out_fd = -1;
	size_t buffer_size = 10UL<<20;
	size_t snapshot_size;
	int opt;

	while ((opt = getopt(argc, argv, "i:o:t:T:s:wph")) != -1) {
		switch (opt) {
			case 'i': /* input file */
				in_fd = open(optarg, O_RDONLY);
				if (in_fd == -1) {
					perror("open log_file:");
					return 1;
				}
				break;
			case 'o': /* output file */
				out_file = optarg;
				out_fd = open_file(optarg);
				if (out_fd == -1)
					return 1;
				break;
			case 't': /* timeout */
				if (read_timeout(optarg, &timeout))
					return 1;
				break;
			case 'T': /* time period */
				use_time_period = 1;
				if (read_timeout(optarg, &time_period))
					return 1;
				break;
			case 's': /* output file size */
				buffer_size = read_size(optarg);
				if (buffer_size == -1UL)
					return 1;
				break;
			case 'w': /* wraparound mode */
				use_wraparound = 1;
				break;
			case 'p': /* use precision (us) timestamps */
				use_us = 1;
				break;
			case 'h':
				usage();
				return 0;
			default:
				usage();
				return 1;
		}
	}

	if (in_fd == -1) {
		fprintf(stderr, "Input file missing\n");
		return 1;
	}

	snapshot_size = snapshot_size_est(in_fd, use_us);
	if (use_time_period) {
		size_t entries;

		entries = (time_period.tv_sec * NS_PER_SEC + time_period.tv_nsec) / (timeout.tv_sec * NS_PER_SEC + timeout.tv_nsec);
		buffer_size = snapshot_size * (entries + 1);
	} else if (out_fd == -1) {
		/* no limit on stdout */
		buffer_size = -1UL;
	}

	printf("file:%s timeout:%lu.%lus buffer_size:%llu\n",
			(out_fd == -1)? "stdout" : out_file,
			timeout.tv_sec, timeout.tv_nsec / NS_PER_MS,
			buffer_size);

	return collect_logs(in_fd, out_fd, &timeout, buffer_size, use_wraparound, use_us);
}