Merge pull request iovisor#198 from iovisor/bblanco_dev

Autoload kprobes for all types of trace_* functions
sjtuhjh · Sep 10, 2015 · 076e8ab · 076e8ab
2 parents d014a29 + b5aba54
commit 076e8ab
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 23 deletions.
diff --git a/examples/disksnoop.c b/examples/disksnoop.c
@@ -18,7 +18,7 @@ struct key_t {
 };
 BPF_HASH(start, struct key_t);
 
-int do_request(struct pt_regs *ctx, struct request *req) {
+int kprobe__blk_start_request(struct pt_regs *ctx, struct request *req) {
 	struct key_t key = {};
 	u64 ts;
 
@@ -30,7 +30,7 @@ int do_request(struct pt_regs *ctx, struct request *req) {
 	return 0;
 }
 
-int do_completion(struct pt_regs *ctx, struct request *req) {
+int kprobe__blk_update_request(struct pt_regs *ctx, struct request *req) {
 	struct key_t key = {};
 	u64 *tsp, delta;
 

diff --git a/examples/disksnoop.py b/examples/disksnoop.py
@@ -17,8 +17,6 @@
 
 # load BPF program
 b = BPF(src_file="disksnoop.c")
-b.attach_kprobe(event="blk_start_request", fn_name="do_request")
-b.attach_kprobe(event="blk_update_request", fn_name="do_completion")
 
 # header
 print("%-18s %-2s %-7s %8s" % ("TIME(s)", "T", "BYTES", "LAT(ms)"))

diff --git a/src/python/bcc/__init__.py b/src/python/bcc/__init__.py
@@ -539,13 +539,23 @@ def detach_kretprobe(event):
             raise Exception("Failed to detach BPF from kprobe")
         del open_kprobes[ev_name]
 
-    @staticmethod
-    def trace_open(nonblocking=False):
+    def _trace_autoload(self):
+        # Cater to one-liner case where attach_kprobe is omitted and C function
+        # name matches that of the kprobe.
+        if len(open_kprobes) == 0:
+            fns = self.load_funcs(BPF.KPROBE)
+            for fn in fns:
+                if fn.name.startswith("kprobe__"):
+                    self.attach_kprobe(event=fn.name[8:], fn_name=fn.name)
+                elif fn.name.startswith("kretprobe__"):
+                    self.attach_kretprobe(event=fn.name[11:], fn_name=fn.name)
+
+    def trace_open(self, nonblocking=False):
         """trace_open(nonblocking=False)
 
         Open the trace_pipe if not already open
         """
-
+        self._trace_autoload()
         global tracefile
         if not tracefile:
             tracefile = open("%s/trace_pipe" % TRACEFS)
@@ -555,16 +565,15 @@ def trace_open(nonblocking=False):
                 fcntl.fcntl(fd, fcntl.F_SETFL, fl | os.O_NONBLOCK)
         return tracefile
 
-    @staticmethod
-    def trace_fields(nonblocking=False):
+    def trace_fields(self, nonblocking=False):
         """trace_fields(nonblocking=False)
 
         Read from the kernel debug trace pipe and return a tuple of the
         fields (task, pid, cpu, flags, timestamp, msg) or None if no
         line was read (nonblocking=True)
         """
         while True:
-            line = BPF.trace_readline(nonblocking)
+            line = self.trace_readline(nonblocking)
             if not line and nonblocking: return (None,) * 6
             # don't print messages related to lost events
             if line.startswith("CPU:"): continue
@@ -576,15 +585,14 @@ def trace_fields(nonblocking=False):
             msg = line[ts_end + 4:]
             return (task, int(pid), int(cpu), flags, float(ts), msg)
 
-    @staticmethod
-    def trace_readline(nonblocking=False):
+    def trace_readline(self, nonblocking=False):
         """trace_readline(nonblocking=False)
 
         Read from the kernel debug trace pipe and return one line
         If nonblocking is False, this will block until ctrl-C is pressed.
         """
 
-        trace = BPF.trace_open(nonblocking)
+        trace = self.trace_open(nonblocking)
 
         line = None
         try:
@@ -604,16 +612,6 @@ def trace_print(self, fmt=None):
         example: trace_print(fmt="pid {1}, msg = {5}")
         """
 
-        # Cater to one-liner case where attach_kprobe is omitted and C function
-        # name matches that of the kprobe.
-        if len(open_kprobes) == 0:
-            fns = self.load_funcs(BPF.KPROBE)
-            for fn in fns:
-                if fn.name.startswith("kprobe__"):
-                    self.attach_kprobe(event=fn.name[8:], fn_name=fn.name)
-                elif fn.name.startswith("kretprobe__"):
-                    self.attach_kprobe(event=fn.name[11:], fn_name=fn.name)
-
         while True:
             if fmt:
                 fields = self.trace_fields(nonblocking=False)