Address review comments

.bundle/config

@@ -0,0 +1 @@
+--- {}

agent.py

@@ -157,8 +157,11 @@ def refresh_specific_checks(self, hostname, checksd, checks):
 
             # this is an error dict
             # checks that failed to load are added to init_failed_checks
+            # and poped from initialized_checks
             if isinstance(fresh_check, dict) and 'error' in fresh_check.keys():
                 checksd['init_failed_checks'][fresh_check.keys()[0]] = fresh_check.values()[0]
+                if idx:
+                    checksd['initialized_checks'].pop(idx)
 
             elif not fresh_check:
                 # no instance left of it to monitor so the check was not loaded

utils/dockerutil.py

@@ -8,13 +8,15 @@
 import time
 
 # 3rd party
-from docker import Client
-from docker import tls
+from docker import Client, tls
+from docker.errors import NullResource
 
 # project
 from utils.singleton import Singleton
 from utils.service_discovery.config_stores import get_config_store
 
+DATADOG_ID = 'com.datadoghq.sd.check.id'
+
 
 class MountException(Exception):
     pass
@@ -111,7 +113,11 @@ def get_events(self):
             try:
                 if self.config_store and event.get('status') in CONFIG_RELOAD_STATUS:
                     image = event.get('from', '')
-                    checks = self._get_checks_from_image(image)
+                    try:
+                        inspect = self.client.inspect_container(event.get('id'))
+                    except NullResource:
+                        inspect = {}
+                    checks = self._get_checks_from_inspect(inspect)
                     if checks:
                         conf_reload_set.update(set(checks))
                 self.events.append(event)
@@ -124,9 +130,13 @@ def get_events(self):
 
         return self.events, conf_reload_set
 
-    def _get_checks_from_image(self, image):
-        """Get the list of checks applied to an image from the image_to_checks cache in the config store"""
-        return self.config_store.image_to_checks[image]
+    def _get_checks_from_inspect(self, inspect):
+        """Get the list of checks applied to a container from the identifier_to_checks cache in the config store.
+        Use the DATADOG_ID label or the image."""
+        identifier = inspect.get('Config', {}).get('Labels', {}).get(DATADOG_ID) or \
+            inspect.get('Config', {}).get('Image')
+
+        return self.config_store.identifier_to_checks[identifier]
 
     def get_hostname(self):
         """Return the `Name` param from `docker info` to use as the hostname"""

utils/service_discovery/abstract_config_store.py

@@ -45,7 +45,9 @@ def __init__(self, agentConfig):
         self.auto_conf_images = get_auto_conf_images(agentConfig)
 
         # cache used by dockerutil to determine which check to reload based on the image linked to an event
-        self.image_to_checks = self._populate_image_to_checks()
+        # this is a defaultdict(set) and some calls to it rely on this property
+        # so if you're planning on changing that, track its references
+        self.identifier_to_checks = self._populate_identifier_to_checks()
 
     @classmethod
     def _drop(cls):
@@ -65,27 +67,27 @@ def client_read(self, path, **kwargs):
     def dump_directory(self, path, **kwargs):
         raise NotImplementedError()
 
-    def _populate_image_to_checks(self):
-        """Populate the image_to_checks cache with templates pulled
+    def _populate_identifier_to_checks(self):
+        """Populate the identifier_to_checks cache with templates pulled
         from the config store and from the auto-config folder"""
-        images_to_checks = defaultdict(set)
+        identifier_to_checks = defaultdict(set)
         # config store templates
         try:
             templates = self.client_read(self.sd_template_dir.lstrip('/'), all=True)
         except (NotImplementedError, TimeoutError, AttributeError):
             templates = []
         for tpl in templates:
             split_tpl = tpl[0].split('/')
-            image, var = split_tpl[-2], split_tpl[-1]
+            ident, var = split_tpl[-2], split_tpl[-1]
             if var == CHECK_NAMES:
-                images_to_checks[image].update(set(json.loads(tpl[1])))
+                identifier_to_checks[ident].update(set(json.loads(tpl[1])))
 
         # auto-config templates
         templates = get_auto_conf_images(self.agentConfig)
         for image, check in templates.iteritems():
-            images_to_checks[image].add(check)
+            identifier_to_checks[image].add(check)
 
-        return images_to_checks
+        return identifier_to_checks
 
     def _get_auto_config(self, image_name):
         ident = self._get_image_ident(image_name)
@@ -136,8 +138,8 @@ def get_check_tpls(self, identifier, **kwargs):
                       'will not be configured by the service discovery'.format(identifier))
             return []
 
-        # Update the image_to_checks cache
-        self._update_image_to_checks(identifier, check_names)
+        # Try to update the identifier_to_checks cache
+        self._update_identifier_to_checks(identifier, check_names)
 
         for idx, c_name in enumerate(check_names):
             if trace_config:
@@ -237,18 +239,18 @@ def crawl_config_template(self):
             self.previous_config_index = config_index
             return False
         # Config has been modified since last crawl
-        # in this case a full config reload is triggered and the image_to_checks cache is rebuilt
+        # in this case a full config reload is triggered and the identifier_to_checks cache is rebuilt
         if config_index != self.previous_config_index:
             log.info('Detected an update in config templates, reloading check configs...')
             self.previous_config_index = config_index
-            self.image_to_checks = self._populate_image_to_checks()
+            self.identifier_to_checks = self._populate_identifier_to_checks()
             return True
         return False
 
-    def _update_image_to_checks(self, image, check_names):
-        """Try to insert in the image_to_checks cache the mapping between an image and its check names"""
-        if image not in self.image_to_checks:
-            self.image_to_checks[image] = set(check_names)
-        elif self.image_to_checks[image] != set(check_names):
-            log.warning("Trying to cache check names for image %s but a different value is already there. "
-                        "This should not happen. Not updating." % image)
+    def _update_identifier_to_checks(self, identifier, check_names):
+        """Try to insert in the identifier_to_checks cache the mapping between an identifier and its check names"""
+        if identifier not in self.identifier_to_checks:
+            self.identifier_to_checks[identifier] = set(check_names)
+        elif self.identifier_to_checks[identifier] != set(check_names):
+            log.warning("Trying to cache check names for ident %s but a different value is already there."
+                        "This should not happen. Not updating." % identifier)

utils/service_discovery/consul_config_store.py

@@ -46,12 +46,12 @@ def get_client(self, reset=False):
 
     def client_read(self, path, **kwargs):
         """Retrieve a value from a consul key."""
-        recurse = kwargs.get('recursive', False or kwargs.get('all, False'))
+        recurse = kwargs.get('recursive') or kwargs.get('all', False)
         res = self.client.kv.get(path, recurse=recurse)
         if kwargs.get('watch', False):
             return res[0]
         elif kwargs.get('all', False):
-            # we use it in _populate_image_to_checks
+            # we use it in _populate_identifier_to_checks
             return [(child.get('Key'), child.get('Value')) for child in res[1]]
         else:
             if res[1] is not None:

utils/service_discovery/etcd_config_store.py

@@ -44,11 +44,11 @@ def client_read(self, path, **kwargs):
             res = self.client.read(
                 path,
                 timeout=kwargs.get('timeout', DEFAULT_TIMEOUT),
-                recursive=kwargs.get('recursive', False) or kwargs.get('all', False))
+                recursive=kwargs.get('recursive') or kwargs.get('all', False))
             if kwargs.get('watch', False):
                 return res.etcd_index
             elif kwargs.get('all', False):
-                # we use it in _populate_image_to_checks
+                # we use it in _populate_identifier_to_checks
                 return [(child.key, child.value) for child in res.children]
             else:
                 return res.value