PR6930: stapio: support file switching

Add file-switching option(-S size[,N]) to stapio. This option has two
arguments, 'size' and 'N', and requires -o option.

- When the size of output file exceeds specified 'size'MB, staprun switches
 output file to the next file. For this purpose, all output file has a serial
 number as a suffix only when user specifies this option.

- Using this option in bulk mode, the output file name will be
 'FILE_cpuX.SERIAL'.

- When the number of files exceeds specified N, staprun removes the oldest
 file. This argument can be omitted.

5 files changed, 195 insertions, 17 deletions

diff --git a/runtime/staprun/common.c b/runtime/staprun/common.c
index a1b70d3b..194488ef 100644
--- a/runtime/staprun/common.c
+++ b/runtime/staprun/common.c
@@ -28,6 +28,8 @@ int delete_mod;
 int load_only;
 int need_uprobes;
 int daemon_mode;
+off_t fsize_max;
+int fnum_max;
 
 /* module variables */
 char *modname = NULL;
@@ -54,6 +56,7 @@ static char *get_abspath(char *path)
 void parse_args(int argc, char **argv)
 {
 	int c;
+	char *s;
 
 	/* Initialize option variables. */
 	verbose = 0;
@@ -66,8 +69,10 @@ void parse_args(int argc, char **argv)
 	load_only = 0;
 	need_uprobes = 0;
 	daemon_mode = 0;
+	fsize_max = 0;
+	fnum_max = 0;
 
-	while ((c = getopt(argc, argv, "ALuvb:t:dc:o:x:D")) != EOF) {
+	while ((c = getopt(argc, argv, "ALuvb:t:dc:o:x:S:D")) != EOF) {
 		switch (c) {
 		case 'u':
 			need_uprobes = 1;
@@ -105,6 +110,16 @@ void parse_args(int argc, char **argv)
 		case 'D':
 			daemon_mode = 1;
 			break;
+		case 'S':
+			fsize_max = strtoul(optarg, &s, 10);
+			fsize_max <<= 20;
+			if (s[0] == ',')
+				fnum_max = (int)strtoul(&s[1], &s, 10);
+			if (s[0] != '\0') {
+				err("Invalid file size option '%s'.\n", optarg);
+				usage(argv[0]);
+			}
+			break;
 		default:
 			usage(argv[0]);
 		}
@@ -161,12 +176,16 @@ void parse_args(int argc, char **argv)
 		err("You have to specify output FILE with '-D' option.\n");
 		usage(argv[0]);
 	}
+	if (outfile_name == NULL && fsize_max != 0) {
+		err("You have to specify output FILE with '-S' option.\n");
+		usage(argv[0]);
+	}
 }
 
 void usage(char *prog)
 {
 	err("\n%s [-v]  [-c cmd ] [-x pid] [-u user] [-A|-L|-d]\n"
-                "\t[-b bufsize] [-o FILE [-D]] MODULE [module-options]\n", prog);
+                "\t[-b bufsize] [-o FILE [-D] [-S size[,N]]] MODULE [module-options]\n", prog);
 	err("-v              Increase verbosity.\n");
 	err("-c cmd          Command \'cmd\' will be run and staprun will\n");
 	err("                exit when it does.  The '_stp_target' variable\n");
@@ -184,6 +203,13 @@ void usage(char *prog)
 	err("                the user has permission to access will be deleted. Use \"*\"\n");
 	err("                (quoted) to delete all unused modules.\n");
 	err("-D              Run in background. This requires '-o' option.\n");
+	err("-S size[,N]     Switches output file to next file when the size\n");
+	err("                of file reaches the specified size. The value\n");
+	err("                should be an integer greater than 1 which is\n");
+	err("                assumed to be the maximum file size in MB.\n");
+	err("                When the number of output files reaches N, it\n");
+	err("                switches to the first output file. You can omit\n");
+	err("                the second argument.\n");
 	err("MODULE can be either a module name or a module path.  If a\n");
 	err("module name is used, it is looked for in the following\n");
 	err("directory: /lib/modules/`uname -r`/systemtap\n");
diff --git a/runtime/staprun/mainloop.c b/runtime/staprun/mainloop.c
index b0d88073..c80bbba4 100644
--- a/runtime/staprun/mainloop.c
+++ b/runtime/staprun/mainloop.c
@@ -489,18 +489,11 @@ int stp_main_loop(void)
     switch (type) {
 #ifdef STP_OLD_TRANSPORT
     case STP_REALTIME_DATA:
-      {
-        ssize_t bw = write(out_fd[0], data, nb);
-        if (bw >= 0 && bw != nb) {
-          nb = nb - bw;
-          bw = write(out_fd[0], data, nb);
-        }
-        if (bw != nb) {
-          _perr("write error (nb=%ld)", (long)nb);
-          cleanup_and_exit(0);
-        }
-        break;
+      if (write_realtime_data(data, nb)) {
+        _perr("write error (nb=%ld)", (long)nb);
+        cleanup_and_exit(0);
       }
+      break;
 #endif
     case STP_OOB_DATA:
       eprintf("%s", (char *)data);
diff --git a/runtime/staprun/relay.c b/runtime/staprun/relay.c
index 19621933..891913b0 100644
--- a/runtime/staprun/relay.c
+++ b/runtime/staprun/relay.c
@@ -44,6 +44,52 @@ static int ppoll(struct pollfd *fds, nfds_t nfds,
 }
 #endif
 
+int make_outfile_name(char *buf, int max, int fnum, int cpu)
+{
+	if (bulkmode) {
+		/* special case: for testing we sometimes want to write to /dev/null */
+		if (strcmp(outfile_name, "/dev/null") == 0) {
+			strcpy(buf, "/dev/null");
+		} else {
+			if (snprintf_chk(buf, max, "%s_cpu%d.%d",
+					 outfile_name, cpu, fnum))
+				return -1;
+		}
+	} else {
+		/* stream mode */
+		if (snprintf_chk(buf, max, "%s.%d", outfile_name, fnum))
+			return -1;
+	}
+	return 0;
+}
+
+static int open_outfile(int fnum, int cpu, int remove_file)
+{
+	char buf[PATH_MAX];
+	if (!outfile_name) {
+		_err("-S is set without -o. Please file a bug report.\n");
+		return -1;
+	}
+
+	if (remove_file) {
+		 /* remove oldest file */
+		if (make_outfile_name(buf, PATH_MAX, fnum - fnum_max, cpu) < 0)
+			return -1;
+		remove(buf); /* don't care */
+	}
+
+	if (make_outfile_name(buf, PATH_MAX, fnum, cpu) < 0)
+		return -1;
+	out_fd[cpu] = open (buf, O_CREAT|O_TRUNC|O_WRONLY, 0666);
+	if (out_fd[cpu] < 0) {
+		perr("Couldn't open output file %s", buf);
+		return -1;
+	}
+	if (set_clexec(out_fd[cpu]) < 0)
+		return -1;
+	return 0;
+}
+
 /**
  *	reader_thread - per-cpu channel buffer reader
  */
@@ -57,6 +103,9 @@ static void *reader_thread(void *data)
 	struct timespec tim = {.tv_sec=0, .tv_nsec=200000000}, *timeout = &tim;
 	sigset_t sigs;
 	struct sigaction sa;
+	off_t wsize = 0;
+	int fnum = 0;
+	int remove_file = 0;
 
 	sigemptyset(&sigs);
 	sigaddset(&sigs,SIGUSR2);
@@ -99,6 +148,19 @@ static void *reader_thread(void *data)
                         }
                 }
 		while ((rc = read(relay_fd[cpu], buf, sizeof(buf))) > 0) {
+			wsize += rc;
+			/* Switching file */
+			if (fsize_max && wsize > fsize_max) {
+				close(out_fd[cpu]);
+				fnum++;
+				if (fnum_max && fnum == fnum_max)
+					remove_file = 1;
+				if (open_outfile(fnum, cpu, remove_file) < 0) {
+					perr("Couldn't open file for cpu %d, exiting.", cpu);
+					return(NULL);
+				}
+				wsize = 0;
+			}
 			if (write(out_fd[cpu], buf, rc) != rc) {
 				perr("Couldn't write to output %d for cpu %d, exiting.", out_fd[cpu], cpu);
 				return(NULL);
@@ -163,7 +225,12 @@ int init_relayfs(void)
 		return -1;
 	}
 
-	if (bulkmode) {
+	if (fsize_max) {
+		/* switch file mode */
+		for (i = 0; i < ncpus; i++)
+			if (open_outfile(0, i, 0) < 0)
+				return -1;
+	} else if (bulkmode) {
 		for (i = 0; i < ncpus; i++) {
 			if (outfile_name) {
 				/* special case: for testing we sometimes want to write to /dev/null */
diff --git a/runtime/staprun/relay_old.c b/runtime/staprun/relay_old.c
index bd746f19..25ba93bf 100644
--- a/runtime/staprun/relay_old.c
+++ b/runtime/staprun/relay_old.c
@@ -23,6 +23,14 @@ static int bulkmode = 0;
 unsigned subbuf_size = 0;
 unsigned n_subbufs = 0;
 
+struct switchfile_ctrl_block {
+	off_t wsize;
+	int fnum;
+	int rmfile;
+};
+
+static struct switchfile_ctrl_block global_scb = {0, 0, 0};
+
 /* per-cpu buffer info */
 static struct buf_status
 {
@@ -70,6 +78,36 @@ void close_oldrelayfs(int detach)
 		close_relayfs_files(i);
 }
 
+static int open_oldoutfile(int fnum, int cpu, int remove_file)
+{
+	char buf[PATH_MAX];
+	if (outfile_name) {
+		if (remove_file) {
+			 /* remove oldest file */
+			if (make_outfile_name(buf, PATH_MAX, fnum - fnum_max,
+					      cpu) < 0)
+				return -1;
+			remove(buf); /* don't care */
+		}
+		if (make_outfile_name(buf, PATH_MAX, fnum, cpu) < 0)
+			return -1;
+	} else if (bulkmode) {
+		if (sprintf_chk(buf, "stpd_cpu%d.%d", cpu, fnum))
+			return -1;
+	} else { /* stream mode */
+		out_fd[cpu] = STDOUT_FILENO;
+		return 0;
+	}
+
+	out_fd[cpu] = open (buf, O_CREAT|O_TRUNC|O_WRONLY, 0666);
+	if (out_fd[cpu] < 0) {
+		perr("Couldn't open output file %s", buf);
+		return -1;
+	}
+	if (set_clexec(out_fd[cpu]) < 0)
+		return -1;
+	return 0;
+}
 /**
  *	open_relayfs_files - open and mmap buffer and open output file.
  *	Returns -1 on unexpected failure, 0 if file not found, 1 on success.
@@ -104,6 +142,11 @@ static int open_relayfs_files(int cpu, const char *relay_filebase, const char *p
 		return -1;
 	}
 
+	if (fsize_max) {
+		if (open_oldoutfile(0, cpu, 0) < 0)
+			goto err2;
+		goto opened;
+	}
 	if (outfile_name) {
 		/* special case: for testing we sometimes want to
 		 * write to /dev/null */
@@ -126,6 +169,7 @@ static int open_relayfs_files(int cpu, const char *relay_filebase, const char *p
 		perr("Couldn't open output file %s", tmp);
 		goto err2;
 	}
+opened:
 
 	total_bufsize = subbuf_size * n_subbufs;
 	relay_buffer[cpu] = mmap(NULL, total_bufsize, PROT_READ,
@@ -155,7 +199,8 @@ err1:
 /**
  *	process_subbufs - write ready subbufs to disk
  */
-static int process_subbufs(struct _stp_buf_info *info)
+static int process_subbufs(struct _stp_buf_info *info,
+			   struct switchfile_ctrl_block *scb)
 {
 	unsigned subbufs_ready, start_subbuf, end_subbuf, subbuf_idx, i;
 	int len, cpu = info->cpu;
@@ -173,6 +218,18 @@ static int process_subbufs(struct _stp_buf_info *info)
 		padding = *((unsigned *)subbuf_ptr);
 		subbuf_ptr += sizeof(padding);
 		len = (subbuf_size - sizeof(padding)) - padding;
+		scb->wsize += len;
+		if (fsize_max && scb->wsize > fsize_max) {
+			fclose(percpu_tmpfile[cpu]);
+			scb->fnum ++;
+			if (fnum_max && scb->fnum == fnum_max)
+				scb->rmfile = 1;
+			if (open_oldoutfile(scb->fnum, cpu, scb->rmfile) < 0) {
+				perr("Couldn't open file for cpu %d, exiting.", cpu);
+				exit(1);
+			}
+			scb->wsize = 0;
+		}
 		if (len) {
 			if (fwrite_unlocked (subbuf_ptr, len, 1, percpu_tmpfile[cpu]) != 1) {
 				_perr("Couldn't write to output file for cpu %d, exiting:", cpu);
@@ -196,6 +253,7 @@ static void *reader_thread(void *data)
 	struct _stp_consumed_info consumed_info;
 	unsigned subbufs_consumed;
 	cpu_set_t cpu_mask;
+	struct switchfile_ctrl_block scb = {0, 0, 0};
 
 	CPU_ZERO(&cpu_mask);
 	CPU_SET(cpu, &cpu_mask);
@@ -217,7 +275,7 @@ static void *reader_thread(void *data)
 		}
 
 		rc = read(proc_fd[cpu], &status[cpu].info, sizeof(struct _stp_buf_info));
-		subbufs_consumed = process_subbufs(&status[cpu].info);
+		subbufs_consumed = process_subbufs(&status[cpu].info, &scb);
 		if (subbufs_consumed) {
 			if (subbufs_consumed > status[cpu].max_backlog)
 				status[cpu].max_backlog = subbufs_consumed;
@@ -233,6 +291,33 @@ static void *reader_thread(void *data)
 }
 
 /**
+ *	write_realtime_data - write realtime data packet to disk
+ */
+int write_realtime_data(void *data, ssize_t nb)
+{
+	ssize_t bw;
+	global_scb.wsize += nb;
+	if (fsize_max && global_scb.wsize > fsize_max) {
+		close(out_fd[0]);
+		global_scb.fnum++;
+		if (fnum_max && global_scb.fnum == fnum_max)
+			global_scb.rmfile = 1;
+		if (open_oldoutfile(global_scb.fnum, 0,
+				    global_scb.rmfile) < 0) {
+			perr("Couldn't open file, exiting.");
+			return -1;
+		}
+		global_scb.wsize = 0;
+	}
+	bw = write(out_fd[0], data, nb);
+	if (bw >= 0 && bw != nb) {
+		nb = nb - bw;
+		bw = write(out_fd[0], data, nb);
+	}
+	return bw != nb;
+}
+
+/**
  *	init_relayfs - create files and threads for relayfs processing
  *
  *	Returns 0 if successful, negative otherwise
@@ -249,6 +334,9 @@ int init_oldrelayfs(void)
 		bulkmode = 1;
  
 	if (!bulkmode) {
+		if (fsize_max)
+			return open_oldoutfile(0, 0, 0);
+
 		if (outfile_name) {
 			out_fd[0] = open (outfile_name, O_CREAT|O_TRUNC|O_WRONLY, 0666);
 			if (out_fd[0] < 0 || set_clexec(out_fd[0]) < 0) {
diff --git a/runtime/staprun/staprun.h b/runtime/staprun/staprun.h
index b380cebd..6d0f9179 100644
--- a/runtime/staprun/staprun.h
+++ b/runtime/staprun/staprun.h
@@ -9,7 +9,7 @@
  *
  * Copyright (C) 2005-2008 Red Hat Inc.
  */
-
+#define _FILE_OFFSET_BITS 64
 #include <stdio.h>
 #include <stdlib.h>
 #include <ctype.h>
@@ -117,7 +117,9 @@ int init_relayfs(void);
 void close_relayfs(void);
 int init_oldrelayfs(void);
 void close_oldrelayfs(int);
+int write_realtime_data(void *data, ssize_t nb);
 void setup_signals(void);
+int make_outfile_name(char *buf, int max, int fnum, int cpu);
 /* staprun_funcs.c */
 void setup_staprun_signals(void);
 const char *moderror(int err);
@@ -158,6 +160,8 @@ extern int delete_mod;
 extern int load_only;
 extern int need_uprobes;
 extern int daemon_mode;
+extern off_t fsize_max;
+extern int fnum_max;
 
 /* getopt variables */
 extern char *optarg;