Merge pull request #964 from goldshtn/utools-missing-probes

u* tools: Gracefully handle missing probes

src/python/bcc/__init__.py

@@ -1065,4 +1065,4 @@ class BPF(object):
         self.cleanup()
 
 
-from .usdt import USDT
+from .usdt import USDT, USDTException

src/python/bcc/usdt.py

@@ -13,10 +13,14 @@
 # limitations under the License.
 
 import ctypes as ct
+import sys
 from .libbcc import lib, _USDT_CB, _USDT_PROBE_CB, \
                     bcc_usdt_location, bcc_usdt_argument, \
                     BCC_USDT_ARGUMENT_FLAGS
 
+class USDTException(Exception):
+    pass
+
 class USDTProbeArgument(object):
     def __init__(self, argument):
         self.signed = argument.size < 0
@@ -77,8 +81,9 @@ class USDTProbeLocation(object):
         res = lib.bcc_usdt_get_argument(self.probe.context, self.probe.name,
                                         self.index, index, ct.pointer(arg))
         if res != 0:
-            raise Exception("error retrieving probe argument %d location %d" %
-                            (index, self.index))
+            raise USDTException(
+                    "error retrieving probe argument %d location %d" %
+                    (index, self.index))
         return USDTProbeArgument(arg)
 
 class USDTProbe(object):
@@ -103,7 +108,7 @@ class USDTProbe(object):
         res = lib.bcc_usdt_get_location(self.context, self.name,
                                         index, ct.pointer(loc))
         if res != 0:
-            raise Exception("error retrieving probe location %d" % index)
+            raise USDTException("error retrieving probe location %d" % index)
         return USDTProbeLocation(self, index, loc)
 
 class USDT(object):
@@ -112,20 +117,33 @@ class USDT(object):
             self.pid = pid
             self.context = lib.bcc_usdt_new_frompid(pid)
             if self.context == None:
-                raise Exception("USDT failed to instrument PID %d" % pid)
+                raise USDTException("USDT failed to instrument PID %d" % pid)
         elif path:
             self.path = path
             self.context = lib.bcc_usdt_new_frompath(path)
             if self.context == None:
-                raise Exception("USDT failed to instrument path %s" % path)
+                raise USDTException("USDT failed to instrument path %s" % path)
         else:
-            raise Exception("either a pid or a binary path must be specified")
+            raise USDTException(
+                    "either a pid or a binary path must be specified")
 
     def enable_probe(self, probe, fn_name):
         if lib.bcc_usdt_enable_probe(self.context, probe, fn_name) != 0:
-            raise Exception(("failed to enable probe '%s'; a possible cause " +
-                            "can be that the probe requires a pid to enable") %
-                            probe)
+            raise USDTException(
+                    ("failed to enable probe '%s'; a possible cause " +
+                     "can be that the probe requires a pid to enable") %
+                     probe
+                  )
+
+    def enable_probe_or_bail(self, probe, fn_name):
+        if lib.bcc_usdt_enable_probe(self.context, probe, fn_name) != 0:
+            print(
+"""Error attaching USDT probes: the specified pid might not contain the
+given language's runtime, or the runtime was not built with the required
+USDT probes. Look for a configure flag similar to --with-dtrace or
+--enable-dtrace. To check which probes are present in the process, use the
+tplist tool.""")
+            sys.exit(1)
 
     def get_text(self):
         return lib.bcc_usdt_genargs(self.context).decode()

tools/ucalls.py

@@ -219,9 +219,9 @@ int syscall_return(struct pt_regs *ctx) {
 
 if args.language:
     usdt = USDT(pid=args.pid)
-    usdt.enable_probe(entry_probe, "trace_entry")
+    usdt.enable_probe_or_bail(entry_probe, "trace_entry")
     if args.latency:
-        usdt.enable_probe(return_probe, "trace_return")
+        usdt.enable_probe_or_bail(return_probe, "trace_return")
 else:
     usdt = None
 

tools/uflow.py

@@ -109,7 +109,7 @@ def enable_probe(probe_name, func_name, read_class, read_method, is_return):
                              .replace("FILTER_METHOD", filter_method)   \
                              .replace("DEPTH", depth)                   \
                              .replace("UPDATE", update)
-    usdt.enable_probe(probe_name, func_name)
+    usdt.enable_probe_or_bail(probe_name, func_name)
 
 usdt = USDT(pid=args.pid)
 

tools/ugc.py

@@ -103,8 +103,8 @@ int trace_%s(struct pt_regs *ctx) {
         return text
 
     def attach(self):
-        usdt.enable_probe(self.begin, "trace_%s" % self.begin)
-        usdt.enable_probe(self.end, "trace_%s" % self.end)
+        usdt.enable_probe_or_bail(self.begin, "trace_%s" % self.begin)
+        usdt.enable_probe_or_bail(self.end, "trace_%s" % self.end)
 
     def format(self, data):
         return self.formatter(data)

tools/uobjnew.py

@@ -78,7 +78,7 @@ int alloc_entry(struct pt_regs *ctx) {
     return 0;
 }
     """
-    usdt.enable_probe("object__alloc", "alloc_entry")
+    usdt.enable_probe_or_bail("object__alloc", "alloc_entry")
 #
 # Ruby
 #
@@ -107,10 +107,10 @@ int object_alloc_entry(struct pt_regs *ctx) {
     return 0;
 }
     """
-    usdt.enable_probe("object__create", "object_alloc_entry")
+    usdt.enable_probe_or_bail("object__create", "object_alloc_entry")
     for thing in ["string", "hash", "array"]:
         program += create_template.replace("THETHING", thing)
-        usdt.enable_probe("%s__create" % thing, "%s_alloc_entry" % thing)
+        usdt.enable_probe_or_bail("%s__create" % thing, "%s_alloc_entry" % thing)
 #
 # C
 #

tools/uthreads.py

@@ -57,7 +57,7 @@ int trace_pthread(struct pt_regs *ctx) {
     return 0;
 }
 """
-usdt.enable_probe("pthread_start", "trace_pthread")
+usdt.enable_probe_or_bail("pthread_start", "trace_pthread")
 
 if args.language == "java":
     template = """
@@ -78,8 +78,8 @@ int %s(struct pt_regs *ctx) {
     """
     program += template % ("trace_start", "start")
     program += template % ("trace_stop", "stop")
-    usdt.enable_probe("thread__start", "trace_start")
-    usdt.enable_probe("thread__stop", "trace_stop")
+    usdt.enable_probe_or_bail("thread__start", "trace_start")
+    usdt.enable_probe_or_bail("thread__stop", "trace_stop")
 
 if args.verbose:
     print(usdt.get_text())