remove some PL-specific details
[util-vserver.git] / python / bwlimit.py
1 #!/usr/bin/python
2
3 # Bandwidth limit module for PlanetLab nodes. The intent is to use the
4 # Hierarchical Token Bucket (HTB) queueing discipline (qdisc) to allow
5 # slices to fairly share access to available node bandwidth. We
6 # currently define three classes of "available node bandwidth":
7 #
8 # 1. Available hardware bandwidth (bwmax): The maximum rate of the
9 # hardware.
10 #
11 # 2. Available capped bandwidth (bwcap): The maximum rate allowed to
12 # non-exempt destinations. By default, equal to bwmax, but may be
13 # lowered by PIs.
14 #
15 # 3. Available uncapped ("exempt") bandwidth: The difference between
16 # bwmax and what is currently being used of bwcap, or the maximum rate
17 # allowed to destinations exempt from caps (e.g., Internet2).
18 #
19 # All three classes of bandwidth are fairly shared according to the
20 # notion of "shares". For instance, if the node is capped at 5 Mbps,
21 # there are N slices, and each slice has 1 share, then each slice
22 # should get at least 5/N Mbps of bandwidth. How HTB is implemented
23 # makes this statement a little too simplistic. What it really means
24 # is that during any single time period, only a certain number of
25 # bytes can be sent onto the wire. Each slice is guaranteed that at
26 # least some small number of its bytes will be sent. Whatever is left
27 # over from the budget, is split in proportion to the number of shares
28 # each slice has.
29 #
30 # Even if the node is not capped at a particular limit (bwcap ==
31 # bwmax), this module enforces fair share access to bwmax. Also, if
32 # the node is capped at a particular limit, rules may optionally be
33 # defined that classify certain packets into the "exempt" class. This
34 # class receives whatever bandwidth is leftover between bwcap and
35 # bwmax; slices fairly share this bandwidth as well.
36 #
37 # The root context is exempt from sharing and can send as much as it
38 # needs to.
39 #
40 # Some relevant URLs:
41 #
42 # 1. http://lartc.org/howto               for how to use tc
43 # 2. http://luxik.cdi.cz/~devik/qos/htb/  for info on HTB
44 #
45 # Andy Bavier <acb@cs.princeton.edu>
46 # Mark Huang <mlhuang@cs.princeton.edu>
47 # Copyright (C) 2006 The Trustees of Princeton University
48 #
49 # $Id: bwlimit.py,v 1.9 2006/03/01 22:37:24 mlhuang Exp $
50 #
51
52 import sys, os, re, getopt
53 from sets import Set
54 import pwd
55
56
57 # Where the tc binary lives
58 TC = "/sbin/tc"
59
60 # Default interface
61 dev = "eth0"
62
63 # Verbosity level
64 verbose = 0
65
66 # bwmin should be small enough that it can be considered negligibly
67 # slow compared to the hardware. 8 bits/second appears to be the
68 # smallest value supported by tc.
69 bwmin = 8
70
71 # bwmax should be large enough that it can be considered at least as
72 # fast as the hardware.
73 bwmax = 1000*1000*1000
74
75 # quantum is the maximum number of bytes that can be borrowed by a
76 # share (or slice, if each slice gets 1 share) in one time period
77 # (with HZ=1000, 1 ms). If multiple slices are competing for bandwidth
78 # above their guarantees, and each is attempting to borrow up to the
79 # node bandwidth cap, quantums control how the excess bandwidth is
80 # distributed. Slices with 2 shares will borrow twice the amount in
81 # one time period as slices with 1 share, so averaged over time, they
82 # will get twice as much of the excess bandwidth. The value should be
83 # as small as possible and at least 1 MTU. By default, it would be
84 # calculated as bwmin/10, but since we use such small a value for
85 # bwmin, it's better to just set it to a value safely above 1 Ethernet
86 # MTU.
87 quantum = 1600
88
89 # cburst is the maximum number of bytes that can be burst onto the
90 # wire in one time period (with HZ=1000, 1 ms). If multiple slices
91 # have data queued for transmission, cbursts control how long each
92 # slice can have the wire for. If not specified, it is set to the
93 # smallest possible value that would enable the slice's "ceil" rate
94 # (usually the node bandwidth cap), to be reached if a slice was able
95 # to borrow enough bandwidth to do so. For now, it's unclear how or if
96 # to relate this to the notion of shares, so just let tc set the
97 # default.
98 cburst = None
99
100 # There is another parameter that controls how bandwidth is allocated
101 # between slices on nodes that is outside the scope of HTB. We enforce
102 # a 16 GByte/day total limit on each slice, which works out to about
103 # 1.5mbit. If a slice exceeds this byte limit before the day finishes,
104 # it is capped at (i.e., its "ceil" rate is set to) the smaller of the
105 # node bandwidth cap or 1.5mbit. pl_mom is in charge of enforcing this
106 # rule and executes this script to override "ceil".
107
108 # We support multiple bandwidth limits, by reserving the top nibble of
109 # the minor classid to be the "subclassid". Theoretically, we could
110 # support up to 15 subclasses, but for now, we only define two: the
111 # "default" subclass 1:10 that is capped at the node bandwidth cap (in
112 # this example, 5mbit) and the "exempt" subclass 1:20 that is capped
113 # at bwmax (i.e., not capped). The 1:1 parent class exists only to
114 # make the borrowing model work. All bandwidth above minimum
115 # guarantees is fairly shared (in this example, slice 2 is guaranteed
116 # at least 1mbit in addition to fair access to the rest), subject to
117 # the restrictions of the class hierarchy: namely, that the total
118 # bandwidth to non-exempt destinations should not exceed the node
119 # bandwidth cap.
120 #
121 #                         1:
122 #                         |
123 #                    1:1 (1gbit)
124 #           ______________|_____________
125 #          |                            |
126 #   1:10 (8bit, 5mbit)           1:20 (8bit, 1gbit)
127 #          |                            |
128 # 1:1000 (8bit, 5mbit),        1:2000 (8bit, 1gbit),
129 # 1:1001 (8bit, 5mbit),        1:2001 (8bit, 1gbit),
130 # 1:1002 (1mbit, 5mbit),       1:2002 (1mbit, 1gbit),
131 # ...                          ...
132 # 1:1FFF (8bit, 5mbit)         1:2FFF (8bit, 1gbit)
133 #
134 default_minor = 0x1000
135 exempt_minor = 0x2000
136
137 # root_xid is for the root context. The root context is exempt from
138 # fair sharing in both the default and exempt subclasses. The root
139 # context gets 5 shares by default.
140 root_xid = 0x0000
141 root_share = 5
142
143 # default_xid is for unclassifiable packets. Packets should not be
144 # classified here very often. They can be if a slice's HTB classes are
145 # deleted before its processes are. Each slice gets 1 share by
146 # default.
147 default_xid = 0x0FFF
148 default_share = 1
149
150 # See tc_util.c and http://physics.nist.gov/cuu/Units/binary.html. Be
151 # warned that older versions of tc interpret "kbps", "mbps", "mbit",
152 # and "kbit" to mean (in this system) "kibps", "mibps", "mibit", and
153 # "kibit" and that if an older version is installed, all rates will
154 # be off by a small fraction.
155 suffixes = {
156     "":         1,
157     "bit":      1,
158     "kibit":    1024,
159     "kbit":     1000,
160     "mibit":    1024*1024,
161     "mbit":     1000000,
162     "gibit":    1024*1024*1024,
163     "gbit":     1000000000,
164     "tibit":    1024*1024*1024*1024,
165     "tbit":     1000000000000,
166     "bps":      8,
167     "kibps":    8*1024,
168     "kbps":     8000,
169     "mibps":    8*1024*1024,
170     "mbps":     8000000,
171     "gibps":    8*1024*1024*1024,
172     "gbps":     8000000000,
173     "tibps":    8*1024*1024*1024*1024,
174     "tbps":     8000000000000
175 }
176
177
178 # Parses an integer or a tc rate string (e.g., 1.5mbit) into bits/second
179 def get_tc_rate(s):
180     if type(s) == int:
181         return s
182     m = re.match(r"([0-9.]+)(\D*)", s)
183     if m is None:
184         return -1
185     suffix = m.group(2).lower()
186     if suffixes.has_key(suffix):
187         return int(float(m.group(1)) * suffixes[suffix])
188     else:
189         return -1
190
191
192 # Prints a tc rate string
193 def format_tc_rate(rate):
194     if rate >= 1000000:
195         return "%.0fmbit" % (rate / 1000000.)
196     elif rate >= 1000:
197         return "%.0fkbit" % (rate / 1000.)
198     else:
199         return "%.0fbit" % rate
200
201
202 # Parse /etc/planetlab/bwcap (or equivalent)
203 def read_bwcap(bwcap_file):
204     bwcap = bwmax
205     try:
206         fp = open(bwcap_file, "r")
207         line = fp.readline().strip()
208         if line:
209             bwcap = get_tc_rate(line)
210     except:
211         pass
212     if bwcap == -1:
213         bwcap = bwmax
214     return bwcap
215
216
217 # Get current (live) value of bwcap
218 def get_bwcap(dev = dev):
219
220     state = tc("-d class show dev %s" % dev)
221     base_re = re.compile(r"class htb 1:10 parent 1:1 .*ceil ([^ ]+) .*")
222     base_classes = filter(None, map(base_re.match, state))
223     if not base_classes:
224         return -1
225     if len(base_classes) > 1:
226         raise Exception, "unable to get current bwcap"
227     return get_tc_rate(base_classes[0].group(1))
228
229
230 # Get slice xid (500) from slice name ("500" or "princeton_mlh") or
231 # slice name ("princeton_mlh") from slice xid (500).
232 def get_slice(xid_or_name):
233
234     if xid_or_name == root_xid:
235         return "root"
236     if xid_or_name == default_xid:
237         return "default"
238     if isinstance(xid_or_name, (int, long)):
239         try:
240             return pwd.getpwuid(xid_or_name).pw_name
241         except KeyError:
242             pass
243     else:
244         try:
245             try:
246                 return int(xid_or_name)
247             except ValueError:
248                 pass
249             return pwd.getpwnam(xid_or_name).pw_uid
250         except KeyError:
251             pass
252
253     return None
254
255
256 # Shortcut for running a command
257 def run(cmd, input = None):
258     try:
259         if verbose:
260             sys.stderr.write("Executing: " + cmd + "\n")
261         if input is None:
262             fileobj = os.popen(cmd, "r")
263             output = fileobj.readlines()
264         else:
265             fileobj = os.popen(cmd, "w")
266             fileobj.write(input)
267             output = None
268         if fileobj.close() is None:
269             return output
270     except Exception, e:
271         pass
272     return None
273
274
275 # Shortcut for running a tc command
276 def tc(cmd):
277     return run(TC + " " + cmd)
278
279
280 # (Re)initialize the bandwidth limits on this node
281 def init(dev, bwcap):
282
283     # Delete root qdisc 1: if it exists. This will also automatically
284     # delete any child classes.
285     for line in tc("qdisc show dev %s" % dev):
286         # Search for the root qdisc 1:
287         m = re.match(r"qdisc htb 1:", line)
288         if m is not None:
289             tc("qdisc del dev %s root handle 1:" % dev)
290             break
291
292     # Initialize HTB. The "default" clause specifies that if a packet
293     # fails classification, it should go into the class with handle
294     # 1FFF.
295     tc("qdisc add dev %s root handle 1: htb default %x" % \
296        (dev, default_minor | default_xid))
297
298     # Set up a parent class from which all subclasses borrow.
299     tc("class add dev %s parent 1: classid 1:1 htb rate %dbit" % \
300        (dev, bwmax))
301
302     # Set up a subclass that represents the node bandwidth cap. We
303     # allow each slice to borrow up to this rate, so it is also
304     # usually the "ceil" rate for each slice.
305     tc("class add dev %s parent 1:1 classid 1:10 htb rate %dbit ceil %dbit" % \
306        (dev, bwmin, bwcap))
307
308     # Set up a subclass that represents "exemption" from the node
309     # bandwidth cap. Once the node bandwidth cap is reached, bandwidth
310     # to exempt destinations can still be fairly shared up to bwmax.
311     tc("class add dev %s parent 1:1 classid 1:20 htb rate %dbit ceil %dbit" % \
312        (dev, bwmin, bwmax))
313
314     # Set up the root class (and tell VNET what it is). Packets sent
315     # by root end up here and are capped at the node bandwidth
316     # cap.
317     on(root_xid, dev, share = root_share)
318     file("/proc/sys/vnet/root_class", "w").write("%d" % ((1 << 16) | default_minor | root_xid))
319
320     # Set up the default class. Packets that fail classification end
321     # up here.
322     on(default_xid, dev, share = default_share)
323
324
325 # Get the bandwidth limits for a particular slice xid as a tuple (xid,
326 # share, minrate, maxrate), or all classes as a list of tuples.
327 def get(xid = None, dev = dev):
328     if xid is None:
329         ret = []
330     else:
331         ret = None
332
333     # class htb 1:1002 parent 1:10 leaf 81b3: prio 1 rate 8bit ceil 5000Kbit burst 1600b cburst 4Kb
334     for line in tc("-d class show dev %s" % dev):
335         # Search for child classes of 1:10
336         m = re.match(r"class htb 1:([0-9a-f]+) parent 1:10", line)
337         if m is None:
338             continue
339
340         # If we are looking for a particular class
341         classid = int(m.group(1), 16) & default_xid
342         if xid is not None and xid != classid:
343             continue
344
345         # Parse share
346         share = 1
347         m = re.search(r"quantum (\d+)", line)
348         if m is not None:
349             share = int(m.group(1)) / quantum
350
351         # Parse minrate
352         minrate = bwmin
353         m = re.search(r"rate (\w+)", line)
354         if m is not None:
355             minrate = get_tc_rate(m.group(1))
356
357         # Parse maxrate 
358         maxrate = bwmax
359         m = re.search(r"ceil (\w+)", line)
360         if m is not None:
361             maxrate = get_tc_rate(m.group(1))
362
363         if xid is None:
364             # Return a list of parameters
365             ret.append((classid, share, minrate, maxrate))
366         else:
367             # Return the parameters for this class
368             ret = (classid, share, minrate, maxrate)
369             break
370
371     return ret
372
373
374 # Apply specified bandwidth limit to the specified slice xid
375 def on(xid, dev = dev, share = None, minrate = None, maxrate = None):
376     # Get defaults from current state if available
377     cap = get(xid, dev)
378     if cap is not None:
379         if share is None:
380             share = cap[1]
381         if minrate is None:
382             minrate = cap[2]
383         if maxrate is None:
384             maxrate = cap[3]
385
386     # Figure out what the current node bandwidth cap is
387     bwcap = bwmax
388     for line in tc("-d class show dev %s" % dev):
389         # Search for 1:10
390         m = re.match(r"class htb 1:10.*ceil (\w+)", line)
391         if m is not None:
392             bwcap = get_tc_rate(m.group(1))
393             break
394
395     # Set defaults
396     if share is None:
397         share = default_share
398     if minrate is None:
399         minrate = bwmin
400     else:
401         minrate = get_tc_rate(minrate)
402     if maxrate is None:
403         maxrate = bwcap
404     else:
405         maxrate = get_tc_rate(maxrate)
406
407     # Sanity checks
408     if maxrate > bwcap:
409         maxrate = bwcap
410     if minrate > maxrate:
411         minrate = maxrate
412
413     # Set up subclasses for the slice
414     tc("class replace dev %s parent 1:10 classid 1:%x htb rate %dbit ceil %dbit quantum %d" % \
415        (dev, default_minor | xid, minrate, maxrate, share * quantum))
416
417     tc("class replace dev %s parent 1:20 classid 1:%x htb rate %dbit ceil %dbit quantum %d" % \
418        (dev, exempt_minor | xid, minrate, bwmax, share * quantum))
419
420     # Attach a FIFO to each subclass, which helps to throttle back
421     # processes that are sending faster than the token buckets can
422     # support.
423     tc("qdisc replace dev %s parent 1:%x handle %x pfifo" % \
424        (dev, default_minor | xid, default_minor | xid))
425
426     tc("qdisc replace dev %s parent 1:%x handle %x pfifo" % \
427        (dev, exempt_minor | xid, exempt_minor | xid))
428
429
430 # Remove class associated with specified slice xid. If further packets
431 # are seen from this slice, they will be classified into the default
432 # class 1:1FFF.
433 def off(xid, dev = dev):
434     cap = get(xid, dev)
435     if cap is not None:
436         tc("class del dev %s classid 1:%x" % (dev, default_minor | xid))
437         tc("class del dev %s classid 1:%x" % (dev, exempt_minor | xid))
438
439
440 def exempt_init(group_name, node_ips):
441
442     # Clean up
443     iptables = "/sbin/iptables -t vnet %s POSTROUTING" 
444     run(iptables % "-F")
445     run("/sbin/ipset -X " + group_name)
446
447     # Create a hashed IP set of all of these destinations
448     run("/sbin/modprobe ip_set_iphash")
449     lines = ["-N %s iphash" % group_name]
450     add_cmd = "-A %s " % group_name
451     lines += [(add_cmd + ip) for ip in node_ips]
452     lines += ["COMMIT"]
453     restore = "\n".join(lines) + "\n"
454     run("/sbin/ipset -R", restore)
455
456     # Add rule to match on destination IP set
457     run((iptables + " -m set --set %s dst -j CLASSIFY --set-class 1:%x") %
458         ("-A", group_name, exempt_minor))
459
460
461 def usage():
462     bwcap_description = format_tc_rate(get_bwcap())
463         
464     print """
465 Usage:
466
467 %s [OPTION]... [COMMAND] [ARGUMENT]...
468
469 Options:
470         -d device       Network interface (default: %s)
471         -r rate         Node bandwidth cap (default: %s)
472         -q quantum      Share multiplier (default: %d bytes)
473         -h              This message
474
475 Commands:
476         init
477                 (Re)initialize bandwidth caps.
478         on slice [share] [minrate] [maxrate]
479                 Set bandwidth cap for the specified slice
480         off slice
481                 Remove bandwidth caps for the specified slice
482         get
483                 Get all bandwidth caps
484         get slice
485                 Get bandwidth caps for the specified slice
486         getcap slice
487                 Get maxrate for the specified slice
488         setcap slice maxrate
489                 Set maxrate for the specified slice
490 """ % (sys.argv[0], dev, bwcap_description, quantum)
491     sys.exit(1)
492     
493
494 def main():
495     global dev, quantum, verbose
496
497     # Defaults
498     bwcap = get_bwcap()
499
500     (opts, argv) = getopt.getopt(sys.argv[1:], "f:d:r:g:q:vh")
501     for (opt, optval) in opts:
502         if opt == '-d':
503             dev = optval
504         elif opt == '-r':
505             bwcap = get_tc_rate(optval)
506         elif opt == '-q':
507             quantum = int(optval)
508         elif opt == '-v':
509             verbose += 1
510         elif opt == '-h':
511             usage()
512
513     if len(argv):
514         if argv[0] == "init" or (argv[0] == "on" and len(argv) == 1):
515             # (Re)initialize
516             init(dev, get_tc_rate(bwcap))
517
518         elif argv[0] == "get" or argv[0] == "show":
519             # Show
520             if len(argv) >= 2:
521                 # Show a particular slice
522                 xid = get_slice(argv[1])
523                 if xid is None:
524                     sys.stderr.write("Error: Invalid slice name or context '%s'\n" % argv[1])
525                     usage()
526                 caps = [get(xid, dev)]
527             else:
528                 # Show all slices
529                 caps = get(None, dev)
530
531             for (xid, share, minrate, maxrate) in caps:
532                 slice = get_slice(xid)
533                 if slice is None:
534                     # Orphaned (not associated with a slice) class
535                     slice = "%d?" % xid
536                 print "%s %d %s %s" % \
537                       (slice, share, format_tc_rate(minrate), format_tc_rate(maxrate))
538
539         elif len(argv) >= 2:
540             # slice, ...
541             xid = get_slice(argv[1])
542             if xid is None:
543                 sys.stderr.write("Error: Invalid slice name or context '%s'\n" % argv[1])
544                 usage()
545
546             if argv[0] == "on" or argv[0] == "add" or argv[0] == "replace":
547                 # Enable cap
548                 args = []
549                 if len(argv) >= 3:
550                     # ... share, minrate, maxrate
551                     casts = [int, get_tc_rate, get_tc_rate]
552                     for i, arg in enumerate(argv[2:]):
553                         if i >= len(casts):
554                             break
555                         args.append(casts[i](arg))
556                 on(xid, dev, *args)
557
558             elif argv[0] == "off" or argv[0] == "del":
559                 # Disable cap
560                 off(xid, dev)
561
562             # Backward compatibility with old resman script
563             elif argv[0] == "getcap":
564                 # Get maxrate
565                 cap = get(xid, dev)
566                 if cap is not None:
567                     (xid, share, minrate, maxrate) = cap
568                     print format_tc_rate(maxrate)
569
570             # Backward compatibility with old resman script
571             elif argv[0] == "setcap":
572                 if len(argv) >= 3:
573                     # Set maxrate
574                     on(xid, dev, maxrate = get_tc_rate(argv[2]))
575                 else:
576                     usage()
577
578             else:
579                 usage()
580
581         else:
582             usage()
583
584
585 if __name__ == '__main__':
586     main()