Merge "Store label info with launcher registration"

nodepool/driver/__init__.py

@@ -217,7 +217,7 @@ class NodeRequestHandler(object):
         # want to make sure we don't continuously grow this array.
         if self.launcher_id not in self.request.declined_by:
             self.request.declined_by.append(self.launcher_id)
-        launchers = set(self.zk.getRegisteredLaunchers())
+        launchers = set([x.id for x in self.zk.getRegisteredLaunchers()])
         if launchers.issubset(set(self.request.declined_by)):
             # All launchers have declined it
             self.log.debug("Failing declined node request %s",

nodepool/launcher.py

@@ -285,7 +285,21 @@ class PoolWorker(threading.Thread):
                 self.log.info("ZooKeeper available. Resuming")
 
             # Make sure we're always registered with ZK
-            self.zk.registerLauncher(self.launcher_id)
+            launcher = zk.Launcher()
+            launcher.id = self.launcher_id
+            for prov_cfg in self.nodepool.config.providers.values():
+                if prov_cfg.driver.name in ('openstack', 'fake'):
+                    for pool_cfg in prov_cfg.pools.values():
+                        launcher.supported_labels.update(
+                            set(pool_cfg.labels.keys()))
+                elif prov_cfg.driver.name == 'static':
+                    for pool_cfg in prov_cfg.pools.values():
+                        launcher.supported_labels.update(pool_cfg.labels)
+                else:
+                    self.log.error(
+                        "Launcher registration unhandled driver: %s",
+                        prov_cfg.driver.name)
+            self.zk.registerLauncher(launcher)
 
             try:
                 if not self.paused_handler:

nodepool/tests/fixtures/launcher_reg1.yaml

@@ -0,0 +1,50 @@
+elements-dir: .
+images-dir: '{images_dir}'
+build-log-dir: '{build_log_dir}'
+build-log-retention: 1
+
+zookeeper-servers:
+  - host: {zookeeper_host}
+    port: {zookeeper_port}
+    chroot: {zookeeper_chroot}
+
+labels:
+  - name: fake-label
+    min-ready: 1
+  - name: fake-label-unused
+
+providers:
+  - name: fake-provider
+    cloud: fake
+    driver: fake
+    region-name: fake-region
+    rate: 0.0001
+    diskimages:
+      - name: fake-image
+        meta:
+          key: value
+          key2: value
+    pools:
+      - name: main
+        max-servers: 96
+        availability-zones:
+          - az1
+        networks:
+          - net-name
+        labels:
+          - name: fake-label
+            diskimage: fake-image
+            min-ram: 8192
+            flavor-name: 'Fake'
+
+diskimages:
+  - name: fake-image
+    elements:
+      - fedora
+      - vm
+    release: 21
+    env-vars:
+      TMPDIR: /opt/dib_tmp
+      DIB_IMAGE_CACHE: /opt/dib_cache
+      DIB_CLOUD_IMAGES: http://download.fedoraproject.org/pub/fedora/linux/releases/test/21-Beta/Cloud/Images/x86_64/
+      BASE_IMAGE_FILE: Fedora-Cloud-Base-20141029-21_Beta.x86_64.qcow2

nodepool/tests/fixtures/launcher_reg2.yaml

@@ -0,0 +1,54 @@
+elements-dir: .
+images-dir: '{images_dir}'
+build-log-dir: '{build_log_dir}'
+build-log-retention: 1
+
+zookeeper-servers:
+  - host: {zookeeper_host}
+    port: {zookeeper_port}
+    chroot: {zookeeper_chroot}
+
+labels:
+  - name: fake-label
+    min-ready: 1
+  - name: fake-label2
+
+providers:
+  - name: fake-provider
+    cloud: fake
+    driver: fake
+    region-name: fake-region
+    rate: 0.0001
+    diskimages:
+      - name: fake-image
+        meta:
+          key: value
+          key2: value
+    pools:
+      - name: main
+        max-servers: 96
+        availability-zones:
+          - az1
+        networks:
+          - net-name
+        labels:
+          - name: fake-label
+            diskimage: fake-image
+            min-ram: 8192
+            flavor-name: 'Fake'
+          - name: fake-label2
+            diskimage: fake-image
+            min-ram: 8192
+            flavor-name: 'Fake'
+
+diskimages:
+  - name: fake-image
+    elements:
+      - fedora
+      - vm
+    release: 21
+    env-vars:
+      TMPDIR: /opt/dib_tmp
+      DIB_IMAGE_CACHE: /opt/dib_cache
+      DIB_CLOUD_IMAGES: http://download.fedoraproject.org/pub/fedora/linux/releases/test/21-Beta/Cloud/Images/x86_64/
+      BASE_IMAGE_FILE: Fedora-Cloud-Base-20141029-21_Beta.x86_64.qcow2

nodepool/tests/test_launcher.py

@@ -1236,3 +1236,27 @@ class TestLauncher(tests.DBTestCase):
         self.waitForNodeDeletion(label1_nodes[0])
         req = self.waitForNodeRequest(req)
         self.assertEqual(req.state, zk.FULFILLED)
+
+    def test_launcher_registers_config_change(self):
+        '''
+        Launchers register themselves and some config info with ZooKeeper.
+        Validate that a config change will propogate to ZooKeeper.
+        '''
+        configfile = self.setup_config('launcher_reg1.yaml')
+        self.useBuilder(configfile)
+        pool = self.useNodepool(configfile, watermark_sleep=1)
+        pool.start()
+
+        self.waitForNodes('fake-label')
+        launchers = self.zk.getRegisteredLaunchers()
+        self.assertEqual(1, len(launchers))
+
+        # the fake-label-unused label should not appear
+        self.assertEqual({'fake-label'}, launchers[0].supported_labels)
+
+        self.replace_config(configfile, 'launcher_reg2.yaml')
+
+        # we should get 1 additional label now
+        while launchers[0].supported_labels != {'fake-label', 'fake-label2'}:
+            time.sleep(1)
+            launchers = self.zk.getRegisteredLaunchers()

nodepool/tests/test_zk.py

@@ -418,19 +418,21 @@ class TestZooKeeper(tests.DBTestCase):
         self.assertIsNone(self.zk.client.exists(path))
 
     def test_registerLauncher(self):
-        name = "launcher-000-001"
-        self.zk.registerLauncher(name)
+        launcher = zk.Launcher()
+        launcher.id = "launcher-000-001"
+        self.zk.registerLauncher(launcher)
         launchers = self.zk.getRegisteredLaunchers()
         self.assertEqual(1, len(launchers))
-        self.assertEqual(name, launchers[0])
+        self.assertEqual(launcher.id, launchers[0].id)
 
     def test_registerLauncher_safe_repeat(self):
-        name = "launcher-000-001"
-        self.zk.registerLauncher(name)
-        self.zk.registerLauncher(name)
+        launcher = zk.Launcher()
+        launcher.id = "launcher-000-001"
+        self.zk.registerLauncher(launcher)
+        self.zk.registerLauncher(launcher)
         launchers = self.zk.getRegisteredLaunchers()
         self.assertEqual(1, len(launchers))
-        self.assertEqual(name, launchers[0])
+        self.assertEqual(launcher.id, launchers[0].id)
 
     def test_getNodeRequests_empty(self):
         self.assertEqual([], self.zk.getNodeRequests())

nodepool/zk.py

@@ -14,6 +14,7 @@
 
 from contextlib import contextmanager
 from copy import copy
+import abc
 import json
 import logging
 import six
@@ -121,7 +122,69 @@ class ZooKeeperWatchEvent(object):
         self.image = image
 
 
-class BaseModel(object):
+class Serializable(abc.ABC):
+    '''
+    Abstract base class for objects that will be stored in ZooKeeper.
+    '''
+
+    @abc.abstractmethod
+    def toDict(self):
+        '''
+        Return a dictionary representation of the object.
+        '''
+        pass
+
+    def serialize(self):
+        '''
+        Return a representation of the object as a string.
+
+        Used for storing the object data in ZooKeeper.
+        '''
+        return json.dumps(self.toDict()).encode('utf8')
+
+
+class Launcher(Serializable):
+    '''
+    Class to describe a nodepool launcher.
+    '''
+
+    def __init__(self):
+        self.id = None
+        self._supported_labels = set()
+
+    def __eq__(self, other):
+        if isinstance(other, Launcher):
+            return (self.id == other.id and
+                    self.supported_labels == other.supported_labels)
+        else:
+            return False
+
+    @property
+    def supported_labels(self):
+        return self._supported_labels
+
+    @supported_labels.setter
+    def supported_labels(self, value):
+        if not isinstance(value, set):
+            raise TypeError("'supported_labels' attribute must be a set")
+        self._supported_labels = value
+
+    def toDict(self):
+        d = {}
+        d['id'] = self.id
+        # sets are not JSON serializable, so use a sorted list
+        d['supported_labels'] = sorted(self.supported_labels)
+        return d
+
+    @staticmethod
+    def fromDict(d):
+        obj = Launcher()
+        obj.id = d.get('id')
+        obj.supported_labels = set(d.get('supported_labels', []))
+        return obj
+
+
+class BaseModel(Serializable):
     VALID_STATES = set([])
 
     def __init__(self, o_id):
@@ -177,14 +240,6 @@ class BaseModel(object):
         if 'state_time' in d:
             self.state_time = d['state_time']
 
-    def serialize(self):
-        '''
-        Return a representation of the object as a string.
-
-        Used for storing the object data in ZooKeeper.
-        '''
-        return json.dumps(self.toDict()).encode('utf8')
-
 
 class ImageBuild(BaseModel):
     '''
@@ -1315,27 +1370,44 @@ class ZooKeeper(object):
         otherwise disconnects from ZooKeeper. It will need to re-register
         after a lost connection. This method is safe to call multiple times.
 
-        :param str launcher: Unique name for the launcher.
+        :param Launcher launcher: Object describing the launcher.
         '''
-        path = self._launcherPath(launcher)
+        path = self._launcherPath(launcher.id)
 
-        try:
-            self.client.create(path, makepath=True, ephemeral=True)
-        except kze.NodeExistsError:
-            pass
+        if self.client.exists(path):
+            data, _ = self.client.get(path)
+            obj = Launcher.fromDict(self._bytesToDict(data))
+            if obj != launcher:
+                self.client.set(path, launcher.serialize())
+                self.log.debug("Updated registration for launcher %s",
+                               launcher.id)
+        else:
+            self.client.create(path, value=launcher.serialize(),
+                               makepath=True, ephemeral=True)
+            self.log.debug("Registered launcher %s", launcher.id)
 
     def getRegisteredLaunchers(self):
         '''
         Get a list of all launchers that have registered with ZooKeeper.
 
-        :returns: A list of launcher names, or empty list if none are found.
+        :returns: A list of Launcher objects, or empty list if none are found.
         '''
         try:
-            launchers = self.client.get_children(self.LAUNCHER_ROOT)
+            launcher_ids = self.client.get_children(self.LAUNCHER_ROOT)
         except kze.NoNodeError:
             return []
 
-        return launchers
+        objs = []
+        for launcher in launcher_ids:
+            path = self._launcherPath(launcher)
+            try:
+                data, _ = self.client.get(path)
+            except kze.NoNodeError:
+                # launcher disappeared
+                continue
+
+            objs.append(Launcher.fromDict(self._bytesToDict(data)))
+        return objs
 
     def getNodeRequests(self):
         '''