Changeset 9988 in genesis for tools/get-link-traffic.py

Timestamp:

Feb 18, 2012, 8:56:43 PM (13 years ago)

Author:

rick

Message:

Use multiprocessing to speedup discovery and running by 10x

File:

• tools/get-link-traffic.py (modified) (4 diffs)

tools/get-link-traffic.py

@@ -10 +10 @@
 import sys
 import yaml
+from multiprocessing import Process, Manager, Pool, Queue, freeze_support
+
 logging.basicConfig(level=logging.DEBUG)
 logger = logging.getLogger()
 
+
+
 DATASTORE='store.yaml'
-
-try:
-  store = yaml.load(open(DATASTORE,'r'))
-except IOError:
-  store = { 'snmp' : {}, 'traffic' : {}, 'uptime' : {}}
-  pass
 
 class ConnectError(Exception):
   pass
 
-def find_right_snmp_ip(data):
-  for k,v in data.iteritems():
-    if k.startswith('iface_'):
-      ip = v['ip'].split('/')[0]
-      logger.info("Trying ip %s", ip)
-      p = subprocess.Popen("snmpwalk -t 1 -r 1 -Oq -c public -v2c %s uptime" % ip,
-        shell=True,
-        stdout=subprocess.PIPE, stderr=subprocess.PIPE)
-      (stdout, stderr) = p.communicate()
-      if p.returncode == 0:
-        return ip
-      else:
-        logger.error(stderr.strip())
-  raise None
-
-def get_snmp_stats(ip,target):
+def get_snmp_stats(logger, ip,target):
   p = subprocess.Popen("snmpwalk -t 1 -r 1 -Oq -c public -v2c %s %s" % (ip, target), 
     shell=True,
@@ -54 +37 @@
   return r
 
-def get_snmp_value(ip, target):
+def get_snmp_value(logger, ip, target):
   p = subprocess.Popen("snmpget -t 1 -r 1 -Ovt -c public -v2c %s %s" % (ip, target), 
     shell=True,
@@ -64 +47 @@
   return stdout.strip()
 
+def find_right_snmp_ip(logger, data):
+  for k,v in data.iteritems():
+    if k.startswith('iface_'):
+      ip = v['ip'].split('/')[0]
+      logger.info("Trying ip %s", ip)
+      try:
+        uptime = get_snmp_value(logger, ip, 'DISMAN-EVENT-MIB::sysUpTimeInstance')
+        return ip
+      except ConnectError:
+        pass
+  return None
+
     
 
@@ -71 +66 @@
   ff = ''
 
-try:
+def process_file(logger, store, nf, rescan):
+  data = yaml.load(open(nf,'r'))
+  nodename = data['nodename']
+
+  if store['snmp'].has_key(nodename):
+    ip = store['snmp'][nodename]
+    if not ip and rescan:
+      logger.info("Re-scanning for new valid IP")
+      ip = find_right_snmp_ip(logger, data)
+      store['snmp'][nodename] = ip
+    else:
+      try:
+        uptime = get_snmp_value(logger, ip, 'DISMAN-EVENT-MIB::sysUpTimeInstance')
+      except ConnectError:
+        logger.info("Re-scanning for new valid IP")
+        ip = find_right_snmp_ip(logger, data)
+        store['snmp'][nodename] = ip
+      
+  else:
+    logger.info("Running discovery for %s", nodename)
+    ip = find_right_snmp_ip(logger, data)
+    store['snmp'][nodename] = ip
+
+  if ip == None:
+    logger.error("No valid ip found for node %s", nodename)
+    return
+  
+  logger.info("Processing %s via %s", nodename, ip)
+  target = 'IF-MIB::ifDescr'
+  
+  try:
+    iface = get_snmp_stats(logger, ip, 'IF-MIB::ifDescr')
+    ifout = get_snmp_stats(logger, ip, 'IF-MIB::ifOutOctets')
+    ifin = get_snmp_stats(logger, ip, 'IF-MIB::ifInOctets')
+
+    uptime = get_snmp_value(logger, ip, 'DISMAN-EVENT-MIB::sysUpTimeInstance')
+    store['uptime'][nodename] = int(uptime)
+    
+    traffic = {}
+    for i,f in iface.iteritems():
+      traffic[f] = (int(ifin[i]), int(ifout[i]))
+    store['traffic'][nodename] = traffic 
+  except ConnectError:
+    logger.error("Unable to get all data")
+    pass
+
+def worker(i, input, m_store):
+  logger = logging.getLogger('Worker%s' % i)
+  logger.info("Worker")
+  for (nf, rescan) in iter(input.get, 'STOP'):
+    process_file(logger, m_store, nf, rescan)
+  logger.info("END")
+
+if __name__ == '__main__':
+  freeze_support()
+  task_queue = Queue()
+  manager = Manager()
+
+  try:
+    store = yaml.load(open(DATASTORE,'r'))
+  except IOError:
+    store = { 'snmp' : {}, 'traffic' : {}, 'uptime' : {}}
+    pass
+  m_store = manager.dict(store)
+
+  NUMBER_OF_PROCESSES = 10
+  RESCAN = True
+  plist = {}
+  for i in range(NUMBER_OF_PROCESSES):
+    plist[i] = Process(target=worker, args=(i, task_queue,m_store))
+    plist[i].start()
+
   for nf in sorted(glob.glob('nodes/*%s*/wleiden.yaml' % ff)):
-    data = yaml.load(open(nf,'r'))
-    nodename = data['nodename']
-    
-    if store['snmp'].has_key(nodename):
-      ip = store['snmp'][nodename]
-    else:
-      logger.info("Running discovery for %s", nodename)
-      ip = find_right_snmp_ip(data)
-      store['snmp'][nodename] = ip
+    task_queue.put((nf, RESCAN))
+ 
+  for i in range(NUMBER_OF_PROCESSES):
+    task_queue.put('STOP')
 
-    if ip == None:
-      logger.error("No valid ip found for node %s", nodename)
-      continue
-    
-    logger.info("Processing %s via %s", nodename, ip)
-    target = 'IF-MIB::ifDescr'
-    
-    try:
-      iface = get_snmp_stats(ip, 'IF-MIB::ifDescr')
-      ifout = get_snmp_stats(ip, 'IF-MIB::ifOutOctets')
-      ifin = get_snmp_stats(ip, 'IF-MIB::ifInOctets')
-
-      uptime = get_snmp_value(ip, 'DISMAN-EVENT-MIB::sysUpTimeInstance')
-      store['uptime'][nodename] = int(uptime)
-      
-      traffic = {}
-      for i,f in iface.iteritems():
-        traffic[f] = (int(ifin[i]), int(ifout[i]))
-      store['traffic'][nodename] = traffic 
-    except ConnectError:
-      logger.error("Unable to get all data")
-      pass
-except KeyboardInterrupt:
-  pass
-
-yaml.dump(store,open(DATASTORE,'w'))
+  for i in range(NUMBER_OF_PROCESSES):
+    plist[i].join()
 
 
+  yaml.dump(store,open(DATASTORE,'w'))
+
+
+