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