ustat: added basic race condition protection (iovisor#2103)

added basic race condition protection for ustat.py

Author: SaveTheRbtz

tools/lib/ustat.py

@@ -20,8 +20,9 @@
 
 from __future__ import print_function
 import argparse
-from bcc import BPF, USDT
+from bcc import BPF, USDT, USDTException
 import os
+import sys
 from subprocess import call
 from time import sleep, strftime
 
@@ -62,7 +63,12 @@ def _find_targets(self):
     def _enable_probes(self):
         self.usdts = []
         for pid in self.targets:
-            usdt = USDT(pid=pid)
+            try:
+                usdt = USDT(pid=pid)
+            except USDTException:
+                # avoid race condition on pid going away.
+                print("failed to instrument %d" % pid, file=sys.stderr)
+                continue
             for event in self.events:
                 try:
                     usdt.enable_probe(event, "%s_%s" % (self.language, event))
@@ -111,6 +117,9 @@ def get_counts(self, bpf):
         for event, category in self.events.items():
             counts = bpf["%s_%s_counts" % (self.language, event)]
             for pid, count in counts.items():
+                if pid.value not in result:
+                    print("result was not found for %d" % pid.value, file=sys.stderr)
+                    continue
                 result[pid.value][category] = count.value
             counts.clear()
         return result