Move shared example code into simulation.py

* Create a class for the examples to share, helps with IPDB cleanup and
  namespace creation.
Signed-off-by: Brenden Blanco <bblanco@plumgrid.com>

examples/simulation.py

+import atexit
+from pyroute2 import IPRoute, NetNS, IPDB, NSPopen
+
+class Simulation(object):
+    """
+    Helper class for controlling multiple namespaces. Inherit from
+    this class and setup your namespaces.
+    """
+
+    def __init__(self, ipdb):
+        self.ipdb = ipdb
+        self.ipdbs = {}
+        self.namespaces = []
+        self.processes = []
+        self.released = False
+
+    # helper function to create a namespace and a veth connecting it
+    def _create_ns(self, name, in_ifc=None, out_ifc=None, ipaddr=None,
+                   macaddr=None, fn=None, cmd=None, action="ok"):
+        ns_ipdb = IPDB(nl=NetNS(name))
+        if in_ifc:
+            in_ifname = in_ifc.ifname
+        else:
+            out_ifc = self.ipdb.create(ifname="%sa" % name, kind="veth",
+                                       peer="%sb" % name).commit()
+            in_ifc = self.ipdb.interfaces[out_ifc.peer]
+            in_ifname = in_ifc.ifname
+        with in_ifc as v:
+            # move half of veth into namespace
+            v.net_ns_fd = ns_ipdb.nl.netns
+        in_ifc = ns_ipdb.interfaces[in_ifname]
+        if out_ifc: out_ifc.up().commit()
+        with in_ifc as v:
+            v.ifname = "eth0"
+            if ipaddr: v.add_ip("%s" % ipaddr)
+            if macaddr: v.address = macaddr
+            v.up()
+        if fn and out_ifc:
+            self.ipdb.nl.tc("add", "ingress", out_ifc["index"], "ffff:")
+            self.ipdb.nl.tc("add-filter", "bpf", out_ifc["index"], ":1",
+                            fd=fn.fd, name=fn.name, parent="ffff:",
+                            action=action, classid=1)
+        self.ipdbs[ns_ipdb.nl.netns] = ns_ipdb
+        self.namespaces.append(ns_ipdb.nl)
+        if cmd:
+            self.processes.append(NSPopen(ns_ipdb.nl.netns, cmd))
+        return (ns_ipdb, out_ifc, in_ifc)
+
+    def release(self):
+        if self.released: return
+        self.released = True
+        for p in self.processes:
+            if p.released: continue
+            p.kill(); p.wait(); p.release()
+        for name, db in self.ipdbs.items(): db.release()
+        for ns in self.namespaces: ns.remove()
+

examples/tc_neighbor_sharing.py

@@ -27,6 +27,7 @@
 
 from bpf import BPF
 from pyroute2 import IPRoute, NetNS, IPDB, NSPopen
+from simulation import Simulation
 import sys
 from time import sleep
 
@@ -42,16 +43,14 @@ num_neighbors = 3
 num_locals = 2
 
 # class to build the simulation network
-class SharedNetSimulation(object):
+class SharedNetSimulation(Simulation):
 
-    def __init__(self):
-        self.ipdbs = []
-        self.namespaces = []
-        self.processes = []
+    def __init__(self, ipdb):
+        super(SharedNetSimulation, self).__init__(ipdb)
 
         # Create the wan namespace, and attach an ingress filter for throttling
         # inbound (download) traffic
-        (self.wan, wan_if) = self._create_ns("wan0", "172.16.1.5/24", None)
+        wan_if = self._create_ns("wan0", ipaddr="172.16.1.5/24")[1]
         ipr.tc("add", "ingress", wan_if["index"], "ffff:")
         ipr.tc("add-filter", "bpf", wan_if["index"], ":1", fd=wan_fn.fd,
                prio=1, name=wan_fn.name, parent="ffff:", action="drop",
@@ -61,39 +60,22 @@ class SharedNetSimulation(object):
                classid=2, rate="1024kbit", burst=1024 * 32, mtu=16 * 1024)
         self.wan_if = wan_if
 
-    # helper function to create a namespace and a veth connecting it
-    def _create_ns(self, name, ipaddr, fn):
-        ns_ipdb = IPDB(nl=NetNS(name))
-        ipdb.create(ifname="%sa" % name, kind="veth", peer="%sb" % name).commit()
-        with ipdb.interfaces["%sb" % name] as v:
-            # move half of veth into namespace
-            v.net_ns_fd = ns_ipdb.nl.netns
-        with ipdb.interfaces["%sa" % name] as v:
-            v.up()
-        with ns_ipdb.interfaces["%sb" % name] as v:
-            v.ifname = "eth0"
-            v.add_ip("%s" % ipaddr)
-            v.up()
-        ifc = ipdb.interfaces["%sa" % name]
-        if fn:
-            ipr.tc("add", "ingress", ifc["index"], "ffff:")
-            ipr.tc("add-filter", "bpf", ifc["index"], ":1", fd=fn.fd, name=fn.name,
-                   parent="ffff:", action="ok", classid=1)
-        self.ipdbs.append(ns_ipdb)
-        self.namespaces.append(ns_ipdb.nl)
-        cmd = ["netserver", "-D"]
-        self.processes.append(NSPopen(ns_ipdb.nl.netns, cmd))
-        return (ns_ipdb, ifc)
-
-    # start the namespaces that compose the network, interconnect them with the bridge,
-    # and attach the tc filters
+    # start the namespaces that compose the network, interconnect them with the
+    # bridge, and attach the tc filters
     def start(self):
         neighbor_list = []
         local_list = []
+        cmd = ["netserver", "-D"]
         for i in range(0, num_neighbors):
-            neighbor_list.append(self._create_ns("neighbor%d" % i, "172.16.1.%d/24" % (i + 100), neighbor_fn))
+            ipaddr = "172.16.1.%d/24" % (i + 100)
+            ret = self._create_ns("neighbor%d" % i, ipaddr=ipaddr,
+                                  fn=neighbor_fn, cmd=cmd)
+            neighbor_list.append(ret)
         for i in range(0, num_locals):
-            local_list.append(self._create_ns("local%d" % i, "172.16.1.%d/24" % (i + 150), pass_fn))
+            ipaddr = "172.16.1.%d/24" % (i + 150)
+            ret = self._create_ns("local%d" % i, ipaddr=ipaddr,
+                                  fn=pass_fn, cmd=cmd)
+            local_list.append(ret)
 
         with ipdb.create(ifname="br100", kind="bridge") as br100:
             for x in neighbor_list:
@@ -103,13 +85,8 @@ class SharedNetSimulation(object):
             br100.add_port(self.wan_if)
             br100.up()
 
-    def release(self):
-        for p in self.processes: p.kill(); p.release()
-        for db in self.ipdbs: db.release()
-        for ns in self.namespaces: ns.remove()
-
 try:
-    sim = SharedNetSimulation()
+    sim = SharedNetSimulation(ipdb)
     sim.start()
     print("Network ready. Create a shell in the wan0 namespace and test with netperf")
     print("   (Neighbors are 172.16.1.100-%d, and LAN clients are 172.16.1.150-%d)"
@@ -119,7 +96,6 @@ try:
 finally:
     if "sim" in locals(): sim.release()
     if "br100" in ipdb.interfaces: ipdb.interfaces.br100.remove().commit()
-    sleep(2)
     ipdb.release()
 
 

examples/vlan_learning.py

@@ -28,10 +28,10 @@ from ctypes import c_uint, c_int, c_ulonglong, Structure
 from pyroute2 import IPRoute, NetNS, IPDB, NSPopen
 from random import shuffle
 from time import sleep
+from simulation import Simulation
 import sys
 
 ipr = IPRoute()
-ipr.bind(async=True)
 ipdb = IPDB(nl=ipr)
 
 num_workers = 3
@@ -39,128 +39,96 @@ num_clients = 9
 num_vlans = 16
 
 # load the bpf program
-b = BPF(src_file="examples/vlan_learning.c", debug=0)
+b = BPF(src_file="vlan_learning.c", debug=0)
 phys_fn = b.load_func("handle_phys2virt", BPF.SCHED_CLS)
 virt_fn = b.load_func("handle_virt2phys", BPF.SCHED_CLS)
 
 ingress = b.get_table("ingress")
 egress = b.get_table("egress")
 
-ipdb_workers = []
-ipdb_clients = []
-ns_workers = []
-ns_clients = []
-worker_processes = []
-client_processes = []
-
-# start the worker namespaces: 1 veth pair, 1 http daemon
-for i in range(0, num_workers):
-    worker = IPDB(nl=NetNS("worker%d" % i))
-    ipdb.create(ifname="wrk%dp0" % i, kind="veth", peer="wrk%dp1" % i).commit()
-    with ipdb.interfaces["wrk%dp0" % i] as v:
-        v.net_ns_fd = worker.nl.netns
-    with ipdb.interfaces["wrk%dp1" % i] as v:
-        ipr.tc("add", "ingress", v["index"], "ffff:")
-        ipr.tc("add-filter", "bpf", v["index"], ":1", fd=virt_fn.fd,
-               name=virt_fn.name, parent="ffff:", action="drop", classid=1)
-        v.up()
-    # use the same ip address in each namespace, clients only need to know
-    # one destination IP!
-    with worker.interfaces["wrk%dp0" % i] as v:
-        v.ifname = "eth0"
-        v.add_ip("172.16.1.5/24")
+class VlanSimulation(Simulation):
+    def __init__(self, ipdb):
+        super(VlanSimulation, self).__init__(ipdb)
+
+    def start(self):
+        # start identical workers each in a namespace
+        for i in range(0, num_workers):
+            httpmod = ("SimpleHTTPServer" if sys.version_info[0] < 3
+                       else "http.server")
+            cmd = ["python", "-m", httpmod, "80"]
+            self._create_ns("worker%d" % i, cmd=cmd, fn=virt_fn, action="drop",
+                            ipaddr="172.16.1.5/24")
+
+        # simulate a physical eth vlan trunk
+        with self.ipdb.create(ifname="eth0a", kind="veth", peer="eth0b") as v:
             v.up()
-    httpmod = "SimpleHTTPServer" if sys.version_info[0] < 3 else "http.server"
-    worker_processes.append(NSPopen(worker.nl.netns, ["python", "-m", httpmod, "80"]))
-    ipdb_workers.append(worker)
-    ns_workers.append(worker.nl)
+        self.ipdb.interfaces.eth0b.up().commit()
 
-# simulate a physical eth vlan trunk
-with ipdb.create(ifname="eth0a", kind="veth", peer="eth0b") as v:
-    v.up()
-ipdb.interfaces.eth0b.up().commit()
-# connect the veth to the bridge
-with ipdb.create(ifname="br100", kind="bridge") as br100:
-    br100.add_port(ipdb.interfaces.eth0b)
+        # connect the trunk to the bridge
+        with self.ipdb.create(ifname="br100", kind="bridge") as br100:
+            br100.add_port(self.ipdb.interfaces.eth0b)
             br100.up()
 
-# for each vlan, create a subinterface on the eth...most of these will be
-# unused, but still listening and waiting for a client to send traffic on
-for i in range(2, 2 + num_vlans):
-    with ipdb.create(ifname="eth0a.%d" % i, kind="vlan",
+        # for each vlan, create a subinterface on the eth...most of these will be
+        # unused, but still listening and waiting for a client to send traffic on
+        for i in range(2, 2 + num_vlans):
+            with self.ipdb.create(ifname="eth0a.%d" % i, kind="vlan",
                                   link=ipdb.interfaces.eth0a, vlan_id=i) as v:
                 v.up()
-    v = ipdb.interfaces["eth0a.%d" % i]
+            v = self.ipdb.interfaces["eth0a.%d" % i]
             # add the bpf program for demuxing phys2virt packets
             ipr.tc("add", "ingress", v["index"], "ffff:")
             ipr.tc("add-filter", "bpf", v["index"], ":1", fd=phys_fn.fd,
                    name=phys_fn.name, parent="ffff:", action="drop", classid=1)
 
-# allocate vlans randomly
-available_vlans = [i for i in range(2, 2 + num_vlans)]
-shuffle(available_vlans)
-available_ips = [[i for i in range(100, 105)] for i in range(0, num_workers)]
-
-# these are simulations of physical clients
-for i in range(0, num_clients):
-    worker_choice = i % num_workers
-    client = IPDB(nl=NetNS("client%d" % i))
-    with ipdb.create(ifname="br100.%d" % i, kind="vlan",
-                     link=br100, vlan_id=available_vlans.pop(0)) as v:
-        v.net_ns_fd = client.nl.netns
-    ipaddr = "172.16.1.%d" % available_ips[worker_choice].pop(0)
-    with client.interfaces["br100.%d" % i] as v:
-        v.add_ip("%s/24" % ipaddr)
-        v.ifname = "eth0"
-        v.address = "02:00:00:%.2x:%.2x:%.2x" % ((i >> 16) & 0xff, (i >> 8) & 0xff, i & 0xff)
-        v.up()
-    macaddr = client.interfaces.eth0.address
+        # allocate vlans randomly
+        available_vlans = [i for i in range(2, 2 + num_vlans)]
+        shuffle(available_vlans)
+        available_ips = [[i for i in range(100, 105)] for i in range(0, num_workers)]
+
+        # these are simulations of physical clients
+        for i in range(0, num_clients):
+            worker_i = i % num_workers
+            ipaddr = "172.16.1.%d" % available_ips[worker_i].pop(0)
+            macaddr = ("02:00:00:%.2x:%.2x:%.2x" %
+                       ((i >> 16) & 0xff, (i >> 8) & 0xff, i & 0xff))
+
             # program arp manually
-    p = NSPopen(ipdb_workers[worker_choice].nl.netns, ["arp", "-s", ipaddr, macaddr])
-    p.communicate()
+            p = NSPopen("worker%d" % worker_i, ["arp", "-s", ipaddr, macaddr])
+            p.communicate(); p.wait(); p.release()
 
             # assign this client to the given worker
-    idx = ipdb.interfaces["wrk%dp1" % worker_choice]["index"]
+            idx = self.ipdb.interfaces["worker%da" % worker_i]["index"]
             mac = int(macaddr.replace(":", ""), 16)
             ingress[ingress.Key(mac)] = ingress.Leaf(idx, 0, 0)
 
-    cmd = ["bash", "-c", "for i in {1..8}; do curl 172.16.1.5 -o /dev/null; sleep 1; done"]
-    client_processes.append(NSPopen(client.nl.netns, cmd))
-
-    ipdb_clients.append(client)
-    ns_clients.append(client.nl)
-
-# IPDBs are no longer needed
-for db in ipdb_workers: db.release()
-for db in ipdb_clients: db.release()
-
-sleep(10)
-input("Press enter to exit: ")
-
-stats_collect = {}
-for key, leaf in ingress.items():
+            # test traffic with curl loop
+            cmd = ["bash", "-c",
+                   "for i in {1..8}; do curl 172.16.1.5 -o /dev/null; sleep 1; done"]
+            br_ifc = self.ipdb.create(ifname="br100.%d" % i, kind="vlan", link=br100,
+                                      vlan_id=available_vlans.pop(0)).commit()
+            (out_ifc, in_ifc) = self._create_ns("client%d" % i, in_ifc=br_ifc,
+                                                ipaddr=ipaddr + "/24", macaddr=macaddr,
+                                                cmd=cmd)[1:3]
+
+try:
+    sim = VlanSimulation(ipdb)
+    sim.start()
+    sleep(10)
+    input("Press enter to exit: ")
+
+    stats_collect = {}
+    for key, leaf in ingress.items():
         stats_collect[key.value] = [leaf.tx_pkts, leaf.tx_bytes, 0, 0]
-for key, leaf in egress.items():
+    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
-for k, v in stats_collect.items():
+    for k, v in stats_collect.items():
         print("mac %.12x rx pkts = %u, rx bytes = %u" % (k, v[0], v[1]))
         print("                 tx pkts = %u, tx bytes = %u" % (v[2], v[3]))
-
-print("Killing worker processes")
-for w in worker_processes:
-    w.kill()
-    w.release()
-
-for c in client_processes:
-    c.kill()
-    c.release()
-
-print("Removing namespaces and simulation interfaces")
-for ns in ns_workers: ns.remove()
-for ns in ns_clients: ns.remove()
-ipdb.interfaces.br100.remove().commit()
-ipdb.interfaces.eth0a.remove().commit()
-sleep(2)
-ipdb.release()
+finally:
+    if "eth0a" in ipdb.interfaces: ipdb.interfaces.eth0a.remove().commit()
+    if "br100" in ipdb.interfaces: ipdb.interfaces.br100.remove().commit()
+    if "sim" in locals(): sim.release()
+    ipdb.release()