Merge pull request iovisor#152 from iovisor/bblanco_dev

Add getitem api to BPF object to reference tables
egggHang · Aug 20, 2015 · d3c27a7 · d3c27a7
2 parents ef67bb2 + 01ef55f
commit d3c27a7
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 6 deletions.
diff --git a/examples/vfsreadlat.py b/examples/vfsreadlat.py
@@ -41,7 +41,6 @@ def usage():
 b = BPF(src_file = "vfsreadlat.c")
 b.attach_kprobe(event="vfs_read", fn_name="do_entry")
 b.attach_kretprobe(event="vfs_read", fn_name="do_return")
-dist = b.get_table("dist")
 dist_max = 64
 
 # header
@@ -69,7 +68,7 @@ def print_log2_hist(d, val_type):
 	val_max = 0
 	for i in range(1, dist_max + 1):
 		try:
-			val = dist[c_int(i)].value - last[i]
+			val = b["dist"][c_int(i)].value - last[i]
 			if (val > 0):
 				idx_max = i
 			if (val > val_max):
@@ -84,10 +83,10 @@ def print_log2_hist(d, val_type):
 		if (low == high):
 			low -= 1
 		try:
-			val = dist[c_int(i)].value - last[i]
+			val = b["dist"][c_int(i)].value - last[i]
 			print("%8d -> %-8d : %-8d |%-*s|" % (low, high, val,
 			    stars_max, stars(val, val_max, stars_max)))
-			last[i] = dist[c_int(i)].value
+			last[i] = b["dist"][c_int(i)].value
 		except:
 			break
 
@@ -105,6 +104,6 @@ def print_log2_hist(d, val_type):
 		pass; do_exit = 1
 
 	print
-	print_log2_hist(dist, "usecs")
+	print_log2_hist(b["dist"], "usecs")
 	if do_exit:
 		exit()
diff --git a/src/python/bpf/__init__.py b/src/python/bpf/__init__.py
@@ -223,6 +223,7 @@ def _find_file(filename):
     def __init__(self, src_file="", hdr_file="", text=None, debug=0):
         self.debug = debug
         self.funcs = {}
+        self.tables = {}
         if text:
             self.module = lib.bpf_module_create_c_from_string(text.encode("ascii"), self.debug)
         else:
@@ -297,7 +298,7 @@ def get_table(self, name, keytype=None, leaftype=None):
         map_id = lib.bpf_table_id(self.module, name.encode("ascii"))
         map_fd = lib.bpf_table_fd(self.module, name.encode("ascii"))
         if map_fd < 0:
-            raise Exception("Failed to find BPF Table %s" % name)
+            raise KeyError
         if not keytype:
             key_desc = lib.bpf_table_key_desc(self.module, name.encode("ascii"))
             if not key_desc:
@@ -310,6 +311,23 @@ def get_table(self, name, keytype=None, leaftype=None):
             leaftype = BPF._decode_table_type(json.loads(leaf_desc.decode()))
         return BPF.Table(self, map_id, map_fd, keytype, leaftype)
 
+    def __getitem__(self, key):
+        if key not in self.tables:
+            self.tables[key] = self.get_table(key)
+        return self.tables[key]
+
+    def __setitem__(self, key, leaf):
+        self.tables[key] = leaf
+
+    def __len__(self):
+        return len(self.tables)
+
+    def __delitem__(self, key):
+        del self.tables[key]
+
+    def __iter__(self):
+        return self.tables.__iter__()
+
     @staticmethod
     def attach_raw_socket(fn, dev):
         if not isinstance(fn, BPF.Function):