Merge pull request #67 from iovisor/bblanco_dev

Make python table accesses more pythonic

examples/vlan_learning.py

@@ -138,11 +138,9 @@ sleep(10)
 input("Press enter to exit: ")
 
 stats_collect = {}
-for key in ingress.iter():
-    leaf = ingress.lookup(key)
+for key, leaf in ingress.items():
     stats_collect[key.value] = [leaf.tx_pkts, leaf.tx_bytes, 0, 0]
-for key in egress.iter():
-    leaf = egress.lookup(key)
+for key, leaf in egress.items():
     x = stats_collect.get(key.value, [0, 0, 0, 0])
     x[2] = leaf.tx_pkts
     x[3] = leaf.tx_bytes

src/python/bpf/__init__.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import atexit
+from collections import MutableMapping
 import ctypes as ct
 import json
 from past.builtins import basestring
@@ -85,14 +86,14 @@ class BPF(object):
             self.name = name
             self.fd = fd
 
-    class Table(object):
+    class Table(MutableMapping):
         def __init__(self, bpf, map_fd, keytype, leaftype):
             self.bpf = bpf
             self.map_fd = map_fd
             self.Key = keytype
             self.Leaf = leaftype
 
-        def lookup(self, key):
+        def __getitem__(self, key):
             key_p = ct.pointer(key)
             leaf = self.Leaf()
             leaf_p = ct.pointer(leaf)
@@ -100,18 +101,35 @@ class BPF(object):
                     ct.cast(key_p, ct.c_void_p),
                     ct.cast(leaf_p, ct.c_void_p))
             if res < 0:
-                raise Exception("Could not lookup in table")
+                raise KeyError
             return leaf
 
-        def update(self, key, leaf, flags=0):
+        def __setitem__(self, key, leaf):
             key_p = ct.pointer(key)
             leaf_p = ct.pointer(leaf)
             res = lib.bpf_update_elem(self.map_fd,
                     ct.cast(key_p, ct.c_void_p),
-                    ct.cast(leaf_p, ct.c_void_p), flags)
+                    ct.cast(leaf_p, ct.c_void_p), 0)
             if res < 0:
                 raise Exception("Could not update table")
 
+        def __len__(self):
+            i = 0
+            for k in self: i += 1
+            return i
+
+        def __delitem__(self, key):
+            key_p = ct.pointer(key)
+            res = lib.bpf_delete_elem(self.map_fd, ct.cast(key_p, ct.c_void_p))
+            if res < 0:
+                raise KeyError
+
+        def __iter__(self):
+            return BPF.Table.Iter(self, self.Key)
+
+        def iter(self): return self.__iter__()
+        def keys(self): return self.__iter__()
+
         class Iter(object):
             def __init__(self, table, keytype):
                 self.Key = keytype
@@ -125,9 +143,6 @@ class BPF(object):
                 self.key = self.table.next(self.key)
                 return self.key
 
-        def iter(self):
-            return BPF.Table.Iter(self, self.Key)
-
         def next(self, key):
             next_key = self.Key()
             next_key_p = ct.pointer(next_key)

tests/cc/test_brb.py

@@ -151,14 +151,14 @@ class TestBPFSocket(TestCase):
                       ip, br_dest_map, br_mac_map,
                       ns_eth_out, vm_mac, vm_ip):
         val = Bpf_Dest(prog_id_br, curr_br_pid)
-        self.pem_dest.update(c_uint(curr_pem_pid), val)
+        self.pem_dest[c_uint(curr_pem_pid)] = val
         val = Bpf_Dest(prog_id_pem, curr_pem_pid)
-        br_dest_map.update(c_uint(curr_br_pid), val)
+        br_dest_map[c_uint(curr_br_pid)] = val
         ifindex = ip.link_lookup(ifname=ns_eth_out)[0]
-        self.pem_port.update(c_uint(curr_pem_pid), c_uint(ifindex))
-        self.pem_ifindex.update(c_uint(ifindex), c_uint(curr_pem_pid))
+        self.pem_port[c_uint(curr_pem_pid)] = c_uint(ifindex)
+        self.pem_ifindex[c_uint(ifindex)] = c_uint(curr_pem_pid)
         mac_addr = Eth_Addr(int(EUI(vm_mac.decode())))
-        br_mac_map.update(mac_addr, c_uint(curr_br_pid))
+        br_mac_map[mac_addr] = c_uint(curr_br_pid)
 
     def attach_filter(self, ip, ifname, fd, name):
         ifindex = ip.link_lookup(ifname=ifname)[0]
@@ -185,9 +185,9 @@ class TestBPFSocket(TestCase):
         self.get_table(b)
 
         # configure jump table
-        self.jump.update(c_uint(prog_id_pem), c_uint(pem_fn.fd))
-        self.jump.update(c_uint(prog_id_br1), c_uint(br1_fn.fd))
-        self.jump.update(c_uint(prog_id_br2), c_uint(br2_fn.fd))
+        self.jump[c_uint(prog_id_pem)] = c_uint(pem_fn.fd)
+        self.jump[c_uint(prog_id_br1)] = c_uint(br1_fn.fd)
+        self.jump[c_uint(prog_id_br2)] = c_uint(br2_fn.fd)
 
         # connect pem and br1
         curr_pem_pid = curr_pem_pid + 1
@@ -205,9 +205,9 @@ class TestBPFSocket(TestCase):
 
         # connect <br1, rtr> and <br2, rtr>
         ifindex = ip.link_lookup(ifname=self.nsrtr_eth0_out)[0]
-        self.br1_rtr.update(c_uint(0), c_uint(ifindex))
+        self.br1_rtr[c_uint(0)] = c_uint(ifindex)
         ifindex = ip.link_lookup(ifname=self.nsrtr_eth1_out)[0]
-        self.br2_rtr.update(c_uint(0), c_uint(ifindex))
+        self.br2_rtr[c_uint(0)] = c_uint(ifindex)
 
         # tc filter setup with bpf programs attached
         self.attach_filter(ip, self.ns1_eth_out, pem_fn.fd, pem_fn.name)
@@ -246,7 +246,7 @@ class TestBPFSocket(TestCase):
         # ping
         subprocess.call(["ip", "netns", "exec", self.ns1, "ping", self.vm2_ip, "-c", "2"])
         # minimum one arp reply, 5 icmp reply
-        self.assertGreater(self.pem_stats.lookup(c_uint(0)).value, 5)
+        self.assertGreater(self.pem_stats[c_uint(0)].value, 5)
 
         # iperf, run server on the background
         subprocess.Popen(["ip", "netns", "exec", self.ns2, "iperf", "-s", "-xSCD"])

tests/cc/test_brb2.py

@@ -153,10 +153,10 @@ class TestBPFSocket(TestCase):
         ns2_ifindex = ip.link_lookup(ifname=self.ns2_eth_out)[0]
         br1_ifindex = ip.link_lookup(ifname=self.veth_br1_2_pem)[0]
         br2_ifindex = ip.link_lookup(ifname=self.veth_br2_2_pem)[0]
-        self.pem_dest.update(c_uint(ns1_ifindex), c_uint(br1_ifindex))
-        self.pem_dest.update(c_uint(br1_ifindex), c_uint(ns1_ifindex))
-        self.pem_dest.update(c_uint(ns2_ifindex), c_uint(br2_ifindex))
-        self.pem_dest.update(c_uint(br2_ifindex), c_uint(ns2_ifindex))
+        self.pem_dest[c_uint(ns1_ifindex)] = c_uint(br1_ifindex)
+        self.pem_dest[c_uint(br1_ifindex)] = c_uint(ns1_ifindex)
+        self.pem_dest[c_uint(ns2_ifindex)] = c_uint(br2_ifindex)
+        self.pem_dest[c_uint(br2_ifindex)] = c_uint(ns2_ifindex)
 
         # tc filter setup with bpf programs attached
         self.attach_filter(ip, self.ns1_eth_out, pem_fn.fd, pem_fn.name)
@@ -190,7 +190,7 @@ class TestBPFSocket(TestCase):
         # ping
         subprocess.call(["ip", "netns", "exec", self.ns1, "ping", self.vm2_ip, "-c", "2"])
         # minimum one arp request/reply, 5 icmp request/reply
-        self.assertGreater(self.pem_stats.lookup(c_uint(0)).value, 11)
+        self.assertGreater(self.pem_stats[c_uint(0)].value, 11)
 
         # iperf, run server on the background
         subprocess.Popen(["ip", "netns", "exec", self.ns2, "iperf", "-s", "-xSCD"])

tests/cc/test_call1.py

@@ -31,18 +31,18 @@ class TestBPFSocket(TestCase):
         ip.tc("add-filter", "bpf", ifindex, ":1", fd=ether_fn.fd,
               name=ether_fn.name, parent="1:", action="ok", classid=1)
         self.jump = b.get_table("jump", c_int, c_int)
-        self.jump.update(c_int(S_ARP), c_int(arp_fn.fd))
-        self.jump.update(c_int(S_IP), c_int(ip_fn.fd))
-        self.jump.update(c_int(S_EOP), c_int(eop_fn.fd))
+        self.jump[c_int(S_ARP)] = c_int(arp_fn.fd)
+        self.jump[c_int(S_IP)] = c_int(ip_fn.fd)
+        self.jump[c_int(S_EOP)] = c_int(eop_fn.fd)
         self.stats = b.get_table("stats", c_int, c_ulonglong)
 
     def test_jumps(self):
         udp = socket(AF_INET, SOCK_DGRAM)
         udp.sendto(b"a" * 10, ("172.16.1.1", 5000))
         udp.close()
-        self.assertGreater(self.stats.lookup(c_int(S_IP)).value, 0)
-        self.assertGreater(self.stats.lookup(c_int(S_ARP)).value, 0)
-        self.assertGreater(self.stats.lookup(c_int(S_EOP)).value, 1)
+        self.assertGreater(self.stats[c_int(S_IP)].value, 0)
+        self.assertGreater(self.stats[c_int(S_ARP)].value, 0)
+        self.assertGreater(self.stats[c_int(S_EOP)].value, 1)
 
 if __name__ == "__main__":
     main()

tests/cc/test_stat1.py

@@ -37,14 +37,18 @@ class TestBPFSocket(TestCase):
     def test_ping(self):
         cmd = ["ping", "-f", "-c", "100", "172.16.1.1"]
         check_call(cmd)
-        #for key in self.stats.iter():
-        #    leaf = self.stats.lookup(key)
+        #for key, leaf in self.stats.items():
         #    print(IPAddress(key.sip), "=>", IPAddress(key.dip),
         #          "rx", leaf.rx_pkts, "tx", leaf.tx_pkts)
         key = self.stats.Key(IPAddress("172.16.1.2").value, IPAddress("172.16.1.1").value)
-        leaf = self.stats.lookup(key)
+        leaf = self.stats[key]
         self.assertEqual(leaf.rx_pkts, 100)
         self.assertEqual(leaf.tx_pkts, 100)
+        del self.stats[key]
+        with self.assertRaises(KeyError):
+            x = self.stats[key]
+        with self.assertRaises(KeyError):
+            del self.stats[key]
 
 if __name__ == "__main__":
     main()

tests/cc/test_trace1.py

@@ -41,8 +41,7 @@ class TestKprobe(TestCase):
         with open("/etc/services", "r") as f:
             for i in range(0, 200):
                 os.read(f.fileno(), 1)
-        for key in self.stats.iter():
-            leaf = self.stats.lookup(key)
+        for key, leaf in self.stats.items():
             print("fd %x:" % key.fd, "stat1 %d" % leaf.stat1, "stat2 %d" % leaf.stat2)
 
 if __name__ == "__main__":

tests/cc/test_trace2.py

@@ -32,8 +32,7 @@ class TestTracingEvent(TestCase):
     def test_sched1(self):
         for i in range(0, 100):
             sleep(0.01)
-        for key in self.stats.iter():
-            leaf = self.stats.lookup(key)
+        for key, leaf in self.stats.items():
             print("ptr %x:" % key.ptr, "stat1 %x" % leaf.stat1)
 
 if __name__ == "__main__":

tests/cc/test_trace3.py

@@ -32,8 +32,7 @@ class TestBlkRequest(TestCase):
                              "count=1024", "bs=4096"])
             subprocess.call(["sync"])
         os.unlink("/opt/trace3.txt")
-        for key in self.latency.iter():
-            leaf = self.latency.lookup(key)
+        for key, leaf in self.latency.items():
             print("latency %u:" % key.value, "count %u" % leaf.value)
         sys.stdout.flush()
 

tests/cc/test_xlate1.py

@@ -38,15 +38,15 @@ class TestBPFFilter(TestCase):
     def test_xlate(self):
         key1 = self.xlate.Key(IPAddress("172.16.1.2").value, IPAddress("172.16.1.1").value)
         leaf1 = self.xlate.Leaf(IPAddress("192.168.1.2").value, IPAddress("192.168.1.1").value, 0, 0)
-        self.xlate.update(key1, leaf1)
+        self.xlate[key1] = leaf1
         key2 = self.xlate.Key(IPAddress("192.168.1.1").value, IPAddress("192.168.1.2").value)
         leaf2 = self.xlate.Leaf(IPAddress("172.16.1.1").value, IPAddress("172.16.1.2").value, 0, 0)
-        self.xlate.update(key2, leaf2)
+        self.xlate[key2] = leaf2
         call(["ping", "-c1", "192.168.1.1"])
-        leaf = self.xlate.lookup(key1)
+        leaf = self.xlate[key1]
         self.assertGreater(leaf.ip_xlated_pkts, 0)
         self.assertGreater(leaf.arp_xlated_pkts, 0)
-        leaf = self.xlate.lookup(key2)
+        leaf = self.xlate[key2]
         self.assertGreater(leaf.ip_xlated_pkts, 0)
         self.assertGreater(leaf.arp_xlated_pkts, 0)