Move open_{kprobes,uprobes,tracepoints} into BPF object

* for #605, this diff moves probe storage from the BPF module to the BPF object,
  letting each instantiation clean up its own probes. A module-level counter
  for all open probes is provided for the quota check. It also adds a
  `cleanup()` function to force cleanup before the `atexit` handler runs.

* for #614, it removes the `len(open_kprobes) == 0` check that prevented more
  than one autoload probe working. It fixes the tests that this change breaks by
  calling the `cleanup()` function added.

src/python/bcc/table.py

@@ -393,15 +393,15 @@ class PerfEventArray(ArrayBase):
             raise Exception("Could not open perf buffer")
         fd = lib.perf_reader_fd(reader)
         self[self.Key(cpu)] = self.Leaf(fd)
-        self.bpf.open_kprobes()[(id(self), cpu)] = reader
+        self.bpf._add_kprobe((id(self), cpu), reader)
         # keep a refcnt
         self._cbs[cpu] = fn
 
     def close_perf_buffer(self, key):
-        reader = self.bpf.open_kprobes().get((id(self), key))
+        reader = self.bpf.open_kprobes.get((id(self), key))
         if reader:
             lib.perf_reader_free(reader)
-            del(self.bpf.open_kprobes()[(id(self), key)])
+            self.bpf._del_kprobe((id(self), key))
         del self._cbs[key]
 
 class PerCpuHash(HashTable):

tests/python/test_histogram.py

@@ -31,6 +31,7 @@ int kprobe__htab_map_delete_elem(struct pt_regs *ctx, struct bpf_map *map, u64 *
                 try: del b["stub"][c_ulonglong(1 << i)]
                 except: pass
         b["hist1"].print_log2_hist()
+        b.cleanup()
 
     def test_struct(self):
         b = BPF(text="""
@@ -52,6 +53,7 @@ int kprobe__htab_map_delete_elem(struct pt_regs *ctx, struct bpf_map *map, u64 *
                 try: del b["stub2"][c_ulonglong(1 << i)]
                 except: pass
         b["hist1"].print_log2_hist()
+        b.cleanup()
 
     def test_chars(self):
         b = BPF(text="""
@@ -68,6 +70,7 @@ int kprobe__finish_task_switch(struct pt_regs *ctx, struct task_struct *prev) {
 """)
         for i in range(0, 100): time.sleep(0.01)
         b["hist1"].print_log2_hist()
+        b.cleanup()
 
 
 if __name__ == "__main__":

tests/python/test_probe_count.py

@@ -2,7 +2,7 @@
 # Copyright (c) Suchakra Sharma <suchakrapani.sharma@polymtl.ca>
 # Licensed under the Apache License, Version 2.0 (the "License")
 
-from bcc import BPF
+from bcc import BPF, _get_num_open_probes
 import os
 import sys
 from unittest import main, TestCase
@@ -25,6 +25,39 @@ class TestKprobeCnt(TestCase):
         open_cnt = self.b.num_open_kprobes()
         self.assertEqual(actual_cnt, open_cnt)
 
+    def tearDown(self):
+        self.b.cleanup()
+
+
+class TestProbeGlobalCnt(TestCase):
+    def setUp(self):
+        self.b1 = BPF(text="""int count(void *ctx) { return 0; }""")
+        self.b2 = BPF(text="""int count(void *ctx) { return 0; }""")
+
+    def test_probe_quota(self):
+        self.b1.attach_kprobe(event="schedule", fn_name="count")
+        self.b2.attach_kprobe(event="submit_bio", fn_name="count")
+        self.assertEqual(1, self.b1.num_open_kprobes())
+        self.assertEqual(1, self.b2.num_open_kprobes())
+        self.assertEqual(2, _get_num_open_probes())
+        self.b1.cleanup()
+        self.b2.cleanup()
+        self.assertEqual(0, _get_num_open_probes())
+
+
+class TestAutoKprobe(TestCase):
+    def setUp(self):
+        self.b = BPF(text="""
+        int kprobe__schedule(void *ctx) { return 0; }
+        int kretprobe__schedule(void *ctx) { return 0; }
+        """)
+
+    def test_count(self):
+        self.assertEqual(2, self.b.num_open_kprobes())
+
+    def tearDown(self):
+        self.b.cleanup()
+
 
 class TestProbeQuota(TestCase):
     def setUp(self):
@@ -34,6 +67,10 @@ class TestProbeQuota(TestCase):
         with self.assertRaises(Exception):
             self.b.attach_kprobe(event_re=".*", fn_name="count")
 
+    def tearDown(self):
+        self.b.cleanup()
+
+
 class TestProbeNotExist(TestCase):
     def setUp(self):
         self.b = BPF(text="""int count(void *ctx) { return 0; }""")
@@ -42,6 +79,9 @@ class TestProbeNotExist(TestCase):
         with self.assertRaises(Exception):
             b.attach_kprobe(event="___doesnotexist", fn_name="count")
 
+    def tearDown(self):
+        self.b.cleanup()
+
 
 if __name__ == "__main__":
     main()

tools/argdist.py

@@ -640,8 +640,8 @@ struct __string_t { char s[%d]; };
                 for probe in self.probes:
                         probe.attach(self.bpf)
                 if self.args.verbose:
-                        print("open uprobes: %s" % BPF.open_uprobes())
-                        print("open kprobes: %s" % BPF.open_kprobes())
+                        print("open uprobes: %s" % self.bpf.open_uprobes)
+                        print("open kprobes: %s" % self.bpf.open_kprobes)
 
         def _main_loop(self):
                 count_so_far = 0