Merge pull request #1873 from andihit/tcptop-improvements

tcptop: fix display of received bytes, reduce syscalls (fixes #1871)

tools/tcptop.py

@@ -32,6 +32,7 @@ from struct import pack
 from time import sleep, strftime
 from subprocess import call
 import ctypes as ct
+from collections import namedtuple, defaultdict
 
 # arguments
 def range_check(string):
@@ -188,6 +189,8 @@ if debug or args.ebpf:
     if args.ebpf:
         exit()
 
+TCPSessionKey = namedtuple('TCPSession', ['pid', 'laddr', 'lport', 'daddr', 'dport'])
+
 def pid_to_comm(pid):
     try:
         comm = open("/proc/%d/comm" % pid, "r").read().rstrip()
@@ -195,6 +198,20 @@ def pid_to_comm(pid):
     except IOError:
         return str(pid)
 
+def get_ipv4_session_key(k):
+    return TCPSessionKey(pid=k.pid,
+                         laddr=inet_ntop(AF_INET, pack("I", k.saddr)),
+                         lport=k.lport,
+                         daddr=inet_ntop(AF_INET, pack("I", k.daddr)),
+                         dport=k.dport)
+
+def get_ipv6_session_key(k):
+    return TCPSessionKey(pid=k.pid,
+                         laddr=inet_ntop(AF_INET6, pack("QQ", k.saddr0, k.saddr1)),
+                         lport=k.lport,
+                         daddr=inet_ntop(AF_INET6, pack("QQ", k.daddr0, k.daddr1)),
+                         dport=k.dport)
+
 # initialize BPF
 b = BPF(text=bpf_text)
 
@@ -223,63 +240,57 @@ while i != args.count and not exiting:
         with open(loadavg) as stats:
             print("%-8s loadavg: %s" % (strftime("%H:%M:%S"), stats.read()))
 
-    # IPv4:  build dict of all seen keys
-    keys = ipv4_recv_bytes
+    # IPv4: build dict of all seen keys
+    ipv4_throughput = defaultdict(lambda: [0, 0])
     for k, v in ipv4_send_bytes.items():
-        if k not in keys:
-            keys[k] = v
+        key = get_ipv4_session_key(k)
+        ipv4_throughput[key][0] = v.value
+    ipv4_send_bytes.clear()
 
-    if keys:
+    for k, v in ipv4_recv_bytes.items():
+        key = get_ipv4_session_key(k)
+        ipv4_throughput[key][1] = v.value
+    ipv4_recv_bytes.clear()
+
+    if ipv4_throughput:
         print("%-6s %-12s %-21s %-21s %6s %6s" % ("PID", "COMM",
             "LADDR", "RADDR", "RX_KB", "TX_KB"))
 
     # output
-    for k, v in reversed(sorted(keys.items(), key=lambda keys: keys[1].value)):
-        send_kbytes = 0
-        if k in ipv4_send_bytes:
-            send_kbytes = int(ipv4_send_bytes[k].value / 1024)
-        recv_kbytes = 0
-        if k in ipv4_recv_bytes:
-            recv_kbytes = int(ipv4_recv_bytes[k].value / 1024)
-
+    for k, (send_bytes, recv_bytes) in sorted(ipv4_throughput.items(),
+                                              key=lambda kv: sum(kv[1]),
+                                              reverse=True):
         print("%-6d %-12.12s %-21s %-21s %6d %6d" % (k.pid,
             pid_to_comm(k.pid),
-            inet_ntop(AF_INET, pack("I", k.saddr)) + ":" + str(k.lport),
-            inet_ntop(AF_INET, pack("I", k.daddr)) + ":" + str(k.dport),
-            recv_kbytes, send_kbytes))
-
-    ipv4_send_bytes.clear()
-    ipv4_recv_bytes.clear()
+            k.laddr + ":" + str(k.lport),
+            k.daddr + ":" + str(k.dport),
+            int(recv_bytes / 1024), int(send_bytes / 1024)))
 
     # IPv6: build dict of all seen keys
-    keys = ipv6_recv_bytes
+    ipv6_throughput = defaultdict(lambda: [0, 0])
     for k, v in ipv6_send_bytes.items():
-        if k not in keys:
-            keys[k] = v
+        key = get_ipv6_session_key(k)
+        ipv6_throughput[key][0] = v.value
+    ipv6_send_bytes.clear()
+
+    for k, v in ipv6_recv_bytes.items():
+        key = get_ipv6_session_key(k)
+        ipv6_throughput[key][1] = v.value
+    ipv6_recv_bytes.clear()
 
-    if keys:
+    if ipv6_throughput:
         # more than 80 chars, sadly.
         print("\n%-6s %-12s %-32s %-32s %6s %6s" % ("PID", "COMM",
             "LADDR6", "RADDR6", "RX_KB", "TX_KB"))
 
     # output
-    for k, v in reversed(sorted(keys.items(), key=lambda keys: keys[1].value)):
-        send_kbytes = 0
-        if k in ipv6_send_bytes:
-            send_kbytes = int(ipv6_send_bytes[k].value / 1024)
-        recv_kbytes = 0
-        if k in ipv6_recv_bytes:
-            recv_kbytes = int(ipv6_recv_bytes[k].value / 1024)
-
+    for k, (send_bytes, recv_bytes) in sorted(ipv6_throughput.items(),
+                                              key=lambda kv: sum(kv[1]),
+                                              reverse=True):
         print("%-6d %-12.12s %-32s %-32s %6d %6d" % (k.pid,
             pid_to_comm(k.pid),
-            inet_ntop(AF_INET6, pack("QQ", k.saddr0, k.saddr1)) + ":" +
-            str(k.lport),
-            inet_ntop(AF_INET6, pack("QQ", k.daddr0, k.daddr1)) + ":" +
-            str(k.dport),
-            recv_kbytes, send_kbytes))
-
-    ipv6_send_bytes.clear()
-    ipv6_recv_bytes.clear()
+            k.laddr + ":" + str(k.lport),
+            k.daddr + ":" + str(k.dport),
+            int(recv_bytes / 1024), int(send_bytes / 1024)))
 
     i += 1