[PATCH] blktrace: allow quit when waiting for repeat connection
[blktrace.git] / blktrace.c
1 /*
2  * block queue tracing application
3  *
4  * Copyright (C) 2005 Jens Axboe <axboe@suse.de>
5  *
6  *  This program is free software; you can redistribute it and/or modify
7  *  it under the terms of the GNU General Public License as published by
8  *  the Free Software Foundation; either version 2 of the License, or
9  *  (at your option) any later version.
10  *
11  *  This program is distributed in the hope that it will be useful,
12  *  but WITHOUT ANY WARRANTY; without even the implied warranty of
13  *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
14  *  GNU General Public License for more details.
15  *
16  *  You should have received a copy of the GNU General Public License
17  *  along with this program; if not, write to the Free Software
18  *  Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
19  *
20  */
21 #include <pthread.h>
22 #include <sys/types.h>
23 #include <sys/stat.h>
24 #include <unistd.h>
25 #include <locale.h>
26 #include <signal.h>
27 #include <fcntl.h>
28 #include <string.h>
29 #include <sys/ioctl.h>
30 #include <sys/param.h>
31 #include <sys/statfs.h>
32 #include <sys/poll.h>
33 #include <sys/mman.h>
34 #include <sys/socket.h>
35 #include <stdio.h>
36 #include <stdlib.h>
37 #include <sched.h>
38 #include <ctype.h>
39 #include <getopt.h>
40 #include <errno.h>
41 #include <netinet/in.h>
42 #include <arpa/inet.h>
43 #include <netdb.h>
44 #include <sys/sendfile.h>
45
46 #include "blktrace.h"
47 #include "barrier.h"
48
49 static char blktrace_version[] = "0.99";
50
51 /*
52  * You may want to increase this even more, if you are logging at a high
53  * rate and see skipped/missed events
54  */
55 #define BUF_SIZE        (512 * 1024)
56 #define BUF_NR          (4)
57
58 #define OFILE_BUF       (128 * 1024)
59
60 #define RELAYFS_TYPE    0xF0B4A981
61
62 #define S_OPTS  "d:a:A:r:o:kw:Vb:n:D:lh:p:s"
63 static struct option l_opts[] = {
64         {
65                 .name = "dev",
66                 .has_arg = required_argument,
67                 .flag = NULL,
68                 .val = 'd'
69         },
70         {
71                 .name = "act-mask",
72                 .has_arg = required_argument,
73                 .flag = NULL,
74                 .val = 'a'
75         },
76         {
77                 .name = "set-mask",
78                 .has_arg = required_argument,
79                 .flag = NULL,
80                 .val = 'A'
81         },
82         {
83                 .name = "relay",
84                 .has_arg = required_argument,
85                 .flag = NULL,
86                 .val = 'r'
87         },
88         {
89                 .name = "output",
90                 .has_arg = required_argument,
91                 .flag = NULL,
92                 .val = 'o'
93         },
94         {
95                 .name = "kill",
96                 .has_arg = no_argument,
97                 .flag = NULL,
98                 .val = 'k'
99         },
100         {
101                 .name = "stopwatch",
102                 .has_arg = required_argument,
103                 .flag = NULL,
104                 .val = 'w'
105         },
106         {
107                 .name = "version",
108                 .has_arg = no_argument,
109                 .flag = NULL,
110                 .val = 'V'
111         },
112         {
113                 .name = "buffer-size",
114                 .has_arg = required_argument,
115                 .flag = NULL,
116                 .val = 'b'
117         },
118         {
119                 .name = "num-sub-buffers",
120                 .has_arg = required_argument,
121                 .flag = NULL,
122                 .val = 'n'
123         },
124         {
125                 .name = "output-dir",
126                 .has_arg = required_argument,
127                 .flag = NULL,
128                 .val = 'D'
129         },
130         {
131                 .name = "listen",
132                 .has_arg = no_argument,
133                 .flag = NULL,
134                 .val = 'l'
135         },
136         {
137                 .name = "host",
138                 .has_arg = required_argument,
139                 .flag = NULL,
140                 .val = 'h'
141         },
142         {
143                 .name = "port",
144                 .has_arg = required_argument,
145                 .flag = NULL,
146                 .val = 'p'
147         },
148         {
149                 .name = "sendfile",
150                 .has_arg = no_argument,
151                 .flag = NULL,
152                 .val = 's'
153         },
154         {
155                 .name = NULL,
156         }
157 };
158
159 struct tip_subbuf {
160         void *buf;
161         unsigned int len;
162         unsigned int max_len;
163         off_t offset;
164 };
165
166 #define FIFO_SIZE       (1024)  /* should be plenty big! */
167 #define CL_SIZE         (128)   /* cache line, any bigger? */
168
169 struct tip_subbuf_fifo {
170         int tail __attribute__((aligned(CL_SIZE)));
171         int head __attribute__((aligned(CL_SIZE)));
172         struct tip_subbuf *q[FIFO_SIZE];
173 };
174
175 struct thread_information {
176         int cpu;
177         pthread_t thread;
178
179         int fd;
180         void *fd_buf;
181         char fn[MAXPATHLEN + 64];
182
183         FILE *ofile;
184         char *ofile_buffer;
185         off_t ofile_offset;
186         int ofile_stdout;
187         int ofile_mmap;
188
189         int (*get_subbuf)(struct thread_information *, unsigned int);
190         int (*flush_subbuf)(struct thread_information *, struct tip_subbuf *);
191         int (*read_data)(struct thread_information *, void *, unsigned int);
192
193         unsigned long events_processed;
194         unsigned long long data_read;
195         struct device_information *device;
196
197         int exited;
198
199         /*
200          * piped fifo buffers
201          */
202         struct tip_subbuf_fifo fifo;
203         struct tip_subbuf *leftover_ts;
204
205         /*
206          * mmap controlled output files
207          */
208         unsigned long long fs_size;
209         unsigned long long fs_max_size;
210         unsigned long fs_off;
211         void *fs_buf;
212         unsigned long fs_buf_len;
213 };
214
215 struct device_information {
216         int fd;
217         char *path;
218         char buts_name[32];
219         volatile int trace_started;
220         unsigned long drop_count;
221         struct thread_information *threads;
222 };
223
224 static int ncpus;
225 static struct thread_information *thread_information;
226 static int ndevs;
227 static struct device_information *device_information;
228
229 /* command line option globals */
230 static char *relay_path;
231 static char *output_name;
232 static char *output_dir;
233 static int act_mask = ~0U;
234 static int kill_running_trace;
235 static unsigned long buf_size = BUF_SIZE;
236 static unsigned long buf_nr = BUF_NR;
237 static unsigned int page_size;
238
239 #define is_done()       (*(volatile int *)(&done))
240 static volatile int done;
241
242 #define is_trace_stopped()      (*(volatile int *)(&trace_stopped))
243 static volatile int trace_stopped;
244
245 #define is_stat_shown() (*(volatile int *)(&stat_shown))
246 static volatile int stat_shown;
247
248 int data_is_native = -1;
249
250 static void exit_trace(int status);
251
252 #define dip_tracing(dip)        (*(volatile int *)(&(dip)->trace_started))
253 #define dip_set_tracing(dip, v) ((dip)->trace_started = (v))
254
255 #define __for_each_dip(__d, __i, __e)   \
256         for (__i = 0, __d = device_information; __i < __e; __i++, __d++)
257
258 #define for_each_dip(__d, __i)  __for_each_dip(__d, __i, ndevs)
259 #define for_each_tip(__d, __t, __j)     \
260         for (__j = 0, __t = (__d)->threads; __j < ncpus; __j++, __t++)
261
262 /*
263  * networking stuff follows. we include a magic number so we know whether
264  * to endianness convert or not
265  */
266 struct blktrace_net_hdr {
267         u32 magic;              /* same as trace magic */
268         char buts_name[32];     /* trace name */
269         u32 cpu;                /* for which cpu */
270         u32 max_cpus;
271         u32 len;                /* length of following trace data */
272 };
273
274 #define TRACE_NET_PORT          (8462)
275
276 enum {
277         Net_none = 0,
278         Net_server,
279         Net_client,
280 };
281
282 /*
283  * network cmd line params
284  */
285 static char hostname[MAXHOSTNAMELEN];
286 static int net_port = TRACE_NET_PORT;
287 static int net_mode = 0;
288 static int net_sendfile;
289
290 static int net_in_fd = -1;
291 static int net_out_fd = -1;
292
293 static void handle_sigint(__attribute__((__unused__)) int sig)
294 {
295         done = 1;
296 }
297
298 static int get_dropped_count(const char *buts_name)
299 {
300         int fd;
301         char tmp[MAXPATHLEN + 64];
302
303         snprintf(tmp, sizeof(tmp), "%s/block/%s/dropped",
304                  relay_path, buts_name);
305
306         fd = open(tmp, O_RDONLY);
307         if (fd < 0) {
308                 /*
309                  * this may be ok, if the kernel doesn't support dropped counts
310                  */
311                 if (errno == ENOENT)
312                         return 0;
313
314                 fprintf(stderr, "Couldn't open dropped file %s\n", tmp);
315                 return -1;
316         }
317
318         if (read(fd, tmp, sizeof(tmp)) < 0) {
319                 perror(tmp);
320                 close(fd);
321                 return -1;
322         }
323
324         close(fd);
325
326         return atoi(tmp);
327 }
328
329 static int start_trace(struct device_information *dip)
330 {
331         struct blk_user_trace_setup buts;
332
333         memset(&buts, 0, sizeof(buts));
334         buts.buf_size = buf_size;
335         buts.buf_nr = buf_nr;
336         buts.act_mask = act_mask;
337
338         if (ioctl(dip->fd, BLKTRACESETUP, &buts) < 0) {
339                 perror("BLKTRACESETUP");
340                 return 1;
341         }
342
343         if (ioctl(dip->fd, BLKTRACESTART) < 0) {
344                 perror("BLKTRACESTART");
345                 return 1;
346         }
347
348         memcpy(dip->buts_name, buts.name, sizeof(dip->buts_name));
349         dip_set_tracing(dip, 1);
350         return 0;
351 }
352
353 static void stop_trace(struct device_information *dip)
354 {
355         if (dip_tracing(dip) || kill_running_trace) {
356                 dip_set_tracing(dip, 0);
357
358                 if (ioctl(dip->fd, BLKTRACESTOP) < 0)
359                         perror("BLKTRACESTOP");
360                 if (ioctl(dip->fd, BLKTRACETEARDOWN) < 0)
361                         perror("BLKTRACETEARDOWN");
362
363                 close(dip->fd);
364                 dip->fd = -1;
365         }
366 }
367
368 static void stop_all_traces(void)
369 {
370         struct device_information *dip;
371         int i;
372
373         for_each_dip(dip, i) {
374                 dip->drop_count = get_dropped_count(dip->buts_name);
375                 stop_trace(dip);
376         }
377 }
378
379 static void wait_for_data(struct thread_information *tip)
380 {
381         struct pollfd pfd = { .fd = tip->fd, .events = POLLIN };
382
383         do {
384                 poll(&pfd, 1, 100);
385                 if (pfd.revents & POLLIN)
386                         break;
387                 if (tip->ofile_stdout)
388                         break;
389         } while (!is_done());
390 }
391
392 static int read_data_file(struct thread_information *tip, void *buf,
393                           unsigned int len)
394 {
395         int ret = 0;
396
397         do {
398                 wait_for_data(tip);
399
400                 ret = read(tip->fd, buf, len);
401                 if (!ret)
402                         continue;
403                 else if (ret > 0)
404                         return ret;
405                 else {
406                         if (errno != EAGAIN) {
407                                 perror(tip->fn);
408                                 fprintf(stderr,"Thread %d failed read of %s\n",
409                                         tip->cpu, tip->fn);
410                                 break;
411                         }
412                         continue;
413                 }
414         } while (!is_done());
415
416         return ret;
417
418 }
419
420 static int read_data_net(struct thread_information *tip, void *buf,
421                          unsigned int len)
422 {
423         unsigned int bytes_left = len;
424         int ret = 0;
425
426         do {
427                 ret = recv(net_in_fd, buf, bytes_left, MSG_WAITALL);
428
429                 if (!ret)
430                         continue;
431                 else if (ret < 0) {
432                         if (errno != EAGAIN) {
433                                 perror(tip->fn);
434                                 fprintf(stderr, "server: failed read\n");
435                                 return 0;
436                         }
437                         continue;
438                 } else {
439                         buf += ret;
440                         bytes_left -= ret;
441                 }
442         } while (!is_done() && bytes_left);
443
444         return len - bytes_left;
445 }
446
447 static int read_data(struct thread_information *tip, void *buf,
448                      unsigned int len)
449 {
450         int ret = tip->read_data(tip, buf, len);
451
452         if (ret > 0)
453                 tip->data_read += ret;
454
455         return ret;
456 }
457
458 static inline struct tip_subbuf *
459 subbuf_fifo_dequeue(struct thread_information *tip)
460 {
461         const int head = tip->fifo.head;
462         const int next = (head + 1) & (FIFO_SIZE - 1);
463
464         if (head != tip->fifo.tail) {
465                 struct tip_subbuf *ts = tip->fifo.q[head];
466
467                 store_barrier();
468                 tip->fifo.head = next;
469                 return ts;
470         }
471
472         return NULL;
473 }
474
475 static inline int subbuf_fifo_queue(struct thread_information *tip,
476                                     struct tip_subbuf *ts)
477 {
478         const int tail = tip->fifo.tail;
479         const int next = (tail + 1) & (FIFO_SIZE - 1);
480
481         if (next != tip->fifo.head) {
482                 tip->fifo.q[tail] = ts;
483                 store_barrier();
484                 tip->fifo.tail = next;
485                 return 0;
486         }
487
488         fprintf(stderr, "fifo too small!\n");
489         return 1;
490 }
491
492 /*
493  * For file output, truncate and mmap the file appropriately
494  */
495 static int mmap_subbuf(struct thread_information *tip, unsigned int maxlen)
496 {
497         int ofd = fileno(tip->ofile);
498         int ret;
499
500         /*
501          * extend file, if we have to. use chunks of 16 subbuffers.
502          */
503         if (tip->fs_off + buf_size > tip->fs_buf_len) {
504                 if (tip->fs_buf) {
505                         munlock(tip->fs_buf, tip->fs_buf_len);
506                         munmap(tip->fs_buf, tip->fs_buf_len);
507                         tip->fs_buf = NULL;
508                 }
509
510                 tip->fs_off = tip->fs_size & (page_size - 1);
511                 tip->fs_buf_len = (16 * buf_size) - tip->fs_off;
512                 tip->fs_max_size += tip->fs_buf_len;
513
514                 if (ftruncate(ofd, tip->fs_max_size) < 0) {
515                         perror("ftruncate");
516                         return -1;
517                 }
518
519                 tip->fs_buf = mmap(NULL, tip->fs_buf_len, PROT_WRITE,
520                                    MAP_SHARED, ofd, tip->fs_size - tip->fs_off);
521                 if (tip->fs_buf == MAP_FAILED) {
522                         perror("mmap");
523                         return -1;
524                 }
525                 mlock(tip->fs_buf, tip->fs_buf_len);
526         }
527
528         ret = read_data(tip, tip->fs_buf + tip->fs_off, maxlen);
529         if (ret >= 0) {
530                 tip->fs_size += ret;
531                 tip->fs_off += ret;
532                 return 0;
533         }
534
535         return -1;
536 }
537
538 static int get_subbuf_sendfile(struct thread_information *tip,
539                                unsigned int maxlen)
540 {
541         struct tip_subbuf *ts = malloc(sizeof(*ts));
542         struct stat sb;
543
544         ts->buf = malloc(buf_size);
545         ts->max_len = maxlen;
546         ts->buf = NULL;
547
548         if (fstat(tip->fd, &sb) < 0) {
549                 perror("trace stat");
550                 return 1;
551         }
552
553         ts->len = sb.st_size - tip->ofile_offset;
554         ts->max_len = ts->len;
555         ts->offset = tip->ofile_offset;
556         tip->ofile_offset += ts->len;
557         return subbuf_fifo_queue(tip, ts);
558 }
559
560 /*
561  * Use the copy approach for pipes and network
562  */
563 static int get_subbuf(struct thread_information *tip, unsigned int maxlen)
564 {
565         struct tip_subbuf *ts = malloc(sizeof(*ts));
566         int ret;
567
568         ts->buf = malloc(buf_size);
569         ts->max_len = maxlen;
570
571         ret = read_data(tip, ts->buf, ts->max_len);
572         if (ret > 0) {
573                 ts->len = ret;
574                 return subbuf_fifo_queue(tip, ts);
575         }
576
577         return ret;
578 }
579
580 static void close_thread(struct thread_information *tip)
581 {
582         if (tip->fd != -1)
583                 close(tip->fd);
584         if (tip->ofile)
585                 fclose(tip->ofile);
586         if (tip->ofile_buffer)
587                 free(tip->ofile_buffer);
588         if (tip->fd_buf)
589                 free(tip->fd_buf);
590
591         tip->fd = -1;
592         tip->ofile = NULL;
593         tip->ofile_buffer = NULL;
594         tip->fd_buf = NULL;
595 }
596
597 static void tip_ftrunc_final(struct thread_information *tip)
598 {
599         /*
600          * truncate to right size and cleanup mmap
601          */
602         if (tip->ofile_mmap) {
603                 int ofd = fileno(tip->ofile);
604
605                 if (tip->fs_buf)
606                         munmap(tip->fs_buf, tip->fs_buf_len);
607
608                 ftruncate(ofd, tip->fs_size);
609         }
610 }
611
612 static void *thread_main(void *arg)
613 {
614         struct thread_information *tip = arg;
615         pid_t pid = getpid();
616         cpu_set_t cpu_mask;
617
618         CPU_ZERO(&cpu_mask);
619         CPU_SET((tip->cpu), &cpu_mask);
620
621         if (sched_setaffinity(pid, sizeof(cpu_mask), &cpu_mask) == -1) {
622                 perror("sched_setaffinity");
623                 exit_trace(1);
624         }
625
626         snprintf(tip->fn, sizeof(tip->fn), "%s/block/%s/trace%d",
627                         relay_path, tip->device->buts_name, tip->cpu);
628         tip->fd = open(tip->fn, O_RDONLY);
629         if (tip->fd < 0) {
630                 perror(tip->fn);
631                 fprintf(stderr,"Thread %d failed open of %s\n", tip->cpu,
632                         tip->fn);
633                 exit_trace(1);
634         }
635
636         while (!is_done()) {
637                 if (tip->get_subbuf(tip, buf_size))
638                         break;
639         }
640
641         tip_ftrunc_final(tip);
642         tip->exited = 1;
643         return NULL;
644 }
645
646 static int write_data_net(int fd, void *buf, unsigned int buf_len)
647 {
648         unsigned int bytes_left = buf_len;
649         int ret;
650
651         while (bytes_left) {
652                 ret = send(fd, buf, bytes_left, 0);
653                 if (ret < 0) {
654                         perror("send");
655                         return 1;
656                 }
657
658                 buf += ret;
659                 bytes_left -= ret;
660         }
661
662         return 0;
663 }
664
665 static int net_send_header(struct thread_information *tip, unsigned int len)
666 {
667         struct blktrace_net_hdr hdr;
668
669         hdr.magic = BLK_IO_TRACE_MAGIC;
670         strcpy(hdr.buts_name, tip->device->buts_name);
671         hdr.cpu = tip->cpu;
672         hdr.max_cpus = ncpus;
673         hdr.len = len;
674
675         return write_data_net(net_out_fd, &hdr, sizeof(hdr));
676 }
677
678 /*
679  * send header with 0 length to signal end-of-run
680  */
681 static void net_client_send_close(void)
682 {
683         struct blktrace_net_hdr hdr;
684
685         hdr.magic = BLK_IO_TRACE_MAGIC;
686         hdr.cpu = 0;
687         hdr.max_cpus = ncpus;
688         hdr.len = 0;
689
690         write_data_net(net_out_fd, &hdr, sizeof(hdr));
691 }
692
693 static int flush_subbuf_net(struct thread_information *tip,
694                             struct tip_subbuf *ts)
695 {
696         if (net_send_header(tip, ts->len))
697                 return 1;
698         if (write_data_net(net_out_fd, ts->buf, ts->len))
699                 return 1;
700
701         free(ts->buf);
702         free(ts);
703         return 0;
704 }
705
706 static int flush_subbuf_sendfile(struct thread_information *tip,
707                                  struct tip_subbuf *ts)
708 {
709         if (net_send_header(tip, ts->len))
710                 return 1;
711         if (sendfile(net_out_fd, tip->fd, &ts->offset, ts->len) < 0) {
712                 perror("sendfile");
713                 return 1;
714         }
715
716         free(ts);
717         return 0;
718 }
719
720 static int write_data(struct thread_information *tip, void *buf,
721                       unsigned int buf_len)
722 {
723         int ret;
724
725         if (!buf_len)
726                 return 0;
727
728         while (1) {
729                 ret = fwrite(buf, buf_len, 1, tip->ofile);
730                 if (ret == 1)
731                         break;
732
733                 if (ret < 0) {
734                         perror("write");
735                         return 1;
736                 }
737         }
738
739         if (tip->ofile_stdout)
740                 fflush(tip->ofile);
741
742         return 0;
743 }
744
745 static int flush_subbuf_file(struct thread_information *tip,
746                              struct tip_subbuf *ts)
747 {
748         unsigned int offset = 0;
749         struct blk_io_trace *t;
750         int pdu_len, events = 0;
751
752         /*
753          * surplus from last run
754          */
755         if (tip->leftover_ts) {
756                 struct tip_subbuf *prev_ts = tip->leftover_ts;
757
758                 if (prev_ts->len + ts->len > prev_ts->max_len) {
759                         prev_ts->max_len += ts->len;
760                         prev_ts->buf = realloc(prev_ts->buf, prev_ts->max_len);
761                 }
762
763                 memcpy(prev_ts->buf + prev_ts->len, ts->buf, ts->len);
764                 prev_ts->len += ts->len;
765
766                 free(ts->buf);
767                 free(ts);
768
769                 ts = prev_ts;
770                 tip->leftover_ts = NULL;
771         }
772
773         while (offset + sizeof(*t) <= ts->len) {
774                 t = ts->buf + offset;
775
776                 if (verify_trace(t)) {
777                         write_data(tip, ts->buf, offset);
778                         return -1;
779                 }
780
781                 pdu_len = t->pdu_len;
782
783                 if (offset + sizeof(*t) + pdu_len > ts->len)
784                         break;
785
786                 offset += sizeof(*t) + pdu_len;
787                 tip->events_processed++;
788                 tip->data_read += sizeof(*t) + pdu_len;
789                 events++;
790         }
791
792         if (write_data(tip, ts->buf, offset))
793                 return -1;
794
795         /*
796          * leftover bytes, save them for next time
797          */
798         if (offset != ts->len) {
799                 tip->leftover_ts = ts;
800                 ts->len -= offset;
801                 memmove(ts->buf, ts->buf + offset, ts->len);
802         } else {
803                 free(ts->buf);
804                 free(ts);
805         }
806
807         return events;
808 }
809
810 static int write_tip_events(struct thread_information *tip)
811 {
812         struct tip_subbuf *ts = subbuf_fifo_dequeue(tip);
813
814         if (ts)
815                 return tip->flush_subbuf(tip, ts);
816
817         return 0;
818 }
819
820 /*
821  * scans the tips we know and writes out the subbuffers we accumulate
822  */
823 static void get_and_write_events(void)
824 {
825         struct device_information *dip;
826         struct thread_information *tip;
827         int i, j, events, ret, tips_running;
828
829         while (!is_done()) {
830                 events = 0;
831
832                 for_each_dip(dip, i) {
833                         for_each_tip(dip, tip, j) {
834                                 ret = write_tip_events(tip);
835                                 if (ret > 0)
836                                         events += ret;
837                         }
838                 }
839
840                 if (!events)
841                         usleep(10);
842         }
843
844         /*
845          * reap stored events
846          */
847         do {
848                 events = 0;
849                 tips_running = 0;
850                 for_each_dip(dip, i) {
851                         for_each_tip(dip, tip, j) {
852                                 ret = write_tip_events(tip);
853                                 if (ret > 0)
854                                         events += ret;
855                                 tips_running += !tip->exited;
856                         }
857                 }
858                 usleep(10);
859         } while (events || tips_running);
860 }
861
862 static void wait_for_threads(void)
863 {
864         /*
865          * for piped or network output, poll and fetch data for writeout.
866          * for files, we just wait around for trace threads to exit
867          */
868         if ((output_name && !strcmp(output_name, "-")) ||
869             net_mode == Net_client)
870                 get_and_write_events();
871         else {
872                 struct device_information *dip;
873                 struct thread_information *tip;
874                 int i, j, tips_running;
875
876                 do {
877                         tips_running = 0;
878                         usleep(1000);
879
880                         for_each_dip(dip, i)
881                                 for_each_tip(dip, tip, j)
882                                         tips_running += !tip->exited;
883                 } while (tips_running);
884         }
885
886         if (net_mode == Net_client)
887                 net_client_send_close();
888 }
889
890 static void fill_ofname(char *dst, char *buts_name, int cpu)
891 {
892         int len = 0;
893
894         if (output_dir)
895                 len = sprintf(dst, "%s/", output_dir);
896
897         if (output_name)
898                 sprintf(dst + len, "%s.blktrace.%d", output_name, cpu);
899         else
900                 sprintf(dst + len, "%s.blktrace.%d", buts_name, cpu);
901 }
902
903 static void fill_ops(struct thread_information *tip)
904 {
905         /*
906          * setup ops
907          */
908         if (net_mode == Net_client) {
909                 if (net_sendfile) {
910                         tip->get_subbuf = get_subbuf_sendfile;
911                         tip->flush_subbuf = flush_subbuf_sendfile;
912                 } else {
913                         tip->get_subbuf = get_subbuf;
914                         tip->flush_subbuf = flush_subbuf_net;
915                 }
916         } else {
917                 if (tip->ofile_mmap)
918                         tip->get_subbuf = mmap_subbuf;
919                 else
920                         tip->get_subbuf = get_subbuf;
921
922                 tip->flush_subbuf = flush_subbuf_file;
923         }
924                         
925         if (net_mode == Net_server)
926                 tip->read_data = read_data_net;
927         else
928                 tip->read_data = read_data_file;
929 }
930
931 static int start_threads(struct device_information *dip)
932 {
933         struct thread_information *tip;
934         int j, pipeline = output_name && !strcmp(output_name, "-");
935         int mode, vbuf_size;
936         char op[64];
937
938         for_each_tip(dip, tip, j) {
939                 tip->cpu = j;
940                 tip->device = dip;
941                 tip->events_processed = 0;
942                 memset(&tip->fifo, 0, sizeof(tip->fifo));
943                 tip->leftover_ts = NULL;
944
945                 if (pipeline) {
946                         tip->ofile = fdopen(STDOUT_FILENO, "w");
947                         tip->ofile_stdout = 1;
948                         tip->ofile_mmap = 0;
949                         mode = _IOLBF;
950                         vbuf_size = 512;
951                 } else {
952                         fill_ofname(op, dip->buts_name, tip->cpu);
953                         tip->ofile = fopen(op, "w+");
954                         tip->ofile_stdout = 0;
955                         tip->ofile_mmap = 1;
956                         mode = _IOFBF;
957                         vbuf_size = OFILE_BUF;
958                 }
959
960                 if (tip->ofile == NULL) {
961                         perror(op);
962                         return 1;
963                 }
964
965                 tip->ofile_buffer = malloc(vbuf_size);
966                 if (setvbuf(tip->ofile, tip->ofile_buffer, mode, vbuf_size)) {
967                         perror("setvbuf");
968                         close_thread(tip);
969                         return 1;
970                 }
971
972                 fill_ops(tip);
973
974                 if (pthread_create(&tip->thread, NULL, thread_main, tip)) {
975                         perror("pthread_create");
976                         close_thread(tip);
977                         return 1;
978                 }
979         }
980
981         return 0;
982 }
983
984 static void stop_threads(struct device_information *dip)
985 {
986         struct thread_information *tip;
987         unsigned long ret;
988         int i;
989
990         for_each_tip(dip, tip, i) {
991                 (void) pthread_join(tip->thread, (void *) &ret);
992                 close_thread(tip);
993         }
994 }
995
996 static void stop_all_threads(void)
997 {
998         struct device_information *dip;
999         int i;
1000
1001         for_each_dip(dip, i)
1002                 stop_threads(dip);
1003 }
1004
1005 static void stop_all_tracing(void)
1006 {
1007         struct device_information *dip;
1008         int i;
1009
1010         for_each_dip(dip, i)
1011                 stop_trace(dip);
1012 }
1013
1014 static void exit_trace(int status)
1015 {
1016         if (!is_trace_stopped()) {
1017                 trace_stopped = 1;
1018                 stop_all_threads();
1019                 stop_all_tracing();
1020         }
1021
1022         exit(status);
1023 }
1024
1025 static int resize_devices(char *path)
1026 {
1027         int size = (ndevs + 1) * sizeof(struct device_information);
1028
1029         device_information = realloc(device_information, size);
1030         if (!device_information) {
1031                 fprintf(stderr, "Out of memory, device %s (%d)\n", path, size);
1032                 return 1;
1033         }
1034         device_information[ndevs].path = path;
1035         ndevs++;
1036         return 0;
1037 }
1038
1039 static int open_devices(void)
1040 {
1041         struct device_information *dip;
1042         int i;
1043
1044         for_each_dip(dip, i) {
1045                 dip->fd = open(dip->path, O_RDONLY | O_NONBLOCK);
1046                 if (dip->fd < 0) {
1047                         perror(dip->path);
1048                         return 1;
1049                 }
1050         }
1051
1052         return 0;
1053 }
1054
1055 static int start_devices(void)
1056 {
1057         struct device_information *dip;
1058         int i, j, size;
1059
1060         size = ncpus * sizeof(struct thread_information);
1061         thread_information = malloc(size * ndevs);
1062         if (!thread_information) {
1063                 fprintf(stderr, "Out of memory, threads (%d)\n", size * ndevs);
1064                 return 1;
1065         }
1066
1067         for_each_dip(dip, i) {
1068                 if (start_trace(dip)) {
1069                         close(dip->fd);
1070                         fprintf(stderr, "Failed to start trace on %s\n",
1071                                 dip->path);
1072                         break;
1073                 }
1074         }
1075
1076         if (i != ndevs) {
1077                 __for_each_dip(dip, j, i)
1078                         stop_trace(dip);
1079
1080                 return 1;
1081         }
1082
1083         for_each_dip(dip, i) {
1084                 dip->threads = thread_information + (i * ncpus);
1085                 if (start_threads(dip)) {
1086                         fprintf(stderr, "Failed to start worker threads\n");
1087                         break;
1088                 }
1089         }
1090
1091         if (i != ndevs) {
1092                 __for_each_dip(dip, j, i)
1093                         stop_threads(dip);
1094                 for_each_dip(dip, i)
1095                         stop_trace(dip);
1096
1097                 return 1;
1098         }
1099
1100         return 0;
1101 }
1102
1103 static void show_stats(void)
1104 {
1105         struct device_information *dip;
1106         struct thread_information *tip;
1107         unsigned long long events_processed, data_read;
1108         unsigned long total_drops;
1109         int i, j, no_stdout = 0;
1110
1111         if (is_stat_shown())
1112                 return;
1113
1114         if (output_name && !strcmp(output_name, "-"))
1115                 no_stdout = 1;
1116
1117         stat_shown = 1;
1118
1119         total_drops = 0;
1120         for_each_dip(dip, i) {
1121                 if (!no_stdout)
1122                         printf("Device: %s\n", dip->path);
1123                 events_processed = 0;
1124                 data_read = 0;
1125                 for_each_tip(dip, tip, j) {
1126                         if (!no_stdout)
1127                                 printf("  CPU%3d: %20lu events, %8llu KiB data\n",
1128                                         tip->cpu, tip->events_processed,
1129                                         tip->data_read >> 10);
1130                         events_processed += tip->events_processed;
1131                         data_read += tip->data_read;
1132                 }
1133                 total_drops += dip->drop_count;
1134                 if (!no_stdout)
1135                         printf("  Total:  %20llu events (dropped %lu), %8llu KiB data\n",
1136                                         events_processed, dip->drop_count,
1137                                         data_read >> 10);
1138         }
1139
1140         if (total_drops)
1141                 fprintf(stderr, "You have dropped events, consider using a larger buffer size (-b)\n");
1142 }
1143
1144 static struct device_information *net_get_dip(char *buts_name)
1145 {
1146         struct device_information *dip;
1147         int i;
1148
1149         for (i = 0; i < ndevs; i++) {
1150                 dip = &device_information[i];
1151
1152                 if (!strcmp(dip->buts_name, buts_name))
1153                         return dip;
1154         }
1155
1156         device_information = realloc(device_information, (ndevs + 1) * sizeof(*dip));
1157         dip = &device_information[ndevs];
1158         strcpy(dip->buts_name, buts_name);
1159         strcpy(dip->path, buts_name);
1160         ndevs++;
1161         dip->threads = malloc(ncpus * sizeof(struct thread_information));
1162         memset(dip->threads, 0, ncpus * sizeof(struct thread_information));
1163
1164         /*
1165          * open all files
1166          */
1167         for (i = 0; i < ncpus; i++) {
1168                 struct thread_information *tip = &dip->threads[i];
1169                 char op[64];
1170
1171                 tip->cpu = i;
1172                 tip->ofile_stdout = 0;
1173                 tip->ofile_mmap = 1;
1174                 tip->device = dip;
1175
1176                 fill_ops(tip);
1177
1178                 fill_ofname(op, dip->buts_name, tip->cpu);
1179
1180                 tip->ofile = fopen(op, "w+");
1181                 if (!tip->ofile) {
1182                         perror("fopen");
1183                         return NULL;
1184                 }
1185         }
1186
1187         return dip;
1188 }
1189
1190 static struct thread_information *net_get_tip(struct blktrace_net_hdr *bnh)
1191 {
1192         struct device_information *dip;
1193
1194         ncpus = bnh->max_cpus;
1195         dip = net_get_dip(bnh->buts_name);
1196         return &dip->threads[bnh->cpu];
1197 }
1198
1199 static int net_get_header(struct blktrace_net_hdr *bnh)
1200 {
1201         int fl = fcntl(net_in_fd, F_GETFL);
1202         int bytes_left, ret;
1203         void *p = bnh;
1204
1205         fcntl(net_in_fd, F_SETFL, fl | O_NONBLOCK);
1206         bytes_left = sizeof(*bnh);
1207         while (bytes_left && !is_done()) {
1208                 ret = recv(net_in_fd, p, bytes_left, MSG_WAITALL);
1209                 if (ret < 0) {
1210                         if (errno != EAGAIN) {
1211                                 perror("recv header");
1212                                 return 1;
1213                         }
1214                         usleep(100);
1215                         continue;
1216                 } else if (!ret) {
1217                         usleep(100);
1218                         continue;
1219                 } else {
1220                         p += ret;
1221                         bytes_left -= ret;
1222                 }
1223         }
1224         fcntl(net_in_fd, F_SETFL, fl & ~O_NONBLOCK);
1225         return 0;
1226 }
1227
1228 static int net_server_loop(void)
1229 {
1230         struct thread_information *tip;
1231         struct blktrace_net_hdr bnh;
1232
1233         if (net_get_header(&bnh))
1234                 return 1;
1235
1236         if (data_is_native == -1 && check_data_endianness(bnh.magic)) {
1237                 fprintf(stderr, "server: received data is bad\n");
1238                 return 1;
1239         }
1240
1241         if (!data_is_native) {
1242                 bnh.cpu = be32_to_cpu(bnh.cpu);
1243                 bnh.len = be32_to_cpu(bnh.len);
1244         }
1245
1246         /*
1247          * len == 0 means that the other end signalled end-of-run
1248          */
1249         if (!bnh.len) {
1250                 fprintf(stderr, "server: end of run\n");
1251                 return 1;
1252         }
1253
1254         tip = net_get_tip(&bnh);
1255         if (!tip)
1256                 return 1;
1257
1258         if (mmap_subbuf(tip, bnh.len))
1259                 return 1;
1260
1261         return 0;
1262 }
1263
1264 /*
1265  * Start here when we are in server mode - just fetch data from the network
1266  * and dump to files
1267  */
1268 static int net_server(void)
1269 {
1270         struct device_information *dip;
1271         struct thread_information *tip;
1272         struct sockaddr_in addr;
1273         socklen_t socklen;
1274         int fd, opt, i, j;
1275
1276         fd = socket(AF_INET, SOCK_STREAM, 0);
1277         if (fd < 0) {
1278                 perror("server: socket");
1279                 return 1;
1280         }
1281
1282         opt = 1;
1283         if (setsockopt(fd, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt)) < 0) {
1284                 perror("setsockopt");
1285                 return 1;
1286         }
1287
1288         memset(&addr, 0, sizeof(addr));
1289         addr.sin_family = AF_INET;
1290         addr.sin_addr.s_addr = htonl(INADDR_ANY);
1291         addr.sin_port = htons(net_port);
1292
1293         if (bind(fd, (struct sockaddr *) &addr, sizeof(addr)) < 0) {
1294                 perror("bind");
1295                 return 1;
1296         }
1297
1298         if (listen(fd, 1) < 0) {
1299                 perror("listen");
1300                 return 1;
1301         }
1302
1303 repeat:
1304         signal(SIGINT, NULL);
1305         signal(SIGHUP, NULL);
1306         signal(SIGTERM, NULL);
1307         signal(SIGALRM, NULL);
1308
1309         printf("blktrace: waiting for incoming connection...\n");
1310
1311         socklen = sizeof(addr);
1312         net_in_fd = accept(fd, (struct sockaddr *) &addr, &socklen);
1313         if (net_in_fd < 0) {
1314                 perror("accept");
1315                 return 1;
1316         }
1317
1318         signal(SIGINT, handle_sigint);
1319         signal(SIGHUP, handle_sigint);
1320         signal(SIGTERM, handle_sigint);
1321         signal(SIGALRM, handle_sigint);
1322
1323         printf("blktrace: connected!\n");
1324
1325         while (!is_done()) {
1326                 if (net_server_loop())
1327                         break;
1328         }
1329
1330         for_each_dip(dip, i)
1331                 for_each_tip(dip, tip, j)
1332                         tip_ftrunc_final(tip);
1333
1334         show_stats();
1335
1336         if (is_done())
1337                 return 0;
1338
1339         /*
1340          * cleanup for next run
1341          */
1342         for_each_dip(dip, i) {
1343                 for_each_tip(dip, tip, j)
1344                         fclose(tip->ofile);
1345
1346                 free(dip->threads);
1347         }
1348
1349         free(device_information);
1350         device_information = NULL;
1351         ncpus = ndevs = 0;
1352         goto repeat;
1353 }
1354
1355 /*
1356  * Setup outgoing network connection where we will transmit data
1357  */
1358 static int net_setup_client(void)
1359 {
1360         struct sockaddr_in addr;
1361         int fd;
1362
1363         fd = socket(AF_INET, SOCK_STREAM, 0);
1364         if (fd < 0) {
1365                 perror("client: socket");
1366                 return 1;
1367         }
1368
1369         memset(&addr, 0, sizeof(addr));
1370         addr.sin_family = AF_INET;
1371         addr.sin_port = htons(net_port);
1372
1373         if (inet_aton(hostname, &addr.sin_addr) != 1) {
1374                 struct hostent *hent = gethostbyname(hostname);
1375                 if (!hent) {
1376                         perror("gethostbyname");
1377                         return 1;
1378                 }
1379
1380                 memcpy(&addr.sin_addr, hent->h_addr, 4);
1381                 strcpy(hostname, hent->h_name);
1382         }
1383
1384         printf("blktrace: connecting to %s\n", hostname);
1385
1386         if (connect(fd, (struct sockaddr *) &addr, sizeof(addr)) < 0) {
1387                 perror("client: connect");
1388                 return 1;
1389         }
1390
1391         printf("blktrace: connected!\n");
1392         net_out_fd = fd;
1393         return 0;
1394 }
1395
1396 static char usage_str[] = \
1397         "-d <dev> [ -r relay path ] [ -o <output> ] [-k ] [ -w time ]\n" \
1398         "[ -a action ] [ -A action mask ] [ -v ]\n\n" \
1399         "\t-d Use specified device. May also be given last after options\n" \
1400         "\t-r Path to mounted relayfs, defaults to /relay\n" \
1401         "\t-o File(s) to send output to\n" \
1402         "\t-D Directory to prepend to output file names\n" \
1403         "\t-k Kill a running trace\n" \
1404         "\t-w Stop after defined time, in seconds\n" \
1405         "\t-a Only trace specified actions. See documentation\n" \
1406         "\t-A Give trace mask as a single value. See documentation\n" \
1407         "\t-b Sub buffer size in KiB\n" \
1408         "\t-n Number of sub buffers\n" \
1409         "\t-l Run in network listen mode (blktrace server)\n" \
1410         "\t-h Run in network client mode, connecting to the given host\n" \
1411         "\t-p Network port to use (default 8462)\n" \
1412         "\t-s Make the network client use sendfile() to transfer data\n" \
1413         "\t-V Print program version info\n\n";
1414
1415 static void show_usage(char *program)
1416 {
1417         fprintf(stderr, "Usage: %s %s %s",program, blktrace_version, usage_str);
1418 }
1419
1420 int main(int argc, char *argv[])
1421 {
1422         static char default_relay_path[] = "/relay";
1423         struct statfs st;
1424         int i, c;
1425         int stop_watch = 0;
1426         int act_mask_tmp = 0;
1427
1428         while ((c = getopt_long(argc, argv, S_OPTS, l_opts, NULL)) >= 0) {
1429                 switch (c) {
1430                 case 'a':
1431                         i = find_mask_map(optarg);
1432                         if (i < 0) {
1433                                 fprintf(stderr,"Invalid action mask %s\n",
1434                                         optarg);
1435                                 return 1;
1436                         }
1437                         act_mask_tmp |= i;
1438                         break;
1439
1440                 case 'A':
1441                         if ((sscanf(optarg, "%x", &i) != 1) || 
1442                                                         !valid_act_opt(i)) {
1443                                 fprintf(stderr,
1444                                         "Invalid set action mask %s/0x%x\n",
1445                                         optarg, i);
1446                                 return 1;
1447                         }
1448                         act_mask_tmp = i;
1449                         break;
1450
1451                 case 'd':
1452                         if (resize_devices(optarg) != 0)
1453                                 return 1;
1454                         break;
1455
1456                 case 'r':
1457                         relay_path = optarg;
1458                         break;
1459
1460                 case 'o':
1461                         output_name = optarg;
1462                         break;
1463                 case 'k':
1464                         kill_running_trace = 1;
1465                         break;
1466                 case 'w':
1467                         stop_watch = atoi(optarg);
1468                         if (stop_watch <= 0) {
1469                                 fprintf(stderr,
1470                                         "Invalid stopwatch value (%d secs)\n",
1471                                         stop_watch);
1472                                 return 1;
1473                         }
1474                         break;
1475                 case 'V':
1476                         printf("%s version %s\n", argv[0], blktrace_version);
1477                         return 0;
1478                 case 'b':
1479                         buf_size = strtoul(optarg, NULL, 10);
1480                         if (buf_size <= 0 || buf_size > 16*1024) {
1481                                 fprintf(stderr,
1482                                         "Invalid buffer size (%lu)\n",buf_size);
1483                                 return 1;
1484                         }
1485                         buf_size <<= 10;
1486                         break;
1487                 case 'n':
1488                         buf_nr = strtoul(optarg, NULL, 10);
1489                         if (buf_nr <= 0) {
1490                                 fprintf(stderr,
1491                                         "Invalid buffer nr (%lu)\n", buf_nr);
1492                                 return 1;
1493                         }
1494                         break;
1495                 case 'D':
1496                         output_dir = optarg;
1497                         break;
1498                 case 'h':
1499                         net_mode = Net_client;
1500                         strcpy(hostname, optarg);
1501                         break;
1502                 case 'l':
1503                         net_mode = Net_server;
1504                         break;
1505                 case 'p':
1506                         net_port = atoi(optarg);
1507                         break;
1508                 case 's':
1509                         net_sendfile = 1;
1510                         break;
1511                 default:
1512                         show_usage(argv[0]);
1513                         return 1;
1514                 }
1515         }
1516
1517         setlocale(LC_NUMERIC, "en_US");
1518
1519         page_size = getpagesize();
1520
1521         if (net_mode == Net_server)
1522                 return net_server();
1523
1524         while (optind < argc) {
1525                 if (resize_devices(argv[optind++]) != 0)
1526                         return 1;
1527         }
1528
1529         if (ndevs == 0) {
1530                 show_usage(argv[0]);
1531                 return 1;
1532         }
1533
1534         if (!relay_path)
1535                 relay_path = default_relay_path;
1536
1537         if (act_mask_tmp != 0)
1538                 act_mask = act_mask_tmp;
1539
1540         if (statfs(relay_path, &st) < 0) {
1541                 perror("statfs");
1542                 fprintf(stderr,"%s does not appear to be a valid path\n",
1543                         relay_path);
1544                 return 1;
1545         } else if (st.f_type != (long) RELAYFS_TYPE) {
1546                 fprintf(stderr,"%s does not appear to be a relay filesystem\n",
1547                         relay_path);
1548                 return 1;
1549         }
1550
1551         if (open_devices() != 0)
1552                 return 1;
1553
1554         if (kill_running_trace) {
1555                 stop_all_traces();
1556                 return 0;
1557         }
1558
1559         ncpus = sysconf(_SC_NPROCESSORS_ONLN);
1560         if (ncpus < 0) {
1561                 fprintf(stderr, "sysconf(_SC_NPROCESSORS_ONLN) failed\n");
1562                 return 1;
1563         }
1564
1565         signal(SIGINT, handle_sigint);
1566         signal(SIGHUP, handle_sigint);
1567         signal(SIGTERM, handle_sigint);
1568         signal(SIGALRM, handle_sigint);
1569
1570         if (net_mode == Net_client && net_setup_client())
1571                 return 1;
1572
1573         if (start_devices() != 0)
1574                 return 1;
1575
1576         atexit(stop_all_tracing);
1577
1578         if (stop_watch)
1579                 alarm(stop_watch);
1580
1581         wait_for_threads();
1582
1583         if (!is_trace_stopped()) {
1584                 trace_stopped = 1;
1585                 stop_all_threads();
1586                 stop_all_traces();
1587         }
1588
1589         show_stats();
1590
1591         return 0;
1592 }
1593