Gevent implementation of kafka-python

.gitignore

@@ -5,3 +5,4 @@ build
 dist
 MANIFEST
 env
+*~

kafka/client.py

@@ -38,6 +38,9 @@ def __init__(self, hosts, client_id=CLIENT_ID,
         self.topic_partitions = {}   # topic_id -> [0, 1, 2, ...]
         self.load_metadata_for_topics()  # bootstrap with all metadata
 
+    def create_connection(self, host, port):
+        return KafkaConnection(host, port, timeout=self.timeout)
+
     ##################
     #   Private API  #
     ##################
@@ -56,8 +59,7 @@ def _get_conn_for_broker(self, broker):
         Get or create a connection to a broker
         """
         if (broker.host, broker.port) not in self.conns:
-            self.conns[(broker.host, broker.port)] = \
-                KafkaConnection(broker.host, broker.port, timeout=self.timeout)
+            self.conns[(broker.host, broker.port)] = self.create_connection(broker.host, broker.port)
 
         return self._get_conn(broker.host, broker.port)
 

kafka/conn.py

@@ -54,11 +54,9 @@ def __init__(self, host, port, timeout=DEFAULT_SOCKET_TIMEOUT_SECONDS):
         super(KafkaConnection, self).__init__()
         self.host = host
         self.port = port
-        self._sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
-        self._sock.connect((host, port))
+        self._sock = None
         self.timeout = timeout
-        self._sock.settimeout(self.timeout)
-        self._dirty = False
+        self.reinit()
 
     def __repr__(self):
         return "<KafkaConnection host=%s port=%d>" % (self.host, self.port)

kafka/green/__init__.py

@@ -0,0 +1,11 @@
+from kafka import *
+
+from kafka.green.producer import _Producer, _SimpleProducer, _KeyedProducer
+from kafka.green.conn import _KafkaConnection
+from kafka.green.client import _KafkaClient
+
+Producer=_Producer
+SimpleProducer=_SimpleProducer
+KeyedProducer=_KeyedProducer
+KafkaConnection=_KafkaConnection
+KafkaClient=_KafkaClient

kafka/green/client.py

@@ -0,0 +1,17 @@
+from kafka.client import KafkaClient, DEFAULT_SOCKET_TIMEOUT_SECONDS
+
+from .conn import _KafkaConnection
+
+class _KafkaClient(KafkaClient):
+
+    def __init__(self, hosts, client_id=KafkaClient.CLIENT_ID,
+                 timeout=DEFAULT_SOCKET_TIMEOUT_SECONDS):
+        super(_KafkaClient, self).__init__(hosts=hosts, client_id=client_id, timeout=timeout)
+
+    def copy(self):
+        # have to override this since copy.deepcopy cannot serialize
+        # a gevent.socket
+        return _KafkaClient(['{}:{}'.format(entry[0], entry[1]) for entry in self.hosts], self.client_id, self.timeout)
+
+    def create_connection(self, host, port):
+        return _KafkaConnection(host, port, timeout=self.timeout)

kafka/green/conn.py

@@ -0,0 +1,21 @@
+import gevent.socket as socket
+
+from kafka.conn import KafkaConnection
+
+class _KafkaConnection(KafkaConnection):
+    """
+    Gevent version of kafka.KafkaConnection class. Uses
+    gevent.socket instead of socket.socket.
+    """
+    def __init__(self, host, port, timeout=10):
+        super(_KafkaConnection, self).__init__(host, port, timeout)
+
+    def reinit(self):
+        """
+        Re-initialize the socket connection
+        """
+        self.close()
+        self._sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        self._sock.connect((self.host, self.port))
+        self._sock.settimeout(self.timeout)
+        self._dirty = False

kafka/green/producer.py

@@ -0,0 +1,33 @@
+from kafka.producer import Producer, _send_upstream, STOP_ASYNC_PRODUCER, BATCH_SEND_MSG_COUNT, BATCH_SEND_DEFAULT_INTERVAL, SimpleProducer, KeyedProducer
+
+import gevent
+from gevent.queue import Queue
+
+class _ProducerMixin(object):
+
+    def _setup_async(self, batch_send_every_n, batch_send_every_t):
+        if self.async:
+            self.queue = Queue()  # Messages are sent through this queue
+            self.job = gevent.spawn(_send_upstream,
+                                    self.queue,
+                                    self.client.copy(),
+                                    batch_send_every_t,
+                                    batch_send_every_n,
+                                    self.req_acks,
+                                    self.ack_timeout)
+
+    def stop(self, timeout=1):
+        if self.async:
+            self.queue.put((STOP_ASYNC_PRODUCER, None))
+            self.job.join(timeout)
+            if self.job.dead is False:
+                self.job.kill()
+
+class _Producer(_ProducerMixin, Producer):
+    pass
+
+class _SimpleProducer(_ProducerMixin, SimpleProducer):
+    pass
+
+class _KeyedProducer(_ProducerMixin, KeyedProducer):
+    pass

kafka/producer.py

@@ -117,7 +117,9 @@ def __init__(self, client, async=False,
         self.async = async
         self.req_acks = req_acks
         self.ack_timeout = ack_timeout
+        self._setup_async(batch_send_every_n, batch_send_every_t)
 
+    def _setup_async(self, batch_send_every_n, batch_send_every_t):
         if self.async:
             self.queue = Queue()  # Messages are sent through this queue
             self.proc = Process(target=_send_upstream,

setup.py

@@ -26,7 +26,7 @@ def run(self):
     tests_require=["tox"],
     cmdclass={"test": Tox},
 
-    packages=["kafka"],
+    packages=["kafka", "kafka.green"],
 
     author="David Arthur",
     author_email="[email protected]",