initial mesosutil and orchestrator.Tagger work

tests/core/test_orchestrator.py

@@ -0,0 +1,154 @@
+# stdlib
+import unittest
+import os
+
+# 3rd party
+import mock
+
+# project
+from utils.orchestrator import MesosUtil, BaseUtil
+
+CO_ID = 1234
+
+
+class TestMesosUtil(unittest.TestCase):
+    @mock.patch('docker.Client.__init__')
+    def test_extract_tags(self, mock_init):
+        mock_init.return_value = None
+        mesos = MesosUtil()
+
+        env = ["CHRONOS_JOB_NAME=test-job",
+               "MARATHON_APP_ID=/system/dd-agent",
+               "MESOS_TASK_ID=system_dd-agent.dcc75b42-4b87-11e7-9a62-70b3d5800001"]
+
+        tags = ['chronos_job:test-job', 'marathon_app:/system/dd-agent',
+                'mesos_task:system_dd-agent.dcc75b42-4b87-11e7-9a62-70b3d5800001']
+
+        container = {'Config': {'Env': env}}
+
+        self.assertEqual(sorted(tags), sorted(mesos._get_cacheable_tags(CO_ID, co=container)))
+
+    @mock.patch.dict(os.environ, {"MESOS_TASK_ID": "test"})
+    def test_detect(self):
+        self.assertTrue(MesosUtil.is_detected())
+
+    @mock.patch.dict(os.environ, {})
+    def test_no_detect(self):
+        self.assertFalse(MesosUtil.is_detected())
+
+
+class TestBaseUtil(unittest.TestCase):
+    class DummyUtil(BaseUtil):
+        def _get_cacheable_tags(self, cid, co=None):
+            return ["test:tag"]
+
+    class NeedInspectUtil(BaseUtil):
+        def __init__(self):
+            BaseUtil.__init__(self)
+            self.needs_inspect = True
+
+        def _get_cacheable_tags(self, cid, co=None):
+            return ["test:tag"]
+
+    class NeedEnvUtil(BaseUtil):
+        def __init__(self):
+            BaseUtil.__init__(self)
+            self.needs_env = True
+
+        def _get_cacheable_tags(self, cid, co=None):
+            return ["test:tag"]
+
+    @mock.patch('docker.Client.__init__')
+    def test_extract_tags(self, mock_init):
+        mock_init.return_value = None
+        dummy = self.DummyUtil()
+        dummy.reset_cache()
+
+        self.assertEqual(["test:tag"], dummy.get_container_tags(cid=CO_ID))
+
+    @mock.patch('docker.Client.__init__')
+    def test_cache_invalidation_event(self, mock_init):
+        mock_init.return_value = None
+        dummy = self.DummyUtil()
+        dummy.reset_cache()
+
+        dummy.get_container_tags(cid=CO_ID)
+        self.assertTrue(CO_ID in dummy._container_tags_cache)
+
+        EVENT = {'status': 'die', 'id': CO_ID}
+        dummy.invalidate_cache([EVENT])
+        self.assertFalse(CO_ID in dummy._container_tags_cache)
+
+    @mock.patch('docker.Client.__init__')
+    def test_reset_cache(self, mock_init):
+        mock_init.return_value = None
+        dummy = self.DummyUtil()
+        dummy.reset_cache()
+
+        dummy.get_container_tags(cid=CO_ID)
+        self.assertTrue(CO_ID in dummy._container_tags_cache)
+
+        dummy.reset_cache()
+        self.assertFalse(CO_ID in dummy._container_tags_cache)
+
+    @mock.patch('docker.Client.inspect_container')
+    @mock.patch('docker.Client.__init__')
+    def test_auto_inspect(self, mock_init, mock_inspect):
+        mock_init.return_value = None
+
+        dummy = self.NeedInspectUtil()
+        dummy.reset_cache()
+
+        dummy.get_container_tags(cid=CO_ID)
+        mock_inspect.assert_called_once()
+
+    @mock.patch('docker.Client.inspect_container')
+    @mock.patch('docker.Client.__init__')
+    def test_no_inspect_if_cached(self, mock_init, mock_inspect):
+        mock_init.return_value = None
+
+        dummy = self.NeedInspectUtil()
+        dummy.reset_cache()
+
+        dummy.get_container_tags(cid=CO_ID)
+        mock_inspect.assert_called_once()
+
+        dummy.get_container_tags(cid=CO_ID)
+        mock_inspect.assert_called_once()
+
+    @mock.patch('docker.Client.inspect_container')
+    @mock.patch('docker.Client.__init__')
+    def test_no_useless_inspect(self, mock_init, mock_inspect):
+        mock_init.return_value = None
+
+        dummy = self.NeedInspectUtil()
+        dummy.reset_cache()
+        co = {'Id': CO_ID, 'Created': 1}
+
+        dummy.get_container_tags(cid=1, co=co)
+        mock_inspect.assert_not_called()
+
+        dummy.get_container_tags(co=co)
+        mock_inspect.assert_not_called()
+
+    @mock.patch('docker.Client.inspect_container')
+    @mock.patch('docker.Client.__init__')
+    def test_auto_env_inspect(self, mock_init, mock_inspect):
+        mock_init.return_value = None
+
+        dummy = self.NeedEnvUtil()
+        dummy.reset_cache()
+
+        dummy.get_container_tags(co={'Id': CO_ID})
+        mock_inspect.assert_called_once()
+
+    @mock.patch('docker.Client.inspect_container')
+    @mock.patch('docker.Client.__init__')
+    def test_no_useless_env_inspect(self, mock_init, mock_inspect):
+        mock_init.return_value = None
+
+        dummy = self.NeedEnvUtil()
+        dummy.reset_cache()
+
+        dummy.get_container_tags(co={'Id': CO_ID, 'Config': {'Env': {1: 1}}})
+        mock_inspect.assert_not_called()

utils/orchestrator/__init__.py

@@ -2,5 +2,9 @@
 # All rights reserved
 # Licensed under Simplified BSD License (see LICENSE)
 
-from nomadutil import NomadUtil  # noqa: F401
 from ecsutil import ECSUtil  # noqa: F401
+from mesosutil import MesosUtil  # noqa: F401
+from nomadutil import NomadUtil  # noqa: F401
+from baseutil import BaseUtil  # noqa: F401
+
+from tagger import Tagger  # noqa: F401

utils/orchestrator/baseutil.py

@@ -0,0 +1,85 @@
+# (C) Datadog, Inc. 2017
+# All rights reserved
+# Licensed under Simplified BSD License (see LICENSE)
+
+# stdlib
+import logging
+
+# project
+from utils.dockerutil import DockerUtil
+from utils.singleton import Singleton
+
+
+class BaseUtil:
+    """
+    Base class for orchestrator utils. Only handles container tags for now.
+    Users should go through the orchestrator.Tagger class to simplify the code
+
+    Children classes can implement:
+      - __init__: to change self.needs_inspect
+      - _get_cacheable_tags: tags will be cached for reuse
+      - _get_transient_tags: tags can change and won't be cached (TODO)
+      - invalidate_cache: custom cache invalidation logic
+      - is_detected (staticmethod)
+    """
+    __metaclass__ = Singleton
+
+    def __init__(self):
+        # Whether your get___tags methods need the inspect result
+        self.needs_inspect = False
+        # Whether your methods need the env portion (not in partial inspect)
+        self.needs_env = False
+
+        self.log = logging.getLogger(__name__)
+        self.docker_util = DockerUtil()
+
+        # Tags cache as a dict {co_id: [tags]}
+        self._container_tags_cache = {}
+
+    def get_container_tags(self, cid=None, co=None):
+        """
+        Returns container tags for the given container, inspecting the container if needed
+        :param container: either the container id or container dict returned by docker-py
+        :return: tags as list<string>, cached
+        """
+
+        if (cid is not None) and (co is not None):
+            self.log.error("Can only pass either a container id or object, not both, returning empty tags")
+            return []
+        if (cid is None) and (co is None):
+            self.log.error("Need one container id or container object, returning empty tags")
+            return []
+        elif co is not None:
+            if 'Id' in co:
+                cid = co.get('Id')
+            else:
+                self.log.warning("Invalid container dict, returning empty tags")
+                return []
+
+        if cid in self._container_tags_cache:
+            return self._container_tags_cache[cid]
+        else:
+            if (self.needs_inspect or self.needs_env) and co is None:
+                co = self.docker_util.inspect_container(cid)
+            if self.needs_env and 'Env' not in co.get('Config', {}):
+                co = self.docker_util.inspect_container(cid)
+            self._container_tags_cache[cid] = self._get_cacheable_tags(cid, co)
+            return self._container_tags_cache[cid]
+
+    def invalidate_cache(self, events):
+        """
+        Allows cache invalidation when containers die
+        :param events from self.get_events
+        """
+        try:
+            for ev in events:
+                if ev.get('status') == 'die' and ev.get('id') in self._container_tags_cache:
+                    del self._container_tags_cache[ev.get('id')]
+        except Exception as e:
+            self.log.warning("Error when invalidating tag cache: " + str(e))
+
+    def reset_cache(self):
+        """
+        Empties all caches to reset the singleton to initial state
+        """
+        self._container_tags_cache = {}

utils/orchestrator/mesosutil.py

@@ -0,0 +1,40 @@
+# (C) Datadog, Inc. 2017
+# All rights reserved
+# Licensed under Simplified BSD License (see LICENSE)
+
+import os
+
+from .baseutil import BaseUtil
+
+CHRONOS_JOB_NAME = "CHRONOS_JOB_NAME"
+MARATHON_APP_ID = "MARATHON_APP_ID"
+MESOS_TASK_ID = "MESOS_TASK_ID"
+
+
+class MesosUtil(BaseUtil):
+    def __init__(self):
+        BaseUtil.__init__(self)
+        self.needs_inspect = True
+        self.needs_env = True
+
+    def _get_cacheable_tags(self, cid, co=None):
+        tags = []
+
+        self.log.warning("called")
+
+        envvars = co.get('Config', {}).get('Env', {})
+        self.log.warning(envvars)
+
+        for var in envvars:
+            if var.startswith(CHRONOS_JOB_NAME):
+                tags.append('chronos_job:%s' % var[len(CHRONOS_JOB_NAME) + 1:])
+            elif var.startswith(MARATHON_APP_ID):
+                tags.append('marathon_app:%s' % var[len(MARATHON_APP_ID) + 1:])
+            elif var.startswith(MESOS_TASK_ID):
+                tags.append('mesos_task:%s' % var[len(MESOS_TASK_ID) + 1:])
+
+        return tags
+
+    @staticmethod
+    def is_detected():
+        return MESOS_TASK_ID in os.environ

utils/orchestrator/tagger.py

@@ -0,0 +1,59 @@
+# (C) Datadog, Inc. 2017
+# All rights reserved
+# Licensed under Simplified BSD License (see LICENSE)
+
+
+from .mesosutil import MesosUtil
+from utils.singleton import Singleton
+
+
+class Tagger():
+    """
+    Wraps several BaseUtil classes with autodetection and allows to query
+    them through the same interface as BaseUtil classes
+
+    See BaseUtil for apidoc
+    """
+    __metaclass__ = Singleton
+
+    def __init__(self):
+        self._utils = []  # [BaseUtil object]
+        self._has_detected = False
+        self.reset()
+
+    def get_container_tags(self, cid=None, co=None):
+        concat_tags = []
+        for util in self._utils:
+            tags = util.get_container_tags(cid, co)
+            if tags:
+                concat_tags.extend(tags)
+
+        return concat_tags
+
+    def invalidate_cache(self, events):
+        for util in self._utils:
+            util.invalidate_cache(events)
+
+    def reset_cache(self):
+        for util in self._utils:
+            util.reset_cache()
+
+    def reset(self):
+        """
+        Trigger a new autodetection and reset underlying util classes
+        """
+        self._utils = []
+
+        if MesosUtil.is_detected():
+            m = MesosUtil()
+            m.reset_cache()
+            self._utils.append(m)
+
+        self._has_detected = bool(self._utils)
+
+    def has_detected(self):
+        """
+        Returns whether the tagger has detected orchestrators it handles
+        If false, calling get_container_tags will return an empty list
+        """
+        return self._has_detected

utils/platform.py

@@ -6,6 +6,8 @@
 import os
 import sys
 
+import utils.orchestrator
+
 
 def get_os():
     "Human-friendly OS name"
@@ -22,6 +24,7 @@ def get_os():
     else:
         return sys.platform
 
+
 class Platform(object):
     """
     Return information about the given platform.
@@ -108,3 +111,7 @@ def is_swarm():
     @staticmethod
     def is_nomad():
         return 'NOMAD_ALLOC_ID' in os.environ
+
+    @staticmethod
+    def is_mesos():
+        return utils.orchestrator.MesosUtil.is_detected()