rewrite tasks, add virtioserial support

agentcelery.py

 # from twisted.internet.defer import Deferred
 from twisted.internet import reactor  # threads
+from celery.result import TimeoutError
 from celery import Celery
 from kombu import Queue, Exchange
 from os import getenv
 from socket import gethostname
+from threading import Event
 import logging
 
 logger = logging.getLogger()
@@ -21,98 +23,100 @@ celery.conf.update(CELERY_CACHE_BACKEND=CACHE_URI,
                                         routing_key='agent'), ))
 
 
+def send_command(vm, command, *args, **kwargs):
+    uuid = kwargs.get('uuid', None)
+    timeout = kwargs.get('timeout', 10)
+    if uuid:
+        event = Event()
+        reactor.running_tasks[vm][uuid] = event
+        reactor.ended_tasks[vm][uuid] = None
+
+    for conn in reactor.connections[vm]:
+        logger.info('%s(%s, %s)', command, vm,
+                    ', '.join(map(lambda x: str(x)[:100], kwargs.values())))
+        conn.send_command(command=command, args=kwargs)
+
+    if uuid:
+        success = event.wait(timeout)
+        retval = reactor.ended_tasks[vm][uuid]
+
+        del reactor.ended_tasks[vm][uuid]
+        del reactor.running_tasks[vm][uuid]
+
+        if not success:
+            raise TimeoutError()
+
+        return retval
+
+
 @celery.task(name='agent.change_password')
 def change_password(vm, password):
-    reactor.connections[vm].send_command(command='change_password',
-                                         args={'password':
-                                               password})
-    logger.debug('change_password(%s,%s)', vm, password)
+    send_command(vm, command='change_password', password=password)
 
 
 @celery.task(name='agent.set_hostname')
 def set_hostname(vm, hostname):
-    reactor.connections[vm].send_command(command='set_hostname',
-                                         args={'hostname':
-                                               hostname})
-    logger.debug('set_hostname(%s,%s)', vm, hostname)
+    send_command(vm, command='set_hostname', hostname=hostname)
 
 
 @celery.task(name='agent.restart_networking')
 def restart_networking(vm):
-    reactor.connections[vm].send_command(command='restart_networking',
-                                         args={})
-    logger.debug('restart_networking(%s)', vm)
+    send_command(vm, command='restart_networking')
 
 
 @celery.task(name='agent.set_time')
 def set_time(vm, time):
-    reactor.connections[vm].send_command(command='set_time',
-                                         args={'time': time})
-    logger.debug('set_time(%s,%s)', vm, time)
+    send_command(vm, command='set_time', time=time)
 
 
 @celery.task(name='agent.mount_store')
 def mount_store(vm, host, username, password):
-    reactor.connections[vm].send_command(command='mount_store',
-                                         args={'host': host,
-                                               'username': username,
-                                               'password': password})
-    logger.debug('mount_store(%s,%s,%s)', vm, host, username)
+    send_command(vm, command='mount_store', host=host,
+                 username=username, password=password)
 
 
 @celery.task(name='agent.cleanup')
 def cleanup(vm):
-    reactor.connections[vm].send_command(command='cleanup', args={})
-    logger.debug('cleanup(%s)', vm)
+    send_command(vm, command='cleanup')
 
 
 @celery.task(name='agent.start_access_server')
 def start_access_server(vm):
-    reactor.connections[vm].send_command(
-        command='start_access_server', args={})
-    logger.debug('start_access_server(%s)', vm)
+    send_command(vm, command='start_access_server')
 
 
 @celery.task(name='agent.update')
-def update(vm, data):
-    logger.debug('update(%s)', vm)
-    return reactor.connections[vm].send_command(
-        command='update', args={'data': data}, uuid=update.request.id)
+def update(vm, data, executable=None):
+    kwargs = {'command': 'update', 'data': data, 'uuid': update.request.id}
+    if executable is not None:
+        kwargs['executable'] = executable
+    return send_command(vm, **kwargs)
 
 
 @celery.task(name='agent.add_keys')
 def add_keys(vm, keys):
-    logger.debug('add_keys(%s, %s)', vm, keys)
-    reactor.connections[vm].send_command(
-        command='add_keys', args={'keys': keys})
+    send_command(vm, command='add_keys', keys=keys)
 
 
 @celery.task(name='agent.del_keys')
 def del_keys(vm, keys):
-    logger.debug('del_keys(%s, %s)', vm, keys)
-    reactor.connections[vm].send_command(
-        command='del_keys', args={'keys': keys})
+    send_command(vm, command='del_keys', keys=keys)
 
 
 @celery.task(name='agent.get_keys')
 def get_keys(vm):
-    logger.debug('get_keys(%s)', vm)
-    return reactor.connections[vm].send_command(
-        command='get_keys', args={}, uuid=get_keys.request.id)
+    return send_command(vm, command='get_keys')
 
 
 @celery.task(name='agent.send_expiration')
 def send_expiration(vm, url):
-    logger.debug('send_expiration(%s, %s)', vm, url)
-    return reactor.connections[vm].send_command(
-        command='send_expiration', args={'url': url})
+    return send_command(vm, command='send_expiration',
+                        url=url)
 
 
 @celery.task(name='agent.change_ip')
 def change_ip(vm, interfaces, dns):
-    logger.debug('change_ip(%s, %s, %s)', vm, interfaces, dns)
-    return reactor.connections[vm].send_command(
-        command='change_ip', args={'interfaces': interfaces, 'dns': dns})
+    send_command(vm, command='change_ip', interfaces=interfaces, dns=dns)
 
 
 @celery.task(name='vm.tasks.local_agent_tasks.renew')

agentdriver.py

@@ -25,6 +25,8 @@ Worker.install_platform_tweaks = install_platform_tweaks
 
 
 def reactor_started():
+    reactor.running_tasks = {}
+    reactor.ended_tasks = {}
     for f in listdir(SOCKET_DIR):
         f = path.join(SOCKET_DIR, f)
         inotify_handler(None, filepath.FilePath(f), None)

protocol.py

@@ -5,9 +5,7 @@ import pickle
 import logging
 import time
 import struct
-from threading import Event
 from os import getenv
-from celery.result import TimeoutError
 
 from utils import SerialLineReceiverBase
 
@@ -42,10 +40,11 @@ class GraphiteClientFactory(protocol.ClientFactory):
 
 
 def inotify_handler(self, file, mask):
-    vm = file.basename()
-    logger.info('inotify: %s' % vm)
-    if vm in reactor.connections:
-        return
+    vm = file.basename().replace('vio-', '')
+    logger.info('inotify: %s (%s)', vm, file.path)
+    for conn in reactor.connections.get(vm, []):
+        if file.path == conn.transport.addr:
+            return
     serial = SerialLineReceiverFactory(vm)
     logger.info("connecting to %s (%s)", vm, file.path)
     reactor.connectUNIX(file.path, serial)
@@ -76,47 +75,29 @@ class SerialLineReceiver(SerialLineReceiverBase):
                                args=args)
 
     def handle_response(self, response, args):
+        vm = self.factory.vm
         if response == 'status':
             self.send_to_graphite(args)
         else:
             uuid = args.get('uuid', None)
             if not uuid:
                 return
-            event = self.factory.running_tasks.get(uuid, None)
+            event = reactor.running_tasks[vm].get(uuid, None)
             if event:
-                self.factory.ended_tasks[uuid] = args
+                reactor.ended_tasks[vm][uuid] = args
                 event.set()
 
     def connectionMade(self):
-        logger.info("connected to %s" % self.factory.vm)
-        reactor.connections[self.factory.vm] = self
+        logger.info("connected to %s (%s)", self.factory.vm,
+                    self.transport.addr)
+        if self.factory.vm not in reactor.connections:
+            reactor.connections[self.factory.vm] = set()
+        reactor.connections[self.factory.vm].add(self)
 
     def connectionLost(self, reason):
-        logger.info("disconnected from %s" % self.factory.vm)
-        del reactor.connections[self.factory.vm]
-
-    def send_command(self, command, args, timeout=10.0, uuid=None):
-        if not uuid:
-            super(SerialLineReceiver, self).send_command(command, args)
-            return
-
-        event = Event()
-        args['uuid'] = uuid
-        self.factory.running_tasks[uuid] = event
-        self.factory.ended_tasks[uuid] = None
-
-        super(SerialLineReceiver, self).send_command(command, args)
-
-        success = event.wait(timeout)
-        retval = self.factory.ended_tasks[uuid]
-
-        del self.factory.ended_tasks[uuid]
-        del self.factory.running_tasks[uuid]
-
-        if not success:
-            raise TimeoutError()
-
-        return retval
+        logger.info("disconnected from %s (%s)", self.factory.vm,
+                    self.transport.addr)
+        reactor.connections[self.factory.vm].remove(self)
 
 
 class SerialLineReceiverFactory(protocol.ClientFactory):
@@ -124,5 +105,7 @@ class SerialLineReceiverFactory(protocol.ClientFactory):
 
     def __init__(self, vm):
         self.vm = vm
-        self.running_tasks = {}
-        self.ended_tasks = {}
+        if vm not in reactor.running_tasks:
+            reactor.running_tasks[vm] = {}
+        if vm not in reactor.ended_tasks:
+            reactor.ended_tasks[vm] = {}