scripts/demo.py

   1 #!/usr/local/bin/python
   2
   3 import httplib2
   4 import json
   5 import signal
   6 import time
   7
   8 from threading import Thread
   9
  10 from affinity_control import AffinityControl
  11 from subnet import SubnetControl
  12 from stats import Stats
  13
  14 # If True, SIG_INT has been captured
  15 global sigint
  16 sigint = False
  17
  18 # Handle SIG_INT
  19 def signal_handler(signal, frame):
  20     global sigint
  21     sigint = True
  22
  23 # Monitors statistics
  24 class WaypointMonitor(Thread):
  25
  26     def __init__(self, monitor_type, **kwargs):
  27         Thread.__init__(self)
  28         self.stat = Stats(monitor_type, **kwargs)
  29         self.stat_type = monitor_type
  30         self.waypoint_address = None
  31         print "Created waypoint monitor for %s" % self.stat
  32
  33     def set_waypoint(self, waypoint_ip):
  34         self.waypoint_address = waypoint_ip
  35         print "Registered waypoint for %s.  Any large flows will be redirected to %s." % (self.stat, waypoint_ip)
  36
  37     def set_large_flow_threshold(self, s):
  38         self.stat.set_large_flow_threshold(s)
  39         print "Set threshold for large flows to %d bytes" % s
  40
  41     def run(self):
  42         global sigint
  43         did_waypoint = False
  44         while not sigint:
  45             _, is_big = self.stat.refresh()
  46             if is_big and not did_waypoint:
  47                 print "Large flow detected (%d bytes)" % self.stat.get_bytes()
  48                 ac = AffinityControl()
  49                 # First AG: Sources sending data into this subnet
  50                 src_ag_name = "sources"
  51                 src_ips = self.stat.get_large_incoming_hosts()
  52                 if (self.waypoint_address in src_ips):
  53                     src_ips.remove(self.waypoint_address)
  54                 ac.add_affinity_group(src_ag_name, ips=src_ips)
  55                 # Second AG: This entity
  56                 dst_ag_name = "client"
  57                 if (self.stat_type == Stats.TYPE_PREFIX):
  58                     ac.add_affinity_group(dst_ag_name, subnet=self.stat.subnet)
  59                 elif (self.stat_type == Stats.TYPE_HOST):
  60                     pass
  61                 else:
  62                     print "type", self.stat_type, "not supported for redirection"
  63                 # AL: Between them
  64                 link_name = "inflows"
  65                 ac.add_affinity_link(link_name, src_ag_name, dst_ag_name)
  66                 ac.add_waypoint(link_name, self.waypoint_address)
  67                 ac.enable_waypoint(link_name)
  68                 did_waypoint = True
  69                 raw_input("[Press Enter to disable waypoint redirection] ")
  70                 ac.disable_waypoint(link_name)
  71             time.sleep(1)
  72
  73 def main():
  74
  75     # Default subnet is required for the host tracker to work.
  76     subnet_control = SubnetControl()
  77     subnet_control.add_subnet("defaultSubnet", "10.0.0.254/8")
  78
  79     m = WaypointMonitor(Stats.TYPE_PREFIX, subnet="10.0.0.0/31")
  80     m.set_waypoint("10.0.0.2")
  81     m.set_large_flow_threshold(2000) # 2000 bytes
  82     m.start()
  83
  84     # Register signal-handler to catch SIG_INT
  85     signal.signal(signal.SIGINT, signal_handler)
  86     signal.pause()
  87
  88     # join() won't return until SIG_INT has been captured
  89     m.join()
  90
  91 if __name__ == "__main__":
  92     main()