/*

 * block queue tracing application

 *

 * Copyright (C) 2005 Jens Axboe <axboe@suse.de>

 * Copyright (C) 2006 Jens Axboe <axboe@kernel.dk>

 *

 * Rewrite to have a single thread per CPU (managing all devices on that CPU)

 *	Alan D. Brunelle <alan.brunelle@hp.com> - January 2009

 *

 *  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

 *

 */

#include <errno.h>

#include <stdarg.h>

#include <stdio.h>

#include <stdlib.h>

#include <string.h>

#include <fcntl.h>

#include <getopt.h>

#include <sched.h>

#include <unistd.h>

#include <poll.h>

#include <signal.h>

#include <pthread.h>

#include <locale.h>

#include <sys/ioctl.h>

#include <sys/types.h>

#include <sys/stat.h>

#include <sys/vfs.h>

#include <sys/mman.h>

#include <sys/param.h>

#include <sys/time.h>

#include <sys/resource.h>

#include <sys/socket.h>

#include <netinet/in.h>

#include <arpa/inet.h>

#include <netdb.h>

#include <sys/sendfile.h>

#include "btt/list.h"

#include "blktrace.h"

/*

 * You may want to increase this even more, if you are logging at a high

 * rate and see skipped/missed events

 */

#define BUF_SIZE		(512 * 1024)

#define BUF_NR			(4)

#define FILE_VBUF_SIZE		(128 * 1024)

#define DEBUGFS_TYPE		(0x64626720)

#define TRACE_NET_PORT		(8462)

enum {

	Net_none = 0,

	Net_server,

	Net_client,

};

enum thread_status {

	Th_running,

	Th_leaving,

	Th_error

};

/*

 * Generic stats collected: nevents can be _roughly_ estimated by data_read

 * (discounting pdu...)

 *

 * These fields are updated w/ pdc_dr_update & pdc_nev_update below.

 */

struct pdc_stats {

	unsigned long long data_read;

	unsigned long long nevents;

};

struct devpath {

	struct list_head head;

	char *path;			/* path to device special file */

	char *buts_name;		/* name returned from bt kernel code */

	struct pdc_stats *stats;

	int fd, ncpus;

	unsigned long long drops;

	/*

	 * For piped output only:

	 *

	 * Each tracer will have a tracer_devpath_head that it will add new

	 * data onto. It's list is protected above (tracer_devpath_head.mutex)

	 * and it will signal the processing thread using the dp_cond,

	 * dp_mutex & dp_entries variables above.

	 */

	struct tracer_devpath_head *heads;

	/*

	 * For network server mode only:

	 */

	struct cl_host *ch;

	u32 cl_id;

	time_t cl_connect_time;

	struct io_info *ios;

};

/*

 * For piped output to stdout we will have each tracer thread (one per dev)

 * tack buffers read from the relay queues on a per-device list.

 *

 * The main thread will then collect trace buffers from each of lists in turn.

 *

 * We will use a mutex to guard each of the trace_buf list. The tracers

 * can then signal the main thread using <dp_cond,dp_mutex> and

 * dp_entries. (When dp_entries is 0, and a tracer adds an entry it will

 * signal. When dp_entries is 0, the main thread will wait for that condition

 * to be signalled.)

 *

 * adb: It may be better just to have a large buffer per tracer per dev,

 * and then use it as a ring-buffer. This would certainly cut down a lot

 * of malloc/free thrashing, at the cost of more memory movements (potentially).

 */

struct trace_buf {

	struct list_head head;

	struct devpath *dpp;

	void *buf;

	int cpu, len;

};

struct tracer_devpath_head {

	pthread_mutex_t mutex;

	struct list_head head;

	struct trace_buf *prev;

};

/*

 * Used to handle the mmap() interfaces for output file (containing traces)

 */

struct mmap_info {

	void *fs_buf;

	unsigned long long fs_size, fs_max_size, fs_off, fs_buf_len;

	unsigned long buf_size, buf_nr;

	int pagesize;

};

/*

 * Each thread doing work on a (client) side of blktrace will have one

 * of these. The ios array contains input/output information, pfds holds

 * poll() data. The volatile's provide flags to/from the main executing

 * thread.

 */

struct tracer {

	struct list_head head;

	struct io_info *ios;

	struct pollfd *pfds;

	pthread_t thread;

	int cpu, nios;

	volatile int status, is_done;

};

/*

 * networking stuff follows. we include a magic number so we know whether

 * to endianness convert or not.

 *

 * The len field is overloaded:

 *	0 - Indicates an "open" - allowing the server to set up for a dev/cpu

 *	1 - Indicates a "close" - Shut down connection orderly

 *

 * The cpu field is overloaded on close: it will contain the number of drops.

 */

struct blktrace_net_hdr {

	u32 magic;		/* same as trace magic */

	char buts_name[32];	/* trace name */

	u32 cpu;		/* for which cpu */

	u32 max_cpus;

	u32 len;		/* length of following trace data */

	u32 cl_id;		/* id for set of client per-cpu connections */

	u32 buf_size;		/* client buf_size for this trace  */

	u32 buf_nr;		/* client buf_nr for this trace  */

	u32 page_size;		/* client page_size for this trace  */

};

/*

 * Each host encountered has one of these. The head is used to link this

 * on to the network server's ch_list. Connections associated with this

 * host are linked on conn_list, and any devices traced on that host

 * are connected on the devpaths list.

 */

struct cl_host {

	struct list_head head;

	struct list_head conn_list;

	struct list_head devpaths;

	struct net_server_s *ns;

	char *hostname;

	struct in_addr cl_in_addr;

	int connects, ndevs, cl_opens;

};

/*

 * Each connection (client to server socket ('fd')) has one of these. A

 * back reference to the host ('ch'), and lists headers (for the host

 * list, and the network server conn_list) are also included.

 */

struct cl_conn {

	struct list_head ch_head, ns_head;

	struct cl_host *ch;

	int fd, ncpus;

	time_t connect_time;

};

/*

 * The network server requires some poll structures to be maintained -

 * one per conection currently on conn_list. The nchs/ch_list values

 * are for each host connected to this server. The addr field is used

 * for scratch as new connections are established.

 */

struct net_server_s {

	struct list_head conn_list;

	struct list_head ch_list;

	struct pollfd *pfds;

	int listen_fd, connects, nchs;

	struct sockaddr_in addr;

};

/*

 * This structure is (generically) used to providide information

 * for a read-to-write set of values.

 *

 * ifn & ifd represent input information

 *

 * ofn, ofd, ofp, obuf & mmap_info are used for output file (optionally).

 */

struct io_info {

	struct devpath *dpp;

	FILE *ofp;

	char *obuf;

	struct cl_conn *nc;	/* Server network connection */

	/*

	 * mmap controlled output files

	 */

	struct mmap_info mmap_info;

	/*

	 * Client network fields

	 */

	unsigned int ready;

	unsigned long long data_queued;

	/*

	 * Input/output file descriptors & names

	 */

	int ifd, ofd;

	char ifn[MAXPATHLEN + 64];

	char ofn[MAXPATHLEN + 64];

};

static char blktrace_version[] = "2.0.0";

/*

 * Linkage to blktrace helper routines (trace conversions)

 */

int data_is_native = -1;

static int ndevs;

static int max_cpus;

static int ncpus;

static cpu_set_t *online_cpus;

static int pagesize;

static int act_mask = ~0U;

static int kill_running_trace;

static int stop_watch;

static int piped_output;

static char *debugfs_path = "/sys/kernel/debug";

static char *output_name;

static char *output_dir;

static unsigned long buf_size = BUF_SIZE;

static unsigned long buf_nr = BUF_NR;

static FILE *pfp;

static LIST_HEAD(devpaths);

static LIST_HEAD(tracers);

static volatile int done;

/*

 * tracer threads add entries, the main thread takes them off and processes

 * them. These protect the dp_entries variable.

 */

static pthread_cond_t dp_cond = PTHREAD_COND_INITIALIZER;

static pthread_mutex_t dp_mutex = PTHREAD_MUTEX_INITIALIZER;

static volatile int dp_entries;

/*

 * These synchronize master / thread interactions.

 */

static pthread_cond_t mt_cond = PTHREAD_COND_INITIALIZER;

static pthread_mutex_t mt_mutex = PTHREAD_MUTEX_INITIALIZER;

static volatile int nthreads_running;

static volatile int nthreads_leaving;

static volatile int nthreads_error;

static volatile int tracers_run;

/*

 * network cmd line params

 */

static struct sockaddr_in hostname_addr;

static char hostname[MAXHOSTNAMELEN];

static int net_port = TRACE_NET_PORT;

static int net_use_sendfile = 1;

static int net_mode;

static int *cl_fds;

static int (*handle_pfds)(struct tracer *, int, int);

static int (*handle_list)(struct tracer_devpath_head *, struct list_head *);

#define S_OPTS	"d:a:A:r:o:kw:vVb:n:D:lh:p:sI:"

static struct option l_opts[] = {

	{

		.name = "dev",

		.has_arg = required_argument,

		.flag = NULL,

		.val = 'd'

	},

	{

		.name = "input-devs",

		.has_arg = required_argument,

		.flag = NULL,

		.val = 'I'

	},

	{

		.name = "act-mask",

		.has_arg = required_argument,

		.flag = NULL,

		.val = 'a'

	},

	{

		.name = "set-mask",

		.has_arg = required_argument,

		.flag = NULL,

		.val = 'A'

	},

	{

		.name = "relay",

		.has_arg = required_argument,

		.flag = NULL,

		.val = 'r'

	},

	{

		.name = "output",

		.has_arg = required_argument,

		.flag = NULL,

		.val = 'o'

	},

	{

		.name = "kill",

		.has_arg = no_argument,

		.flag = NULL,

		.val = 'k'

	},

	{

		.name = "stopwatch",

		.has_arg = required_argument,

		.flag = NULL,

		.val = 'w'

	},

	{

		.name = "version",

		.has_arg = no_argument,

		.flag = NULL,

		.val = 'v'

	},

	{

		.name = "version",

		.has_arg = no_argument,

		.flag = NULL,

		.val = 'V'

	},

	{

		.name = "buffer-size",

		.has_arg = required_argument,

		.flag = NULL,

		.val = 'b'

	},

	{

		.name = "num-sub-buffers",

		.has_arg = required_argument,

		.flag = NULL,

		.val = 'n'

	},

	{

		.name = "output-dir",

		.has_arg = required_argument,

		.flag = NULL,

		.val = 'D'

	},

	{

		.name = "listen",

		.has_arg = no_argument,

		.flag = NULL,

		.val = 'l'

	},

	{

		.name = "host",

		.has_arg = required_argument,

		.flag = NULL,

		.val = 'h'

	},

	{

		.name = "port",

		.has_arg = required_argument,

		.flag = NULL,

		.val = 'p'

	},

	{

		.name = "no-sendfile",

		.has_arg = no_argument,

		.flag = NULL,

		.val = 's'

	},

	{

		.name = NULL,

	}

};

static char usage_str[] = "\n\n" \

	"-d <dev>             | --dev=<dev>\n" \

        "[ -r <debugfs path>  | --relay=<debugfs path> ]\n" \

        "[ -o <file>          | --output=<file>]\n" \

        "[ -D <dir>           | --output-dir=<dir>\n" \

        "[ -w <time>          | --stopwatch=<time>]\n" \

        "[ -a <action field>  | --act-mask=<action field>]\n" \

        "[ -A <action mask>   | --set-mask=<action mask>]\n" \

        "[ -b <size>          | --buffer-size]\n" \

        "[ -n <number>        | --num-sub-buffers=<number>]\n" \

        "[ -l                 | --listen]\n" \

        "[ -h <hostname>      | --host=<hostname>]\n" \

        "[ -p <port number>   | --port=<port number>]\n" \

        "[ -s                 | --no-sendfile]\n" \

        "[ -I <devs file>     | --input-devs=<devs file>]\n" \

        "[ -v <version>       | --version]\n" \

        "[ -V <version>       | --version]\n" \

	"\t-d Use specified device. May also be given last after options\n" \

	"\t-r Path to mounted debugfs, defaults to /sys/kernel/debug\n" \

	"\t-o File(s) to send output to\n" \

	"\t-D Directory to prepend to output file names\n" \

	"\t-w Stop after defined time, in seconds\n" \

	"\t-a Only trace specified actions. See documentation\n" \

	"\t-A Give trace mask as a single value. See documentation\n" \

	"\t-b Sub buffer size in KiB (default 512)\n" \

	"\t-n Number of sub buffers (default 4)\n" \

	"\t-l Run in network listen mode (blktrace server)\n" \

	"\t-h Run in network client mode, connecting to the given host\n" \

	"\t-p Network port to use (default 8462)\n" \

	"\t-s Make the network client NOT use sendfile() to transfer data\n" \

	"\t-I Add devices found in <devs file>\n" \

	"\t-v Print program version info\n" \

	"\t-V Print program version info\n\n";

static void clear_events(struct pollfd *pfd)

{

	pfd->events = 0;

	pfd->revents = 0;

}

static inline int net_client_use_sendfile(void)

{

	return net_mode == Net_client && net_use_sendfile;

}

static inline int net_client_use_send(void)

{

	return net_mode == Net_client && !net_use_sendfile;

}

static inline int use_tracer_devpaths(void)

{

	return piped_output || net_client_use_send();

}

static inline int in_addr_eq(struct in_addr a, struct in_addr b)

{

	return a.s_addr == b.s_addr;

}

static inline void pdc_dr_update(struct devpath *dpp, int cpu, int data_read)

{

	dpp->stats[cpu].data_read += data_read;

}

static inline void pdc_nev_update(struct devpath *dpp, int cpu, int nevents)

{

	dpp->stats[cpu].nevents += nevents;

}

static void show_usage(char *prog)

{

	fprintf(stderr, "Usage: %s %s", prog, usage_str);

}

/*

 * Create a timespec 'msec' milliseconds into the future

 */

static inline void make_timespec(struct timespec *tsp, long delta_msec)

{

	struct timeval now;

	gettimeofday(&now, NULL);

	tsp->tv_sec = now.tv_sec;

	tsp->tv_nsec = 1000L * now.tv_usec;

	tsp->tv_nsec += (delta_msec * 1000000L);

	if (tsp->tv_nsec > 1000000000L) {

		long secs = tsp->tv_nsec / 1000000000L;

		tsp->tv_sec += secs;

		tsp->tv_nsec -= (secs * 1000000000L);

	}

}

/*

 * Add a timer to ensure wait ends

 */

static void t_pthread_cond_wait(pthread_cond_t *cond, pthread_mutex_t *mutex)

{

	struct timespec ts;

	make_timespec(&ts, 50);

	pthread_cond_timedwait(cond, mutex, &ts);

}

static void unblock_tracers(void)

{

	pthread_mutex_lock(&mt_mutex);

	tracers_run = 1;

	pthread_cond_broadcast(&mt_cond);

	pthread_mutex_unlock(&mt_mutex);

}

static void tracer_wait_unblock(struct tracer *tp)

{

	pthread_mutex_lock(&mt_mutex);

	while (!tp->is_done && !tracers_run)

		pthread_cond_wait(&mt_cond, &mt_mutex);

	pthread_mutex_unlock(&mt_mutex);

}

static void tracer_signal_ready(struct tracer *tp,

				enum thread_status th_status,

				int status)

{

	pthread_mutex_lock(&mt_mutex);

	tp->status = status;

	if (th_status == Th_running)

		nthreads_running++;

	else if (th_status == Th_error)

		nthreads_error++;

	else

		nthreads_leaving++;

	pthread_cond_signal(&mt_cond);

	pthread_mutex_unlock(&mt_mutex);

}

static void wait_tracers_ready(int ncpus_started)

{

	pthread_mutex_lock(&mt_mutex);

	while ((nthreads_running + nthreads_error) < ncpus_started)

		t_pthread_cond_wait(&mt_cond, &mt_mutex);

	pthread_mutex_unlock(&mt_mutex);

}

static void wait_tracers_leaving(void)

{

	pthread_mutex_lock(&mt_mutex);

	while (nthreads_leaving < nthreads_running)

		t_pthread_cond_wait(&mt_cond, &mt_mutex);

	pthread_mutex_unlock(&mt_mutex);

}

static void init_mmap_info(struct mmap_info *mip)

{

	mip->buf_size = buf_size;

	mip->buf_nr = buf_nr;

	mip->pagesize = pagesize;

}

static void net_close_connection(int *fd)

{

	shutdown(*fd, SHUT_RDWR);

	close(*fd);

	*fd = -1;

}

static void dpp_free(struct devpath *dpp)

{

	if (dpp->stats)

		free(dpp->stats);

	if (dpp->ios)

		free(dpp->ios);

	if (dpp->path)

		free(dpp->path);

	if (dpp->buts_name)

		free(dpp->buts_name);

	free(dpp);

}

static int lock_on_cpu(int cpu)

{

	cpu_set_t * cpu_mask;

	size_t size;

	cpu_mask = CPU_ALLOC(max_cpus);

	size = CPU_ALLOC_SIZE(max_cpus);

	CPU_ZERO_S(size, cpu_mask);

	CPU_SET_S(cpu, size, cpu_mask);

	if (sched_setaffinity(0, size, cpu_mask) < 0) {

		CPU_FREE(cpu_mask);		

		return errno;

	}

	CPU_FREE(cpu_mask);		

	return 0;

}

static int increase_limit(int resource, rlim_t increase)

{

	struct rlimit rlim;

	int save_errno = errno;

	if (!getrlimit(resource, &rlim)) {

		rlim.rlim_cur += increase;

		if (rlim.rlim_cur >= rlim.rlim_max)

			rlim.rlim_max = rlim.rlim_cur + increase;

		if (!setrlimit(resource, &rlim))

			return 1;

	}

	errno = save_errno;

	return 0;

}

static int handle_open_failure(void)

{

	if (errno == ENFILE || errno == EMFILE)

		return increase_limit(RLIMIT_NOFILE, 16);

	return 0;

}

static int handle_mem_failure(size_t length)

{

	if (errno == ENFILE)

		return handle_open_failure();

	else if (errno == ENOMEM)

		return increase_limit(RLIMIT_MEMLOCK, 2 * length);

	return 0;

}

static FILE *my_fopen(const char *path, const char *mode)

{

	FILE *fp;

	do {

		fp = fopen(path, mode);

	} while (fp == NULL && handle_open_failure());

	return fp;

}

static int my_open(const char *path, int flags)

{

	int fd;

	do {

		fd = open(path, flags);

	} while (fd < 0 && handle_open_failure());

	return fd;

}

static int my_socket(int domain, int type, int protocol)

{

	int fd;

	do {

		fd = socket(domain, type, protocol);

	} while (fd < 0 && handle_open_failure());

	return fd;

}

static int my_accept(int sockfd, struct sockaddr *addr, socklen_t *addrlen)

{

	int fd;

	do {

		fd = accept(sockfd, addr, addrlen);

	} while (fd < 0 && handle_open_failure());

	return fd;

}

static void *my_mmap(void *addr, size_t length, int prot, int flags, int fd,

		     off_t offset)

{

	void *new;

	do {

		new = mmap(addr, length, prot, flags, fd, offset);

	} while (new == MAP_FAILED && handle_mem_failure(length));

	return new;

}

static int my_mlock(struct tracer *tp,

		    const void *addr, size_t len)

{

	int ret, retry = 0;

	do {

		ret = mlock(addr, len);

		if ((retry >= 10) && tp && tp->is_done)

			break;

		retry++;

	} while (ret < 0 && handle_mem_failure(len));

	return ret;

}

static int setup_mmap(int fd, unsigned int maxlen,

		      struct mmap_info *mip,

		      struct tracer *tp)

{

	if (mip->fs_off + maxlen > mip->fs_buf_len) {

		unsigned long nr = max(16, mip->buf_nr);

		if (mip->fs_buf) {

			munlock(mip->fs_buf, mip->fs_buf_len);

			munmap(mip->fs_buf, mip->fs_buf_len);

			mip->fs_buf = NULL;

		}

		mip->fs_off = mip->fs_size & (mip->pagesize - 1);

		mip->fs_buf_len = (nr * mip->buf_size) - mip->fs_off;

		mip->fs_max_size += mip->fs_buf_len;

		if (ftruncate(fd, mip->fs_max_size) < 0) {

			perror("setup_mmap: ftruncate");

			return 1;

		}

		mip->fs_buf = my_mmap(NULL, mip->fs_buf_len, PROT_WRITE,

				      MAP_SHARED, fd,

				      mip->fs_size - mip->fs_off);

		if (mip->fs_buf == MAP_FAILED) {

			perror("setup_mmap: mmap");

			return 1;

		}

		if (my_mlock(tp, mip->fs_buf, mip->fs_buf_len) < 0) {

			perror("setup_mlock: mlock");

			return 1;

		}

	}

	return 0;

}

static int __stop_trace(int fd)

{

	/*

	 * Should be stopped, don't complain if it isn't

	 */

	ioctl(fd, BLKTRACESTOP);

	return ioctl(fd, BLKTRACETEARDOWN);

}

static int write_data(char *buf, int len)

{

	int ret;

rewrite:

	ret = fwrite(buf, len, 1, pfp);

	if (ferror(pfp) || ret != 1) {

		if (errno == EINTR) {

			clearerr(pfp);

			goto rewrite;

		}

		if (!piped_output || (errno != EPIPE && errno != EBADF)) {

			fprintf(stderr, "write(%d) failed: %d/%s\n",

				len, errno, strerror(errno));

		}

		goto err;

	}

	fflush(pfp);

	return 0;

err:

	clearerr(pfp);

	return 1;

}

/*

 * Returns the number of bytes read (successfully)

 */

static int __net_recv_data(int fd, void *buf, unsigned int len)

{

	unsigned int bytes_left = len;

	while (bytes_left && !done) {

		int ret = recv(fd, buf, bytes_left, MSG_WAITALL);

		if (ret == 0)

			break;

		else if (ret < 0) {

			if (errno == EAGAIN) {

				usleep(50);

				continue;

			}

			perror("server: net_recv_data: recv failed");

			break;

		} else {

			buf += ret;

			bytes_left -= ret;

		}

	}

	return len - bytes_left;

}

static int net_recv_data(int fd, void *buf, unsigned int len)

{

	return __net_recv_data(fd, buf, len);

}

/*

 * Returns number of bytes written

 */

static int net_send_data(int fd, void *buf, unsigned int buf_len)

{

	int ret;

	unsigned int bytes_left = buf_len;

	while (bytes_left) {

		ret = send(fd, buf, bytes_left, 0);

		if (ret < 0) {

			perror("send");

			break;

		}

		buf += ret;

		bytes_left -= ret;

	}

	return buf_len - bytes_left;

}

static int net_send_header(int fd, int cpu, char *buts_name, int len)

{

	struct blktrace_net_hdr hdr;