Merge pull request #6 from CodersOfTheNight/event_augment

Event augment

Author: zaibacu

AUTHORS.md

@@ -0,0 +1,3 @@
+Maintainers
+============
+Šarūnas Navickas (zaibacu@gmail.com): 2016 - now

CONTRIBUTING.md

@@ -0,0 +1,21 @@
+New agents
+==========
+Most natural way of contributing to this project is creating new agents.
+What is agent? It is basically a metrics collector for specific service, 
+eg. `oshino-redis` is an agent which runs `info` command as a Redis client,
+parses received metrics and gives them to `oshino` as events.
+
+To make life easier, there's a `cookiecutter` template [oshino-cookiecutter](https://github.com/CodersOfTheNight/oshino-cookiecutter)
+which generates all boilerplate code, you just need to modify it.
+
+Already existing custom agents are mostly listed here: [Third party Agents](https://github.com/CodersOfTheNight/oshino/blob/master/docs/thirdparty.md)
+
+Contributing to the Core
+========================
+It should mostly consist of reporting/solving issues. 
+
+General Guidelines
+===================
+- Code must pass `flake8` (PEP8) coding style test
+- All code (with few exceptions) must be non-blocking - asyncio compatible
+- New parameters should be covered in documentation section

README.md

@@ -56,3 +56,7 @@ Documentation about additional agents can be found [here](docs/thirdparty.md)
 More documentation
 ==================
 More documentation can be found under [docs](docs/index.md) directory
+
+Contributing
+============
+Refer to [CONTRIBUTING.md](CONTRIBUTING.md)

docs/config.md

@@ -6,6 +6,8 @@ Base Config
 - `riemann.transport` - set transport protocol for riemann (default: TCPTransport)
 - `loglevel` - level of logging (default: INFO)
 - `sentry-dsn` - sentry address for error reporting (default: `None`)
+- `executor` - allows to define executor class, default one is `concurrent.futures.ThreadPoolExecutor`
+- `executors-count` - set worker count for executor, default: 10
 
 General Config for Agents
 ------------------------

oshino/config.py

@@ -111,6 +111,18 @@ def agents(self):
     def sentry_dsn(self):
         return self._data.get("sentry-dsn", None)
 
+    @property
+    def executors_count(self):
+        return self._data.get("executors-count", 10)
+
+    @property
+    def executor_class(self):
+        raw = self._data.get(
+                "executor",
+                "concurrent.futures.ThreadPoolExecutor"
+        )
+        return dynamic_import(raw)
+
 
 def load(config_file):
     """

oshino/core/heart.py

@@ -86,6 +86,8 @@ async def main_loop(cfg: Config,
     transport = transport_cls(riemann.host, riemann.port)
     client = processor.QClient(transport)
     agents = create_agents(cfg.agents)
+    executor = cfg.executor_class(max_workers=cfg.executors_count)
+    loop.set_default_executor(executor)
 
     init(agents)
 
@@ -105,7 +107,7 @@ async def main_loop(cfg: Config,
                         len(client.queue.events),
                         len(pending))
 
-        processor.flush(client, transport, logger)
+        await processor.flush(client, transport, logger)
         if continue_fn():
             await asyncio.sleep(cfg.interval - int(td), loop=loop)
         else:

oshino/core/processor.py

@@ -1,42 +1,107 @@
+import asyncio
+
+from copy import copy
+from time import sleep
+
+from logbook import Logger
+from queue import Queue
+
 from riemann_client.client import QueuedClient
 
 
+logger = Logger("Processor")
+
+
+class StopEvent(object):
+    pass
+
+
 class AugmentFixture(object):
-    def apply_augment(self, **data):
-        if "service" not in data:
+
+    def apply_augment(self, data):
+        if "service" not in data or hasattr(self, "stopped"):
             return
 
         key = data["service"]
-        if key in self.augments:
-            subscribers = self.augments[key]
+        if key not in self.augments:
+            return
+
+        subscribers = self.augments[key]
+        for sub in subscribers:
+            sub.put_nowait(data)
+
+    def on_stop(self):
+        if hasattr(self, "stopped"):
+            return
+
+        logger.info("Stopping all augments")
+
+        self.stopped = True
+        for _, subscribers in self.augments.items():
+            for sub in subscribers:
+                sub.put_nowait(StopEvent())
+
+        for _, subscribers in self.augments.items():
             for sub in subscribers:
-                sub.send(data)
+                while not sub.empty():
+                    sleep(0.1)
 
 
 class QClient(QueuedClient, AugmentFixture):
     def __init__(self, *args, **kwargs):
         super(QClient, self).__init__(*args, **kwargs)
         self.augments = {}
+        self.tasks = []
 
     def event(self, **data):
-        self.apply_augment(**data)
+        self.tasks.append(self.apply_augment(copy(data)))
         super(QClient, self).event(**data)
 
 
-def flush(client, transport, logger):
-    try:
-        transport.connect()
-        client.flush()
-        transport.disconnect()
-    except ConnectionRefusedError as ce:
-        logger.warn(ce)
+async def flush(client, transport, logger):
+    future = asyncio.Future()
+
+    async def process_async(future):
+        try:
+            transport.connect()
+            client.flush()
+            transport.disconnect()
+
+            future.set_result(True)
+        except ConnectionRefusedError as ce:
+            logger.warn(ce)
 
+            future.set_result(False)
 
-def register_augment(client, key, generator, logger):
+    asyncio.ensure_future(process_async(future))
+    await future
+    return future.result()
+
+
+def register_augment(client, key, augment_fn, logger):
     if key not in client.augments:
         client.augments[key] = []
 
-    next(generator)
+    loop = asyncio.get_event_loop()
+
+    def generator(q):
+        stopped = False
+        while not (stopped and q.empty()):
+            logger.debug("Waiting for event for {0}".format(augment_fn))
+            event = q.get(block=not stopped)
+            if isinstance(event, StopEvent):
+                logger.debug("Stopping event generator")
+                stopped = True
+                continue
+            yield event
+        raise StopIteration
+
+    q = Queue()
+    g = generator(q)
+
+    def execute_in_thread(fn, client, g):
+        fn(client, g)
+
+    loop.run_in_executor(None, execute_in_thread, augment_fn, client, g)
 
-    generator.send(client)
-    client.augments[key].append(generator)
+    client.augments[key].append(q)

pytest.ini

@@ -1,4 +1,6 @@
 [pytest]
 testpaths = tests
 mock_use_standalone_module = true
-
+markers =
+  slow: mark a slow test
+addopts = -m "not webtest"

tests/data/test_config.yml

@@ -1,6 +1,7 @@
 ---
 interval: 10
 loglevel: DEBUG
+executor: concurrent.futures.ProcessPoolExecutor
 riemann:
   host: localhost
   port: 5555

tests/fixtures.py

@@ -1,10 +1,18 @@
+import asyncio
+
+from concurrent.futures import ThreadPoolExecutor
+
 from pytest import fixture
+
 from .mocks import MockClient, MockTransport
 
 
+
 @fixture
-def mock_client():
-    return MockClient()
+def mock_client(request):
+    client = MockClient()
+    request.addfinalizer(client.on_stop)
+    return client
 
 
 @fixture
@@ -15,3 +23,17 @@ def mock_transport():
 @fixture
 def broken_transport():
     return MockTransport(broken=True)
+
+@fixture(scope="session")
+def executor(request):
+    loop = asyncio.get_event_loop()
+    print("Loop: {0}".format(loop))
+    ex = ThreadPoolExecutor(max_workers=3)
+    def on_stop():
+        ex.shutdown(wait=True)
+        loop.close()
+        print("done closing")
+
+    loop.set_default_executor(ex)
+    request.addfinalizer(on_stop)
+    return ex