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

added basic race condition protection for ustat.py

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