Kafka shim

blueox/__init__.py

@@ -30,20 +30,28 @@
 from .errors import Error
 from .logger import LogHandler
 from .timer import timeit
+from .recorders import kafka
 
 log = logging.getLogger(__name__)
 
 
 def configure(host, port, recorder=None):
     """Initialize blueox
 
-    This instructs the blueox system where to send it's logging data. If blueox is not configured, log data will
+    This instructs the blueox system where to send its logging data. If blueox is not configured, log data will
     be silently dropped.
 
     Currently we support logging through the network (and the configured host and port) to a blueoxd instances, or
     to the specified recorder function
     """
-    if recorder:
+    override_kafka_recorder = os.getenv('BLUEOX_OVERRIDE_KAFKA_RECORDER', 0)
+
+    if int(override_kafka_recorder) == 1:
+        log.info("Kafka override set, using kafka recorder")
+        host = ports.default_kafka_host()
+        kafka.init(host)
+        _context_mod._recorder_function = kafka.send
+    elif recorder:
         _context_mod._recorder_function = recorder
     elif host and port:
         network.init(host, port)

blueox/ports.py

@@ -41,3 +41,13 @@ def default_control_host(host=None):
 def default_collect_host(host=None):
     default_host = os.environ.get(ENV_VAR_COLLECT_HOST, DEFAULT_HOST)
     return _default_host(host, default_host, DEFAULT_COLLECT_PORT)
+
+
+# For consistency, we'll abstract kafka connections in the same way
+ENV_VAR_KAFKA_HOST = 'BLUEOX_KAFKA_HOST'
+DEFAULT_KAFKA_PORT = 9092
+
+
+def default_kafka_host(host=None):
+    default_host = os.environ.get(ENV_VAR_KAFKA_HOST, DEFAULT_HOST)
+    return _default_host(host, default_host, DEFAULT_KAFKA_PORT)

blueox/recorders/kafka.py

@@ -0,0 +1,102 @@
+# -*- coding: utf-8 -*-
+"""
+blueox.kafka
+~~~~~~~~
+
+This module provides the interface into Kafka
+
+:copyright: (c) 2018 by Aaron Biller??
+:license: ISC, see LICENSE for more details.
+
+"""
+from __future__ import absolute_import
+
+import atexit
+import logging
+import msgpack
+import threading
+
+from kafka import KafkaProducer
+
+from blueox import utils
+
+log = logging.getLogger(__name__)
+
+# If we have pending outgoing messages, this is how long we'll wait after
+# being told to exit.
+LINGER_SHUTDOWN_MSECS = 2000
+
+
+threadLocal = threading.local()
+
+# Context can be shared between threads
+_kafka_hosts = None
+
+
+def init(host):
+    global _kafka_hosts
+
+    _kafka_hosts = host
+
+
+def _thread_connect():
+    if _kafka_hosts and not getattr(threadLocal, 'kp', None):
+        threadLocal.kp = KafkaProducer(bootstrap_servers=_kafka_hosts)
+
+
+def _serialize_context(context):
+    context_dict = context.to_dict()
+    for key in ('host', 'type'):
+        if len(context_dict.get(key, "")) > 64:
+            raise ValueError("Value too long: %r" % key)
+
+    context_dict = {
+        k: v.encode('utf-8') if isinstance(v, unicode)
+        else v for k, v in context_dict.items()
+    }
+
+    try:
+        context_data = msgpack.packb(context_dict)
+    except TypeError:
+        try:
+            # If we fail to serialize our context, we can try again with an
+            # enhanced packer (it's slower though)
+            context_data = msgpack.packb(context_dict,
+                                         default=utils.msgpack_encode_default)
+        except TypeError:
+            log.exception("Serialization failure (not fatal, dropping data)")
+
+            # One last try after dropping the body
+            context_dict['body'] = None
+            context_data = msgpack.packb(context_dict)
+
+    return context_data
+
+
+def send(context):
+    _thread_connect()
+
+    try:
+        context_data = _serialize_context(context)
+    except Exception:
+        log.exception("Failed to serialize context")
+        return
+
+    if _kafka_hosts and threadLocal.kp is not None:
+        try:
+            log.debug("Sending msg")
+            threadLocal.kp.send('events', context_data)
+        except Exception:
+            log.exception("Failed during publish to kafka.")
+    else:
+        log.info("Skipping sending event %s", context.name)
+
+
+def close():
+    if getattr(threadLocal, 'kp', None):
+        threadLocal.kp.flush()
+        threadLocal.kp.close(timeout=LINGER_SHUTDOWN_MSECS)
+        threadLocal.kp = None
+
+
+atexit.register(close)

requirements.txt

@@ -4,3 +4,4 @@ pyflakes
 tornado==3.2 
 boto 
 yapf
+kafka-python

tests/ports_test.py

@@ -1,5 +1,8 @@
 import os
-from testify import *
+from testify import (
+    TestCase,
+    assert_equal,
+    teardown)
 
 from blueox import ports
 
@@ -71,3 +74,26 @@ def test_env_port(self):
         os.environ['BLUEOX_HOST'] = 'master:123'
         host = ports.default_collect_host()
         assert_equal(host, "master:123")
+
+
+class DefaultKafkaHost(TestCase):
+    @teardown
+    def clear_env(self):
+        try:
+            del os.environ['BLUEOX_KAFKA_HOST']
+        except KeyError:
+            pass
+
+    def test_emtpy(self):
+        host = ports.default_kafka_host()
+        assert_equal(host, '127.0.0.1:9092')
+
+    def test_env(self):
+        os.environ['BLUEOX_KAFKA_HOST'] = 'local.svc.team-me.aws.jk8s'
+        host = ports.default_kafka_host()
+        assert_equal(host, 'local.svc.team-me.aws.jk8s:9092')
+
+    def test_env_port(self):
+        os.environ['BLUEOX_KAFKA_HOST'] = 'local.svc.team-me.aws.jk8s:9002'
+        host = ports.default_kafka_host()
+        assert_equal(host, 'local.svc.team-me.aws.jk8s:9002')

tests/recorders/kafka_test.py

@@ -0,0 +1,120 @@
+import os
+import random
+import decimal
+import datetime
+
+import msgpack
+from testify import (
+    TestCase,
+    setup,
+    teardown,
+    assert_equal)
+
+from blueox import default_configure
+from blueox import utils
+from blueox import network
+from blueox import context
+from blueox.recorders import kafka
+
+
+class MockKafkaProducer(object):
+    last_topic = None
+    last_data = None
+    close_timeout = None
+
+    def __call__(self, bootstrap_servers=None):
+        self.bootstrap_servers = bootstrap_servers
+        return self
+
+    def send(self, topic, data):
+        self.last_topic = topic
+        self.last_data = data
+
+    def flush(self):
+        pass
+
+    def close(self, timeout=None):
+        self.close_timeout = timeout
+
+
+class KafkaOverrideTestCase(TestCase):
+    @teardown
+    def clear_env(self):
+        try:
+            del os.environ['BLUEOX_OVERRIDE_KAFKA_RECORDER']
+        except KeyError:
+            pass
+
+    def test_configure_no_override(self):
+        default_configure()
+        assert_equal(context._recorder_function, network.send)
+
+    def test_configure_override(self):
+        os.environ['BLUEOX_OVERRIDE_KAFKA_RECORDER'] = '1'
+        default_configure()
+        assert_equal(context._recorder_function, kafka.send)
+
+
+class KafkaSendTestCase(TestCase):
+    @setup
+    def build_context(self):
+        self.context = context.Context('test', 1)
+
+    @setup
+    def init_kafka(self):
+        self.port = random.randint(30000, 40000)
+        kafka.init('127.0.0.1:{}'.format(self.port))
+
+    @setup
+    def configure_kafka(self):
+        context._recorder_function = kafka.send
+        self.kp = MockKafkaProducer()
+        kafka.KafkaProducer = self.kp
+
+    @teardown
+    def unconfigure_kafka(self):
+        context._recorder_function = None
+
+    def test(self):
+        with self.context:
+            self.context.set('foo', True)
+            self.context.set('bar.baz', 10.0)
+
+        data = msgpack.unpackb(self.kp.last_data)
+        assert_equal(self.kp.last_topic, 'events')
+        assert_equal(data['id'], 1)
+        assert_equal(data['type'], 'test')
+        assert_equal(utils.get_deep(data['body'], "bar.baz"), 10.0)
+
+        kafka.close()
+        assert_equal(self.kp.close_timeout, kafka.LINGER_SHUTDOWN_MSECS)
+
+
+class SerializeContextTestCase(TestCase):
+    @setup
+    def build_context(self):
+        self.context = context.Context('test', 1)
+
+    def test_types(self):
+        with self.context:
+            self.context.set('decimal_value', decimal.Decimal("6.66"))
+            self.context.set('date_value', datetime.date(2013, 12, 10))
+            self.context.set('datetime_value', datetime.datetime(2013, 12, 10, 12, 12, 12))
+
+        context_data = kafka._serialize_context(self.context)
+        data = msgpack.unpackb(context_data)
+        assert_equal(data['body']['decimal_value'], "6.66")
+        assert_equal(data['body']['date_value'], "2013-12-10")
+        assert_equal(
+            datetime.datetime.fromtimestamp(float(data['body']['datetime_value'])),
+            datetime.datetime(2013, 12, 10, 12, 12, 12))
+
+    def test_exception(self):
+        with self.context:
+            self.context.set('value', Exception('hello'))
+
+        context_data = kafka._serialize_context(self.context)
+        data = msgpack.unpackb(context_data)
+
+        # The serialization should fail, but that just means we don't have any data.
+        assert_equal(data['body'], None)