Define thrift implementation

requirements.txt

@@ -1,4 +1,8 @@
-argparse==1.2.1
+appdirs==1.4.2
 mirakuru==0.8.2
+packaging==16.8
 psutil==5.1.3
-wsgiref==0.1.2
+pyparsing==2.1.10
+requests==2.13.0
+six==1.10.0
+thrift==0.10.0

worker/config.json

@@ -1,4 +1,8 @@
 {
+	"thrift_host": "0.0.0.0",
+	"thrift_port": "9090",
+	"zk_host": "localhost",
+	"zk_port": "2181",
     "droids": [
         {
             "name": "r2-d2",

worker/droid_service.thrift

@@ -1,6 +1,6 @@
 struct ConnParams {
   1: string host,
-  2: i32 port,
+  2: string port,
 }
 
 service DroidService {

worker/helpers.py

@@ -1,38 +1,70 @@
 import os
 
 from worker.emulator import Emulator
+from zk_client import DroidZkClient
 
 
-class EndpointBuilder(object):
-	def __init__(self):
-		self.port = os.environ.get('ADB_PORT', '5554')
-		self.avd = os.environ.get('AVD_NAME', 'nexus6-android7')
+class DroidBuilder(object):
+    def __init__(self):
+        self.port = os.environ.get('ADB_PORT', '5554')
+        self.avd = os.environ.get('AVD_NAME', 'nexus6-android7')
 
-	def set_port(self, port):
-		self.port = port
-		return self
+    def set_port(self, port):
+        self.port = port
+        return self
 
-	def set_avd(self, avd):
-		self.avd = avd
-		return self
+    def set_avd(self, avd):
+        self.avd = avd
+        return self
 
-	def build(self):
-		return Emulator(port=self.port, avd=self.avd)
+    def build(self):
+        return Emulator(port=self.port, avd=self.avd)
 
 
-class EndpointCoordinator(object):
-	def __init__(self):
-		self.instances = {}
+class DroidCoordinator(object):
+    def __init__(self):
+        self.instances = {}
 
-	def set_endpoint(self, id, endpoint):
-		self.instances[id] = endpoint
+    def set_endpoint(self, id, endpoint):
+        self.instances[id] = {
+            'endpoint': endpoint,
+            'zk_client': DroidZkClient(id),
+        }
 
-	def get_endpoint(self, id):
-		return self.instances[id]
+    def get_endpoint(self, id):
+        return self.instances[id]['endpoint']
 
-	def start_all_instances(self):
-		for instance in self.instances.values():
-			instance.start()
+    def get_zk_client(self, id):
+        return self.instances[id]['zl_client']
 
-	def get_num_endpoints(self):
-		return len(self.instances)
+    def get_num_endpoints(self):
+        return len(self.instances)
+
+    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 iter_endpoints(self):
+        return self.instances.iteritems()
+
+    def setup(self):
+        for instance_id, _ in self.iter_instances():
+            self.start_endpoint(instance_id)
+
+    def stop_endpoint(self, id):
+        endpoint = self.get_endpoint(id)
+        # Stop droid
+        endpoint.stop()
+        # Setup Zk client
+        zk_client = self.get_zk_client(id)
+        zk_client.teardown()
+
+    def teardown(self):
+        for instance_id, _ in self.iter_instances():
+            self.stop_endpoint(instance_id)
+
+

worker/thrift_client.py

@@ -0,0 +1,47 @@
+from contextlib import contextmanager
+
+from tgen.droid_service import DroidService
+
+from thrift import Thrift
+from thrift.transport import TSocket
+from thrift.transport import TTransport
+from thrift.protocol import TBinaryProtocol
+
+
+class DroidClient(object):
+	def __init__(self, host, port):
+		self.host = host
+		self.port = port
+
+		self.client = None
+
+	@contextmanager
+	def transport(self):
+		transport = TSocket.TSocket(self.host, self.port)
+		transport = TTransport.TBufferedTransport(transport)
+		protocol = TBinaryProtocol.TBinaryProtocol(transport)
+		self.client = DroidService.Client(protocol)
+		transport.open()
+		yield
+		transport.close()
+		self.client = None
+
+	def ping(self):
+		with self.transport():
+			self.client.ping()
+
+	def get_package_name(self, apk_url):
+		with self.transport():
+			return self.client.get_package_name(apk_url)
+
+	def get_endpoint(self, endpoint_id):
+		with self.transport():
+			return self.client.get_endpoint(endpoint_id)
+
+	def install_apk(self, endpoint_id, apk_url):
+		with self.transport():
+			return self.client.install_apk(endpoint_id, apk_url)
+
+	def start_package(self, endpoint_id, package_name):
+		with self.transport():
+			return self.client.start_package(endpoint_id, package_name)

worker/thrift_server.py

@@ -1,7 +1,7 @@
 from logger import logger
 from worker.factory import EndpointFactory
 from worker.helpers import EndpointCoordinator, EndpointBuilder
-from worker.utils import read_config
+from worker.utils import get_config, get_public_hostname
 
 from tgen.droid_service.ttypes import ConnParams
 from tgen.droid_service import DroidService
@@ -18,49 +18,56 @@ def __init__(self):
 
     def setup(self):
         logger.debug('Setting up droids')
-        config = read_config()
+        config = get_config()
         builder = EndpointBuilder()
         for droid in config['droids']:
             if droid['port']:
                 builder.set_port(droid['port'])
             if droid['avd']:
                 builder.set_avd(droid['avd'])
             self.coordinator.set_endpoint(droid['name'], builder.build())
-        # Start all instances now
-        self.coordinator.start_all_instances()
+        # Start all endpoints now
+        self.coordinator.setup()
 
     def ping(self):
         logger.debug('Ping!')
 
-    def get_package_name(apk_url):
+    def get_package_name(self, apk_url):
         logger.debug("getting package name for apk {}".format(apk_url))
         return "some random package"
 
-    def get_endpoint(endpoint_id):
+    def get_endpoint(self, endpoint_id):
         logger.debug("Getting endpoint for {}".format(endpoint_id))
+        endpoint = self.coordinator.get_endpoint(endpoint_id)
         cp = ConnParams()
-        cp.host = 'deathstart'
-        cp.port = 54321
+        cp.host = get_public_hostname()
+        cp.port = endpoint.port
         return cp
 
-    def install_apk(endpoint_id, apk_url):
+    def install_apk(self, endpoint_id, apk_url):
         logger.debug("installing apk in {}".format(endpoint_id))
         return True
 
-    def start_package(endpoint_id, package_name):
+    def start_package(self, endpoint_id, package_name):
         logger.debug("Starting package {} for {}".format(package_name, endpoint_id))
         return True
 
     def pre_server_start_log(self):
         logger.debug("{} droid(s) at your service".format(
             self.coordinator.get_num_endpoints()))
 
+    def teardown(self):
+        self.coordinator.teardown()
+
 
 def start_server():
     handler = DroidServiceHandler()
     handler.setup()
     processor = DroidService.Processor(handler)
-    transport = TSocket.TServerSocket(port=9090)
+    config = get_config()
+    transport = TSocket.TServerSocket(
+        host=config['thrift_port'],
+        port=int(config['thrift_port']))
     tfactory = TTransport.TBufferedTransportFactory()
     pfactory = TBinaryProtocol.TBinaryProtocolFactory()
 

worker/utils.py

@@ -1,47 +1,44 @@
+from copy import deepcopy
 import json
 import os
 import socket
 
+import requests
+from requests.exceptions import ConnectTimeout
+
 from settings import ROOT_DIR
 
+
 def is_open_port(port):
     sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
     sock.settimeout(2)
     result = sock.connect_ex(('127.0.0.1', port))
     return result != 0  # Implies port is being used
 
 
+_HOSTNAME = None
 def get_public_hostname():
-    hostname = ''
-    resp = requests.get("http://169.254.169.254/latest/meta-data/public-host")
-    if resp.status == 200:
-        hostname = resp.text
-    else:
-        hostname = socket.gethostname()
-
-    return hostname
-
-def get_hostname():
-    return socket.gethostname()
-
-
-def build_config(emulator):
-    return {
-        'hostname': get_hostname(),
-        'adb_port': emulator.port,
-        'ws_port': 'n/a',
-        'vnc_port': 'n/a',
-    }
-
-
-def read_config():
-    config_path = os.path.join(ROOT_DIR, 'worker',  'config.json')
-    assert os.path.exists(config_path), "Config file not provided for worker"
-    config = {}
-    with open(config_path) as source:
-        config = json.load(source)
-
-    return config
+    if _HOSTNAME is None:
+        try:
+            resp = requests.get("http://169.254.169.254/latest/meta-data/public-host",
+                                timeout=(3, 2))
+            if resp.status == 200:
+                _HOSTNAME = resp.text
+        except ConnectTimeout:
+            _HOSTNAME = socket.gethostname()
+
+    return _HOSTNAME
+
+
+_CONFIG = None
+def get_config():
+    if _CONFIG is None:
+        config_path = os.path.join(ROOT_DIR, 'worker',  'config.json')
+        assert os.path.exists(config_path), "Config file not provided for worker"
+        with open(config_path) as source:
+            _CONFIG = json.load(source)
+
+    return deepcopy(_CONFIG)
 
 
 if __name__ == '__main__':