Push everything quick!

master/droid_master.thrift

@@ -0,0 +1,16 @@
+struct ConnParams {
+  1: string host,
+  2: string port,
+}
+
+service MasterService {
+   void ping(),
+
+   string get_package_name(1: string apk_url),
+
+   ConnParams get_endpoint(1: string user),
+
+   bool install_apk(1: string endpoint_id, 2: string apk_url),
+
+   bool start_package(1: string endpoint_id, 2: string package_name),
+}

master/zk_master.py

@@ -0,0 +1,92 @@
+import json
+
+from kazoo.client import KazooClient, KazooState
+from kazoo.protocol.states import EventType
+from kazoo.recipe.watchers import DataWatch
+
+from logger import logger
+from master.settings import ZOOKEEPER_HOST, ZOOKEEPER_PORT
+
+
+class MasterZkClient(object):
+    def __init__(self):
+        host='{}:{}'.format(ZOOKEEPER_HOST, ZOOKEEPER_PORT)
+        zk = KazooClient(host)
+        zk.add_listener(self.conn_listener)
+        self.zk = zk
+        self.on_available_droid = ChildrenWatch(
+            zk,
+            '/droids/available',
+            self.on_available_droid,
+        )
+
+    def setup(self):
+        self.zk.ensure_path('/droids/available')
+        self.zk.ensure_path('/droids/assigned')
+
+    @staticmethod
+    def conn_listener(state):
+        if state == KazooState.CONNECTED:
+            logger.debug('connected...')
+        elif state == KazooState.LOST:
+            logger.debug('connection lost...')
+        elif state == KazooState.SUSPENDED:
+            logger.debug('connection suspended...')
+
+
+    def on_available_droid(self, children):
+        logger.debug('Registering watches on available droids')
+        for child in children:
+            child_p = '/droids/available/{}'.format(child)
+            DataWatch(self.zk, child_p, self.handle_dropped_droid)
+
+
+    def handle_dropped_droid(self, data, stat, event):
+        if event:
+            if event.type == EventType.CREATED:
+                # Implies a new droid became available
+                logger.debug('WooHoo! Another droid joins our ranks!')
+            elif event.type == EventType.DELETED:
+                # Implies an existing droid was disconnected
+                # Check if it was assigned and remove it from there.
+                node_name = event.path.rsplit('/', 1)[1]
+                assigned_emulators = self.zk.get_children('/droids/assigned')
+                if node_name in assigned_emulators:
+                    # Check if droid was being assigned, by comparing creation
+                    # times.
+                    assigned_p = '/droids/assigned/{}'.format(node_name)
+                    adata, astat = self.zk.get(assigned_p)
+                    import ipdb; ipdb.set_trace()
+                    logger.error('Oh no, droid-{} was lost'.format(node_name))
+                    self.zk.delete(assigned_p)
+            else:
+                logger.debug('These are not the events you are looking for.')
+
+    def assign_droid(self, user):
+        available_droids = self.zk.get_children('/droids/available')
+        if not available_droids:
+            logger.error('No droids available')
+            return False
+        child = available_droids[0]
+        child_p = '/droids/available/{}'.format(child)
+        data, stat = self.zk.get(child_p)
+        data = json.loads(data)
+        transaction = self.zk.transaction()
+        transaction.check(child_p, stat.version)
+        transaction.delete(child_p)
+        data['user'] = user
+        transaction.create(
+            '/droids/assigned/{}'.format(child),
+            value=json.dumps(data),
+        )
+        transaction.commit()
+        return True
+
+
+if __name__ == '__main__':
+    zk.start()
+    init_nodes()
+    from twisted.internet import reactor
+
+    reactor.addSystemEventTrigger('before', 'shutdown', lambda: zk.stop())
+    reactor.run()

settings.py

@@ -1,3 +1,5 @@
 import os
 
-ROOT_DIR = os.path.dirname(__file__)
+ROOT_DIR = os.path.dirname(__file__)
+ZK_HOST = "localhost"
+ZK_PORT = "2181"

worker/config.json

@@ -1,11 +1,8 @@
 {
 	"thrift_host": "0.0.0.0",
 	"thrift_port": "9090",
-	"zk_host": "localhost",
-	"zk_port": "2181",
     "droids": [
         {
-            "name": "r2-d2",
             "port": "5554",
             "avd": "Nexus_6_API_25"
         }

worker/helpers.py

@@ -1,7 +1,7 @@
 import os
 
 from worker.emulator import Emulator
-from zk_client import DroidZkClient
+from zk_droid import DroidZkClient
 
 
 class DroidBuilder(object):
@@ -23,48 +23,51 @@ def build(self):
 
 class DroidCoordinator(object):
     def __init__(self):
-        self.instances = {}
+        self.droids_to_start = []
+        self.initialised = {}
 
-    def set_endpoint(self, id, endpoint):
-        self.instances[id] = {
-            'endpoint': endpoint,
-            'zk_client': DroidZkClient(id),
-        }
+    def add_droid(self, droid):
+        self.droids_to_start.append(droid)
 
-    def get_endpoint(self, id):
-        return self.instances[id]['endpoint']
+    def start_droid(self, droid):
+        droid.start()
+        zk_client = DroidZkClient()
+        zk_client.setup()
+        self.initialised[zk_client.nodename] = {
+            'droid': droid,
+            'zk_client': zk_client,
+        }
 
-    def get_zk_client(self, id):
-        return self.instances[id]['zk_client']
+    def setup(self):
+        while True:
+            try:
+                droid = self.droids_to_start.pop()
+                self.start_droid(droid)
+            except IndexError:
+                break
 
     def count(self):
-        return len(self.instances)
+        return len(self.initialised)
 
-    def start_endpoint(self, id):
-        endpoint = self.get_endpoint(id)
-        # Start droid
-        endpoint.start()
-        # Setup Zk client
-        zk_client = self.get_zk_client(id)
-        zk_client.setup()
+    def get_droid(self, id):
+        return self.initialised[id]['droid']
 
-    def iter_endpoints(self):
-        return self.instances.iteritems()
-
-    def setup(self):
-        for instance_id, _ in self.iter_endpoints():
-            self.start_endpoint(instance_id)
+    def get_zk_client(self, id):
+        return self.initialised[id]['zk_client']
 
-    def stop_endpoint(self, id):
-        endpoint = self.get_endpoint(id)
-        # Stop droid
-        endpoint.stop()
-        # Setup Zk client
+    def stop_droid(self, id):
+        # Stop Zk client
         zk_client = self.get_zk_client(id)
         zk_client.teardown()
+        # Stop droid
+        droid = self.get_droid(id)
+        droid.stop()
+
+    def iter_droid_ids(self):
+        return self.initialised.iterkeys()
 
     def teardown(self):
-        for instance_id, _ in self.iter_endpoints():
-            self.stop_endpoint(instance_id)
+        for instance_id, _ in self.iter_droid_ids():
+            self.stop_droid(instance_id)
 
 

worker/thrift_server.py

@@ -3,6 +3,7 @@
 from logger import logger
 from worker.helpers import DroidCoordinator, DroidBuilder
 from worker.utils import get_config, get_public_hostname
+from worker.utils import get_package_name_from_url
 
 from tgen.droid_service.ttypes import ConnParams
 from tgen.droid_service import DroidService
@@ -26,7 +27,7 @@ def setup(self):
                 builder.set_port(droid['port'])
             if droid['avd']:
                 builder.set_avd(droid['avd'])
-            self.coordinator.set_endpoint(droid['name'], builder.build())
+            self.coordinator.add_droid(builder.build())
         # Start all endpoints now
         self.coordinator.setup()
 
@@ -35,11 +36,11 @@ def ping(self):
 
     def get_package_name(self, apk_url):
         logger.debug("getting package name for apk {}".format(apk_url))
-        return "some random package"
+        return get_package_name_from_url(apk_url)
 
     def get_endpoint(self, endpoint_id):
         logger.debug("Getting endpoint for {}".format(endpoint_id))
-        endpoint = self.coordinator.get_endpoint(endpoint_id)
+        endpoint = self.coordinator.get_droid(endpoint_id)
         cp = ConnParams()
         cp.host = get_public_hostname()
         cp.port = endpoint.port