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

the keys variable was a reference to ipv4_recv_bytes, and after merging
the keys of ipv4_send_bytes the original ipv4_recv_bytes BPF table also
contained all keys, therefore the `if k in ipv4_recv_bytes` check a few
lines below always evaluates to true

with this commit all BPF tables are copied to userspace as dicts once
(to reduce the number of syscalls) and the keys are merged in a new dict

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