zhenwei pi | e42ac41 | 2020-08-23 19:17:52 +0800 | [diff] [blame] | 1 | #!/usr/bin/python |
| 2 | # @lint-avoid-python-3-compatibility-imports |
| 3 | # |
| 4 | # tcprtt Summarize TCP RTT as a histogram. For Linux, uses BCC, eBPF. |
| 5 | # |
| 6 | # USAGE: tcprtt [-h] [-T] [-D] [-m] [-i INTERVAL] [-d DURATION] |
zhenwei pi | 5057944 | 2021-06-21 15:06:58 +0800 | [diff] [blame] | 7 | # [-p LPORT] [-P RPORT] [-a LADDR] [-A RADDR] [-b] [-B] [-e] |
Hariharan Ananthakrishnan | 04893e3 | 2021-08-12 05:55:21 -0700 | [diff] [blame] | 8 | # [-4 | -6] |
zhenwei pi | e42ac41 | 2020-08-23 19:17:52 +0800 | [diff] [blame] | 9 | # |
| 10 | # Copyright (c) 2020 zhenwei pi |
| 11 | # Licensed under the Apache License, Version 2.0 (the "License") |
| 12 | # |
| 13 | # 23-AUG-2020 zhenwei pi Created this. |
| 14 | |
| 15 | from __future__ import print_function |
| 16 | from bcc import BPF |
| 17 | from time import sleep, strftime |
zhenwei pi | 4cbcd9a | 2020-09-29 00:07:38 +0800 | [diff] [blame] | 18 | from socket import inet_ntop, AF_INET |
zhenwei pi | e42ac41 | 2020-08-23 19:17:52 +0800 | [diff] [blame] | 19 | import socket, struct |
| 20 | import argparse |
zhenwei pi | 5057944 | 2021-06-21 15:06:58 +0800 | [diff] [blame] | 21 | import ctypes |
zhenwei pi | e42ac41 | 2020-08-23 19:17:52 +0800 | [diff] [blame] | 22 | |
| 23 | # arguments |
| 24 | examples = """examples: |
| 25 | ./tcprtt # summarize TCP RTT |
| 26 | ./tcprtt -i 1 -d 10 # print 1 second summaries, 10 times |
| 27 | ./tcprtt -m -T # summarize in millisecond, and timestamps |
zhenwei pi | 4cbcd9a | 2020-09-29 00:07:38 +0800 | [diff] [blame] | 28 | ./tcprtt -p # filter for local port |
| 29 | ./tcprtt -P # filter for remote port |
| 30 | ./tcprtt -a # filter for local address |
| 31 | ./tcprtt -A # filter for remote address |
| 32 | ./tcprtt -b # show sockets histogram by local address |
| 33 | ./tcprtt -B # show sockets histogram by remote address |
zhenwei pi | e42ac41 | 2020-08-23 19:17:52 +0800 | [diff] [blame] | 34 | ./tcprtt -D # show debug bpf text |
zhenwei pi | 5057944 | 2021-06-21 15:06:58 +0800 | [diff] [blame] | 35 | ./tcprtt -e # show extension summary(average) |
Hariharan Ananthakrishnan | 04893e3 | 2021-08-12 05:55:21 -0700 | [diff] [blame] | 36 | ./tcprtt -4 # trace only IPv4 family |
| 37 | ./tcprtt -6 # trace only IPv6 family |
zhenwei pi | e42ac41 | 2020-08-23 19:17:52 +0800 | [diff] [blame] | 38 | """ |
| 39 | parser = argparse.ArgumentParser( |
| 40 | description="Summarize TCP RTT as a histogram", |
| 41 | formatter_class=argparse.RawDescriptionHelpFormatter, |
| 42 | epilog=examples) |
| 43 | parser.add_argument("-i", "--interval", |
| 44 | help="summary interval, seconds") |
| 45 | parser.add_argument("-d", "--duration", type=int, default=99999, |
| 46 | help="total duration of trace, seconds") |
| 47 | parser.add_argument("-T", "--timestamp", action="store_true", |
| 48 | help="include timestamp on output") |
| 49 | parser.add_argument("-m", "--milliseconds", action="store_true", |
| 50 | help="millisecond histogram") |
zhenwei pi | 4cbcd9a | 2020-09-29 00:07:38 +0800 | [diff] [blame] | 51 | parser.add_argument("-p", "--lport", |
| 52 | help="filter for local port") |
| 53 | parser.add_argument("-P", "--rport", |
| 54 | help="filter for remote port") |
| 55 | parser.add_argument("-a", "--laddr", |
| 56 | help="filter for local address") |
| 57 | parser.add_argument("-A", "--raddr", |
| 58 | help="filter for remote address") |
| 59 | parser.add_argument("-b", "--byladdr", action="store_true", |
| 60 | help="show sockets histogram by local address") |
| 61 | parser.add_argument("-B", "--byraddr", action="store_true", |
| 62 | help="show sockets histogram by remote address") |
zhenwei pi | 5057944 | 2021-06-21 15:06:58 +0800 | [diff] [blame] | 63 | parser.add_argument("-e", "--extension", action="store_true", |
| 64 | help="show extension summary(average)") |
zhenwei pi | e42ac41 | 2020-08-23 19:17:52 +0800 | [diff] [blame] | 65 | parser.add_argument("-D", "--debug", action="store_true", |
| 66 | help="print BPF program before starting (for debugging purposes)") |
Hariharan Ananthakrishnan | 04893e3 | 2021-08-12 05:55:21 -0700 | [diff] [blame] | 67 | group = parser.add_mutually_exclusive_group() |
| 68 | group.add_argument("-4", "--ipv4", action="store_true", |
| 69 | help="trace IPv4 family only") |
| 70 | group.add_argument("-6", "--ipv6", action="store_true", |
| 71 | help="trace IPv6 family only") |
zhenwei pi | e42ac41 | 2020-08-23 19:17:52 +0800 | [diff] [blame] | 72 | parser.add_argument("--ebpf", action="store_true", |
| 73 | help=argparse.SUPPRESS) |
| 74 | args = parser.parse_args() |
| 75 | if not args.interval: |
| 76 | args.interval = args.duration |
| 77 | |
| 78 | # define BPF program |
| 79 | bpf_text = """ |
| 80 | #ifndef KBUILD_MODNAME |
| 81 | #define KBUILD_MODNAME "bcc" |
| 82 | #endif |
| 83 | #include <uapi/linux/ptrace.h> |
| 84 | #include <linux/tcp.h> |
| 85 | #include <net/sock.h> |
| 86 | #include <net/inet_sock.h> |
| 87 | #include <bcc/proto.h> |
| 88 | |
zhenwei pi | 4cbcd9a | 2020-09-29 00:07:38 +0800 | [diff] [blame] | 89 | typedef struct sock_key { |
| 90 | u64 addr; |
| 91 | u64 slot; |
| 92 | } sock_key_t; |
| 93 | |
zhenwei pi | 5057944 | 2021-06-21 15:06:58 +0800 | [diff] [blame] | 94 | typedef struct sock_latenty { |
| 95 | u64 latency; |
| 96 | u64 count; |
| 97 | } sock_latency_t; |
| 98 | |
| 99 | BPF_HISTOGRAM(hist_srtt, sock_key_t); |
| 100 | BPF_HASH(latency, u64, sock_latency_t); |
zhenwei pi | e42ac41 | 2020-08-23 19:17:52 +0800 | [diff] [blame] | 101 | |
| 102 | int trace_tcp_rcv(struct pt_regs *ctx, struct sock *sk, struct sk_buff *skb) |
| 103 | { |
| 104 | struct tcp_sock *ts = tcp_sk(sk); |
| 105 | u32 srtt = ts->srtt_us >> 3; |
| 106 | const struct inet_sock *inet = inet_sk(sk); |
zhenwei pi | 5057944 | 2021-06-21 15:06:58 +0800 | [diff] [blame] | 107 | |
| 108 | /* filters */ |
zhenwei pi | 4cbcd9a | 2020-09-29 00:07:38 +0800 | [diff] [blame] | 109 | u16 sport = 0; |
| 110 | u16 dport = 0; |
| 111 | u32 saddr = 0; |
| 112 | u32 daddr = 0; |
Hariharan Ananthakrishnan | 04893e3 | 2021-08-12 05:55:21 -0700 | [diff] [blame] | 113 | u16 family = 0; |
zhenwei pi | e42ac41 | 2020-08-23 19:17:52 +0800 | [diff] [blame] | 114 | |
zhenwei pi | 5057944 | 2021-06-21 15:06:58 +0800 | [diff] [blame] | 115 | /* for histogram */ |
| 116 | sock_key_t key; |
| 117 | |
| 118 | /* for avg latency, if no saddr/daddr specified, use 0(addr) as key */ |
| 119 | u64 addr = 0; |
| 120 | |
zhenwei pi | 4cbcd9a | 2020-09-29 00:07:38 +0800 | [diff] [blame] | 121 | bpf_probe_read_kernel(&sport, sizeof(sport), (void *)&inet->inet_sport); |
| 122 | bpf_probe_read_kernel(&dport, sizeof(dport), (void *)&inet->inet_dport); |
| 123 | bpf_probe_read_kernel(&saddr, sizeof(saddr), (void *)&inet->inet_saddr); |
| 124 | bpf_probe_read_kernel(&daddr, sizeof(daddr), (void *)&inet->inet_daddr); |
Hariharan Ananthakrishnan | 04893e3 | 2021-08-12 05:55:21 -0700 | [diff] [blame] | 125 | bpf_probe_read_kernel(&family, sizeof(family), (void *)&sk->__sk_common.skc_family); |
zhenwei pi | 4cbcd9a | 2020-09-29 00:07:38 +0800 | [diff] [blame] | 126 | |
| 127 | LPORTFILTER |
| 128 | RPORTFILTER |
| 129 | LADDRFILTER |
| 130 | RADDRFILTER |
Hariharan Ananthakrishnan | 04893e3 | 2021-08-12 05:55:21 -0700 | [diff] [blame] | 131 | FAMILYFILTER |
zhenwei pi | 4cbcd9a | 2020-09-29 00:07:38 +0800 | [diff] [blame] | 132 | |
zhenwei pi | e42ac41 | 2020-08-23 19:17:52 +0800 | [diff] [blame] | 133 | FACTOR |
| 134 | |
zhenwei pi | 5057944 | 2021-06-21 15:06:58 +0800 | [diff] [blame] | 135 | STORE_HIST |
| 136 | key.slot = bpf_log2l(srtt); |
zcy | 80242fb | 2021-07-02 00:12:32 +0800 | [diff] [blame] | 137 | hist_srtt.atomic_increment(key); |
zhenwei pi | 5057944 | 2021-06-21 15:06:58 +0800 | [diff] [blame] | 138 | |
| 139 | STORE_LATENCY |
zhenwei pi | e42ac41 | 2020-08-23 19:17:52 +0800 | [diff] [blame] | 140 | |
| 141 | return 0; |
| 142 | } |
| 143 | """ |
| 144 | |
zhenwei pi | 4cbcd9a | 2020-09-29 00:07:38 +0800 | [diff] [blame] | 145 | # filter for local port |
| 146 | if args.lport: |
Yonghong Song | 6ba4dc1 | 2020-10-16 11:12:53 -0700 | [diff] [blame] | 147 | bpf_text = bpf_text.replace('LPORTFILTER', |
| 148 | """if (ntohs(sport) != %d) |
zhenwei pi | 4cbcd9a | 2020-09-29 00:07:38 +0800 | [diff] [blame] | 149 | return 0;""" % int(args.lport)) |
zhenwei pi | e42ac41 | 2020-08-23 19:17:52 +0800 | [diff] [blame] | 150 | else: |
Yonghong Song | 6ba4dc1 | 2020-10-16 11:12:53 -0700 | [diff] [blame] | 151 | bpf_text = bpf_text.replace('LPORTFILTER', '') |
zhenwei pi | e42ac41 | 2020-08-23 19:17:52 +0800 | [diff] [blame] | 152 | |
zhenwei pi | 4cbcd9a | 2020-09-29 00:07:38 +0800 | [diff] [blame] | 153 | # filter for remote port |
| 154 | if args.rport: |
Yonghong Song | 6ba4dc1 | 2020-10-16 11:12:53 -0700 | [diff] [blame] | 155 | bpf_text = bpf_text.replace('RPORTFILTER', |
| 156 | """if (ntohs(dport) != %d) |
zhenwei pi | 4cbcd9a | 2020-09-29 00:07:38 +0800 | [diff] [blame] | 157 | return 0;""" % int(args.rport)) |
zhenwei pi | e42ac41 | 2020-08-23 19:17:52 +0800 | [diff] [blame] | 158 | else: |
Yonghong Song | 6ba4dc1 | 2020-10-16 11:12:53 -0700 | [diff] [blame] | 159 | bpf_text = bpf_text.replace('RPORTFILTER', '') |
zhenwei pi | e42ac41 | 2020-08-23 19:17:52 +0800 | [diff] [blame] | 160 | |
zhenwei pi | 4cbcd9a | 2020-09-29 00:07:38 +0800 | [diff] [blame] | 161 | # filter for local address |
| 162 | if args.laddr: |
Yonghong Song | 6ba4dc1 | 2020-10-16 11:12:53 -0700 | [diff] [blame] | 163 | bpf_text = bpf_text.replace('LADDRFILTER', |
| 164 | """if (saddr != %d) |
zhenwei pi | 4cbcd9a | 2020-09-29 00:07:38 +0800 | [diff] [blame] | 165 | return 0;""" % struct.unpack("=I", socket.inet_aton(args.laddr))[0]) |
zhenwei pi | e42ac41 | 2020-08-23 19:17:52 +0800 | [diff] [blame] | 166 | else: |
Yonghong Song | 6ba4dc1 | 2020-10-16 11:12:53 -0700 | [diff] [blame] | 167 | bpf_text = bpf_text.replace('LADDRFILTER', '') |
zhenwei pi | e42ac41 | 2020-08-23 19:17:52 +0800 | [diff] [blame] | 168 | |
zhenwei pi | 4cbcd9a | 2020-09-29 00:07:38 +0800 | [diff] [blame] | 169 | # filter for remote address |
| 170 | if args.raddr: |
Yonghong Song | 6ba4dc1 | 2020-10-16 11:12:53 -0700 | [diff] [blame] | 171 | bpf_text = bpf_text.replace('RADDRFILTER', |
| 172 | """if (daddr != %d) |
zhenwei pi | 4cbcd9a | 2020-09-29 00:07:38 +0800 | [diff] [blame] | 173 | return 0;""" % struct.unpack("=I", socket.inet_aton(args.raddr))[0]) |
zhenwei pi | e42ac41 | 2020-08-23 19:17:52 +0800 | [diff] [blame] | 174 | else: |
Yonghong Song | 6ba4dc1 | 2020-10-16 11:12:53 -0700 | [diff] [blame] | 175 | bpf_text = bpf_text.replace('RADDRFILTER', '') |
Hariharan Ananthakrishnan | 04893e3 | 2021-08-12 05:55:21 -0700 | [diff] [blame] | 176 | if args.ipv4: |
| 177 | bpf_text = bpf_text.replace('FAMILYFILTER', |
| 178 | 'if (family != AF_INET) { return 0; }') |
| 179 | elif args.ipv6: |
| 180 | bpf_text = bpf_text.replace('FAMILYFILTER', |
| 181 | 'if (family != AF_INET6) { return 0; }') |
| 182 | else: |
| 183 | bpf_text = bpf_text.replace('FAMILYFILTER', '') |
zhenwei pi | e42ac41 | 2020-08-23 19:17:52 +0800 | [diff] [blame] | 184 | # show msecs or usecs[default] |
| 185 | if args.milliseconds: |
| 186 | bpf_text = bpf_text.replace('FACTOR', 'srtt /= 1000;') |
| 187 | label = "msecs" |
| 188 | else: |
| 189 | bpf_text = bpf_text.replace('FACTOR', '') |
| 190 | label = "usecs" |
| 191 | |
zhenwei pi | 4cbcd9a | 2020-09-29 00:07:38 +0800 | [diff] [blame] | 192 | print_header = "srtt" |
| 193 | # show byladdr/byraddr histogram |
| 194 | if args.byladdr: |
zhenwei pi | 5057944 | 2021-06-21 15:06:58 +0800 | [diff] [blame] | 195 | bpf_text = bpf_text.replace('STORE_HIST', 'key.addr = addr = saddr;') |
| 196 | print_header = "Local Address" |
zhenwei pi | 4cbcd9a | 2020-09-29 00:07:38 +0800 | [diff] [blame] | 197 | elif args.byraddr: |
zhenwei pi | 5057944 | 2021-06-21 15:06:58 +0800 | [diff] [blame] | 198 | bpf_text = bpf_text.replace('STORE_HIST', 'key.addr = addr = daddr;') |
| 199 | print_header = "Remote Addres" |
zhenwei pi | 4cbcd9a | 2020-09-29 00:07:38 +0800 | [diff] [blame] | 200 | else: |
zhenwei pi | 5057944 | 2021-06-21 15:06:58 +0800 | [diff] [blame] | 201 | bpf_text = bpf_text.replace('STORE_HIST', 'key.addr = addr = 0;') |
| 202 | print_header = "All Addresses" |
| 203 | |
| 204 | if args.extension: |
| 205 | bpf_text = bpf_text.replace('STORE_LATENCY', """ |
| 206 | sock_latency_t newlat = {0}; |
| 207 | sock_latency_t *lat; |
| 208 | lat = latency.lookup(&addr); |
| 209 | if (!lat) { |
| 210 | newlat.latency += srtt; |
| 211 | newlat.count += 1; |
| 212 | latency.update(&addr, &newlat); |
| 213 | } else { |
| 214 | lat->latency +=srtt; |
| 215 | lat->count += 1; |
| 216 | } |
| 217 | """) |
| 218 | else: |
| 219 | bpf_text = bpf_text.replace('STORE_LATENCY', '') |
zhenwei pi | 4cbcd9a | 2020-09-29 00:07:38 +0800 | [diff] [blame] | 220 | |
zhenwei pi | e42ac41 | 2020-08-23 19:17:52 +0800 | [diff] [blame] | 221 | # debug/dump ebpf enable or not |
| 222 | if args.debug or args.ebpf: |
| 223 | print(bpf_text) |
| 224 | if args.ebpf: |
| 225 | exit() |
| 226 | |
| 227 | # load BPF program |
| 228 | b = BPF(text=bpf_text) |
| 229 | b.attach_kprobe(event="tcp_rcv_established", fn_name="trace_tcp_rcv") |
| 230 | |
| 231 | print("Tracing TCP RTT... Hit Ctrl-C to end.") |
| 232 | |
zhenwei pi | 4cbcd9a | 2020-09-29 00:07:38 +0800 | [diff] [blame] | 233 | def print_section(addr): |
zhenwei pi | 5057944 | 2021-06-21 15:06:58 +0800 | [diff] [blame] | 234 | addrstr = "*******" |
| 235 | if (addr): |
zhenwei pi | dcef7bc | 2021-06-23 16:24:11 +0800 | [diff] [blame] | 236 | addrstr = inet_ntop(AF_INET, struct.pack("I", addr)) |
zhenwei pi | 5057944 | 2021-06-21 15:06:58 +0800 | [diff] [blame] | 237 | |
| 238 | avglat = "" |
| 239 | if args.extension: |
| 240 | lats = b.get_table("latency") |
| 241 | lat = lats[ctypes.c_ulong(addr)] |
| 242 | avglat = " [AVG %d]" % (lat.latency / lat.count) |
| 243 | |
| 244 | return addrstr + avglat |
zhenwei pi | 4cbcd9a | 2020-09-29 00:07:38 +0800 | [diff] [blame] | 245 | |
zhenwei pi | e42ac41 | 2020-08-23 19:17:52 +0800 | [diff] [blame] | 246 | # output |
| 247 | exiting = 0 if args.interval else 1 |
| 248 | dist = b.get_table("hist_srtt") |
zhenwei pi | 5057944 | 2021-06-21 15:06:58 +0800 | [diff] [blame] | 249 | lathash = b.get_table("latency") |
zhenwei pi | e42ac41 | 2020-08-23 19:17:52 +0800 | [diff] [blame] | 250 | seconds = 0 |
| 251 | while (1): |
| 252 | try: |
| 253 | sleep(int(args.interval)) |
| 254 | seconds = seconds + int(args.interval) |
| 255 | except KeyboardInterrupt: |
| 256 | exiting = 1 |
| 257 | |
| 258 | print() |
| 259 | if args.timestamp: |
| 260 | print("%-8s\n" % strftime("%H:%M:%S"), end="") |
| 261 | |
zhenwei pi | 4cbcd9a | 2020-09-29 00:07:38 +0800 | [diff] [blame] | 262 | dist.print_log2_hist(label, section_header=print_header, section_print_fn=print_section) |
zhenwei pi | e42ac41 | 2020-08-23 19:17:52 +0800 | [diff] [blame] | 263 | dist.clear() |
zhenwei pi | 5057944 | 2021-06-21 15:06:58 +0800 | [diff] [blame] | 264 | lathash.clear() |
zhenwei pi | e42ac41 | 2020-08-23 19:17:52 +0800 | [diff] [blame] | 265 | |
| 266 | if exiting or seconds >= args.duration: |
| 267 | exit() |