- fix description
[util-vserver.git] / python / bwlimit.py
index 1d2d8cf..c906b2a 100644 (file)
 # Mark Huang <mlhuang@cs.princeton.edu>
 # Copyright (C) 2006 The Trustees of Princeton University
 #
-# $Id: bwlimit.py,v 1.5 2006/02/27 01:58:09 mlhuang Exp $
+# $Id: bwlimit.py,v 1.8 2006/03/01 22:02:52 mlhuang Exp $
 #
 
 import sys, os, re, getopt
+from sets import Set
+import plcapi
 
 
 # Where the tc binary lives
@@ -113,37 +115,41 @@ cburst = None
 # "default" subclass 1:10 that is capped at the node bandwidth cap (in
 # this example, 5mbit) and the "exempt" subclass 1:20 that is capped
 # at bwmax (i.e., not capped). The 1:1 parent class exists only to
-# make the borrowing model work. All bandwidth is fairly shared,
-# subject to the restrictions of the class hierarchy: namely, that the
-# total bandwidth to non-exempt destinations should not exceed the
-# node bandwidth cap. The root slice has a higher priority (0) than
-# the others (1) and can thus request all of the bandwidth of that
-# subclass.
+# make the borrowing model work. All bandwidth above minimum
+# guarantees is fairly shared (in this example, slice 2 is guaranteed
+# at least 1mbit in addition to fair access to the rest), subject to
+# the restrictions of the class hierarchy: namely, that the total
+# bandwidth to non-exempt destinations should not exceed the node
+# bandwidth cap.
 #
-#                            1:
-#                            |
-#                       1:1 (1gbit)
-#              ______________|______________
-#             |                             |
-#    1:10 (8bit, 5mbit)           1:20 (8bit, 1gbit)
-#             |                             |
-# 1:1000 (8bit, 5mbit, 0),     1:2000 (8bit, 1gbit, 0),
-# 1:1001 (8bit, 5mbit, 1),     1:2001 (8bit, 1gbit, 1),
-# 1:1002 (8bit, 5mbit, 1),     1:2002 (8bit, 1gbit, 1),
-# ...                           ...
-# 1:1FFF (8bit, 5mbit, 1)      1:2FFF (8bit, 1gbit, 1)
+#                         1:
+#                         |
+#                    1:1 (1gbit)
+#           ______________|_____________
+#          |                            |
+#   1:10 (8bit, 5mbit)           1:20 (8bit, 1gbit)
+#          |                            |
+# 1:1000 (8bit, 5mbit),        1:2000 (8bit, 1gbit),
+# 1:1001 (8bit, 5mbit),        1:2001 (8bit, 1gbit),
+# 1:1002 (1mbit, 5mbit),       1:2002 (1mbit, 1gbit),
+# ...                          ...
+# 1:1FFF (8bit, 5mbit)         1:2FFF (8bit, 1gbit)
 #
 default_minor = 0x1000
 exempt_minor = 0x2000
 
 # root_xid is for the root context. The root context is exempt from
-# fair sharing in both the default and exempt subclasses..
+# fair sharing in both the default and exempt subclasses. The root
+# context gets 5 shares by default.
 root_xid = 0x0000
+root_share = 5
 
 # default_xid is for unclassifiable packets. Packets should not be
 # classified here very often. They can be if a slice's HTB classes are
-# deleted before its processes are.
+# deleted before its processes are. Each slice gets 1 share by
+# default.
 default_xid = 0x0FFF
+default_share = 1
 
 # See tc_util.c and http://physics.nist.gov/cuu/Units/binary.html. Be
 # warned that older versions of tc interpret "kbps", "mbps", "mbit",
@@ -242,13 +248,18 @@ def get_slice(xid_or_name):
     return None
 
 
-# Shortcut for running a tc command
-def tc(cmd):
+# Shortcut for running a command
+def run(cmd, input = None):
     try:
         if verbose:
-            sys.stderr.write("Executing: " + TC + " " + cmd + "\n")
-        fileobj = os.popen(TC + " " + cmd, "r")
-        output = fileobj.readlines()
+            sys.stderr.write("Executing: " + cmd + "\n")
+        if input is None:
+            fileobj = os.popen(cmd, "r")
+            output = fileobj.readlines()
+        else:
+            fileobj = os.popen(cmd, "w")
+            fileobj.write(input)
+            output = None
         if fileobj.close() is None:
             return output
     except Exception, e:
@@ -256,6 +267,11 @@ def tc(cmd):
     return None
 
 
+# Shortcut for running a tc command
+def tc(cmd):
+    return run(TC + " " + cmd)
+
+
 # (Re)initialize the bandwidth limits on this node
 def init(dev = dev, bwcap = None):
     if bwcap is None:
@@ -266,9 +282,6 @@ def init(dev = dev, bwcap = None):
         # Allow bwcap to be specified as a tc rate string
         bwcap = get_tc_rate(bwcap)
 
-    # Save current state (if any)
-    caps = get(dev = dev)
-
     # Delete root qdisc 1: if it exists. This will also automatically
     # delete any child classes.
     for line in tc("qdisc show dev %s" % dev):
@@ -303,20 +316,15 @@ def init(dev = dev, bwcap = None):
     # Set up the root class (and tell VNET what it is). Packets sent
     # by root end up here and are capped at the node bandwidth
     # cap.
-    on(root_xid, dev, prio = 0)
+    on(root_xid, dev, share = root_share)
     file("/proc/sys/vnet/root_class", "w").write("%d" % ((1 << 16) | default_minor | root_xid))
 
     # Set up the default class. Packets that fail classification end
     # up here.
-    on(default_xid, dev)
+    on(default_xid, dev, share = default_share)
 
-    # Reapply bandwidth caps. If the node bandwidth cap is now lower
-    # than it was before, "ceil" for each class will be lowered. If
-    # the node bandwidth cap is now higher than it was before, "ceil"
-    # for each class should be reapplied.
-    for (xid, share, minrate, maxrate) in caps:
-        if xid != root_xid and xid != default_xid:
-            on(xid, dev, share = share, minrate = minrate, maxrate = maxrate)
+    # Set up exemptions.
+    exempt_init()
 
 
 # Get the bandwidth limits for a particular slice xid as a tuple (xid,
@@ -369,7 +377,7 @@ def get(xid = None, dev = dev):
 
 
 # Apply specified bandwidth limit to the specified slice xid
-def on(xid, dev = dev, share = None, minrate = None, maxrate = None, prio = 1):
+def on(xid, dev = dev, share = None, minrate = None, maxrate = None):
     # Get defaults from current state if available
     cap = get(xid, dev)
     if cap is not None:
@@ -391,7 +399,7 @@ def on(xid, dev = dev, share = None, minrate = None, maxrate = None, prio = 1):
 
     # Set defaults
     if share is None:
-        share = 1
+        share = default_share
     if minrate is None:
         minrate = bwmin
     else:
@@ -404,16 +412,15 @@ def on(xid, dev = dev, share = None, minrate = None, maxrate = None, prio = 1):
     # Sanity checks
     if maxrate > bwcap:
         maxrate = bwcap
-
     if minrate > maxrate:
         minrate = maxrate
 
     # Set up subclasses for the slice
-    tc("class replace dev %s parent 1:10 classid 1:%x htb rate %dbit ceil %dbit quantum %d prio %d" % \
-       (dev, default_minor | xid, minrate, maxrate, share * quantum, prio))
+    tc("class replace dev %s parent 1:10 classid 1:%x htb rate %dbit ceil %dbit quantum %d" % \
+       (dev, default_minor | xid, minrate, maxrate, share * quantum))
 
-    tc("class replace dev %s parent 1:20 classid 1:%x htb rate %dbit ceil %dbit quantum %d prio %d" % \
-       (dev, exempt_minor | xid, minrate, bwmax, share * quantum, prio))
+    tc("class replace dev %s parent 1:20 classid 1:%x htb rate %dbit ceil %dbit quantum %d" % \
+       (dev, exempt_minor | xid, minrate, bwmax, share * quantum))
 
     # Attach a FIFO to each subclass, which helps to throttle back
     # processes that are sending faster than the token buckets can
@@ -429,12 +436,56 @@ def on(xid, dev = dev, share = None, minrate = None, maxrate = None, prio = 1):
 # are seen from this slice, they will be classified into the default
 # class 1:1FFF.
 def off(xid, dev = dev):
-    tc("class del dev %s classid 1:%x" % (dev, default_minor | xid))
-    tc("class del dev %s classid 1:%x" % (dev, exempt_minor | xid))
+    cap = get(xid, dev)
+    if cap is not None:
+        tc("class del dev %s classid 1:%x" % (dev, default_minor | xid))
+        tc("class del dev %s classid 1:%x" % (dev, exempt_minor | xid))
+
+
+def exempt_init():
+    # Who are we?
+    try:
+        node_id = int(file('/etc/planetlab/node_id').readline().strip())
+    except:
+        return False
+
+    api = plcapi.PLCAPI()
+
+    # All nodes that have access to Internet2
+    node_ids = []
+    for node_group in api.AnonAdmGetNodeGroups(api.auth):
+        if node_group['name'] == "Internet2":
+            node_ids += api.AnonAdmGetNodeGroupNodes(api.auth, node_group['nodegroup_id'])
+
+    # Remove duplicates
+    node_ids = list(Set(node_ids))
+
+    # Continue only if we ourselves have access to Internet2
+    if node_id not in node_ids:
+        return True
+
+    # Exempt the following destinations from the node bandwidth cap
+    node_ips = [node['ip'] for node in api.AnonAdmGetNodes(api.auth, node_ids, ['ip'])]
+
+    # Clean up
+    run("/sbin/iptables -t vnet -F POSTROUTING")
+    run("/sbin/ipset -X Internet2")
+
+    # Create a hashed IP set of all of these destinations
+    run("/sbin/modprobe ip_set_iphash")
+    lines = ["-N Internet2 iphash"]
+    lines += ["-A Internet2 " + ip for ip in node_ips]
+    lines += ["COMMIT"]
+    restore = "\n".join(lines) + "\n"
+    run("/sbin/ipset -R", restore)
+
+    # Add rule to match on destination IP set
+    run("/sbin/iptables -t vnet -A POSTROUTING -m set --set Internet2 dst -j CLASSIFY --set-class 1:%x" %
+        exempt_minor)
 
 
 def usage():
-    bwcap_description = format_tc_rate(bwmax)
+    bwcap_description = format_tc_rate(get_bwcap())
         
     print """
 Usage:
@@ -508,7 +559,7 @@ def main():
                 if slice is None:
                     # Orphaned (not associated with a slice) class
                     slice = "%d?" % xid
-                print "%s: share %d minrate %s maxrate %s" % \
+                print "%s %d %s %s" % \
                       (slice, share, format_tc_rate(minrate), format_tc_rate(maxrate))
 
         elif len(argv) >= 2: