using new kafka-python interfaces

Author: sibiryakov

docs/source/topics/frontera-settings.rst

@@ -577,15 +577,14 @@ KAFKA_LOCATION
 
 Hostname and port of kafka broker, separated with :. Can be a string with hostname:port pair separated with commas(,).
 
-.. setting:: KAFKA_CODEC
+.. setting:: KAFKA_COMPRESSION
 
-KAFKA_CODEC
-___________
+KAFKA_COMPRESSION
+-----------------
 
-Default:: ``CODEC_NONE``
+Default:: ``None``
 
-Kafka protocol compression codec, see kafka-python documentation for more details. Please use symbols from kafka-python
-package.
+Kafka's producer compression type string, see `kafka-python documentation`_ for more details.
 
 .. setting:: FRONTIER_GROUP
 
@@ -636,8 +635,6 @@ Default: ``frontier-score``
 Kafka topic used for :term:`scoring log` stream.
 
 
-
-
 Default settings
 ================
 
@@ -681,3 +678,5 @@ Values::
 
     EVENT_LOG_MANAGER = 'frontera.logger.events.EventLogManager'
 
+
+.. _`kafka-python documentation`: http://kafka-python.readthedocs.io/en/1.1.1/apidoc/KafkaProducer.html

frontera/contrib/backends/partitioners.py

@@ -18,6 +18,9 @@ def partition_by_hash(self, value, partitions):
         idx = value % size
         return partitions[idx]
 
+    def __call__(self, key, all_partitions, available):
+        return self.partition(key, all_partitions)
+
 
 class FingerprintPartitioner(Partitioner):
     def partition(self, key, partitions=None):
@@ -26,4 +29,7 @@ def partition(self, key, partitions=None):
         digest = unhexlify(key[0:2] + key[5:7] + key[10:12] + key[15:17])
         value = unpack("<I", digest)
         idx = value[0] % len(partitions)
-        return partitions[idx]
+        return partitions[idx]
+
+    def __call__(self, key, all_partitions, available):
+        return self.partition(key, all_partitions)

frontera/contrib/messagebus/kafkabus.py

@@ -1,21 +1,22 @@
 # -*- coding: utf-8 -*-
 from __future__ import absolute_import
-from frontera.core.messagebus import BaseMessageBus, BaseSpiderLogStream, BaseSpiderFeedStream, \
-    BaseStreamConsumer, BaseScoringLogStream, BaseStreamProducer
 
-from kafka import KafkaClient, SimpleConsumer, KeyedProducer as KafkaKeyedProducer, SimpleProducer as KafkaSimpleProducer
-from kafka.common import BrokerResponseError, MessageSizeTooLargeError
+from logging import getLogger
+
+from kafka import KafkaClient, SimpleConsumer
+from kafka import KafkaConsumer, KafkaProducer, TopicPartition
+from kafka.common import BrokerResponseError
 from kafka.protocol import CODEC_NONE
 
 from frontera.contrib.backends.partitioners import FingerprintPartitioner, Crc32NamePartitioner
 from frontera.contrib.messagebus.kafka import OffsetsFetcher
-from logging import getLogger
-from time import sleep
+from frontera.core.messagebus import BaseMessageBus, BaseSpiderLogStream, BaseSpiderFeedStream, \
+    BaseStreamConsumer, BaseScoringLogStream, BaseStreamProducer
 
 logger = getLogger("kafkabus")
 
 
-class Consumer(BaseStreamConsumer):
+class DeprecatedConsumer(BaseStreamConsumer):
     """
     Used in DB and SW worker. SW consumes per partition.
     """
@@ -70,68 +71,75 @@ def get_offset(self):
         return 0
 
 
+class Consumer(BaseStreamConsumer):
+    """
+    Used in DB and SW worker. SW consumes per partition.
+    """
+    def __init__(self, location, topic, group, partition_id):
+        self._location = location
+        self._group = group
+        self._topic = topic
+        self._partition_ids = [partition_id] if partition_id is not None else None
+
+        self._consumer = KafkaConsumer(
+            bootstrap_servers=self._location,
+            group_id=self._group,
+            max_partition_fetch_bytes=10485760)
+        if self._partition_ids:
+            self._consumer.assign([TopicPartition(self._topic, pid) for pid in self._partition_ids])
+        else:
+            self._consumer.subscribe(self._topic)
+
+    def get_messages(self, timeout=0.1, count=1):
+        while True:
+            try:
+                batch = self._consumer.poll(timeout_ms=timeout)
+                for _, records in batch.iteritems():
+                    for record in records:
+                        yield record.value
+            except Exception, err:
+                logger.warning("Error %s" % err)
+            finally:
+                break
+
+    def get_offset(self):
+        return 0
+
+
 class SimpleProducer(BaseStreamProducer):
-    def __init__(self, location, topic, codec):
+    def __init__(self, location, topic, compression):
         self._location = location
         self._topic = topic
-        self._codec = codec
+        self._compression = compression
         self._create()
 
     def _create(self):
-        self._client = KafkaClient(self._location)
-        self._producer = KafkaSimpleProducer(self._client, codec=self._codec)
+        self._producer = KafkaProducer(bootstrap_servers=self._location, retries=5,
+                                       compression_type=self._compression)
 
     def send(self, key, *messages):
-        self._producer.send_messages(self._topic, *messages)
+        for msg in messages:
+            self._producer.send(self._topic, value=msg)
 
-    def get_offset(self, partition_id):
-        # Kafka has it's own offset management
-        raise KeyError
+    def flush(self):
+        self._producer.flush()
 
 
 class KeyedProducer(BaseStreamProducer):
-    def __init__(self, location, topic_done, partitioner_cls, codec):
-        self._prod = None
+    def __init__(self, location, topic_done, partitioner, compression):
         self._location = location
         self._topic_done = topic_done
-        self._partitioner_cls = partitioner_cls
-        self._codec = codec
-
-    def _connect_producer(self):
-        if self._prod is None:
-            try:
-                self._client = KafkaClient(self._location)
-                self._prod = KafkaKeyedProducer(self._client, partitioner=self._partitioner_cls, codec=self._codec)
-            except BrokerResponseError:
-                self._prod = None
-                logger.warning("Could not connect producer to Kafka server")
-                return False
-        return True
+        self._partitioner = partitioner
+        self._compression = compression
+        self._producer = KafkaProducer(bootstrap_servers=self._location, partitioner=partitioner, retries=5,
+                                       compression_type=self._compression)
 
     def send(self, key, *messages):
-        success = False
-        max_tries = 5
-        if self._connect_producer():
-            n_tries = 0
-            while not success and n_tries < max_tries:
-                try:
-                    self._prod.send_messages(self._topic_done, key, *messages)
-                    success = True
-                except MessageSizeTooLargeError, e:
-                    logger.error(str(e))
-                    break
-                except BrokerResponseError:
-                    n_tries += 1
-                    logger.warning(
-                        "Could not send message. Try {0}/{1}".format(
-                            n_tries, max_tries)
-                    )
-                    sleep(1.0)
-        return success
-
-    def get_offset(self, partition_id):
-        # Kafka has it's own offset management
-        raise KeyError
+        for msg in messages:
+            self._producer.send(self._topic_done, key=key, value=msg)
+
+    def flush(self):
+        self._producer.flush()
 
 
 class SpiderLogStream(BaseSpiderLogStream):
@@ -140,10 +148,12 @@ def __init__(self, messagebus):
         self._db_group = messagebus.general_group
         self._sw_group = messagebus.sw_group
         self._topic_done = messagebus.topic_done
-        self._codec = messagebus.codec
+        self._compression_type = messagebus.compression_type
+        self._partitions = messagebus.spider_log_partitions
 
     def producer(self):
-        return KeyedProducer(self._location, self._topic_done, FingerprintPartitioner, self._codec)
+        return KeyedProducer(self._location, self._topic_done, FingerprintPartitioner(self._partitions),
+                             self._compression_type)
 
     def consumer(self, partition_id, type):
         """
@@ -153,7 +163,7 @@ def consumer(self, partition_id, type):
         :return:
         """
         group = self._sw_group if type == 'sw' else self._db_group
-        return Consumer(self._location, self._topic_done, group, partition_id)
+        return DeprecatedConsumer(self._location, self._topic_done, group, partition_id)
 
 
 class SpiderFeedStream(BaseSpiderFeedStream):
@@ -164,10 +174,11 @@ def __init__(self, messagebus):
         self._max_next_requests = messagebus.max_next_requests
         self._hostname_partitioning = messagebus.hostname_partitioning
         self._offset_fetcher = OffsetsFetcher(self._location, self._topic, self._general_group)
-        self._codec = messagebus.codec
+        self._compression_type = messagebus.compression_type
+        self._partitions = messagebus.spider_feed_partitions
 
     def consumer(self, partition_id):
-        return Consumer(self._location, self._topic, self._general_group, partition_id)
+        return DeprecatedConsumer(self._location, self._topic, self._general_group, partition_id)
 
     def available_partitions(self):
         partitions = []
@@ -178,22 +189,23 @@ def available_partitions(self):
         return partitions
 
     def producer(self):
-        partitioner = Crc32NamePartitioner if self._hostname_partitioning else FingerprintPartitioner
-        return KeyedProducer(self._location, self._topic, partitioner, self._codec)
+        partitioner = Crc32NamePartitioner(self._partitions) if self._hostname_partitioning \
+            else FingerprintPartitioner(self._partitions)
+        return KeyedProducer(self._location, self._topic, partitioner, self._compression_type)
 
 
 class ScoringLogStream(BaseScoringLogStream):
     def __init__(self, messagebus):
         self._topic = messagebus.topic_scoring
         self._group = messagebus.general_group
         self._location = messagebus.kafka_location
-        self._codec = messagebus.codec
+        self._compression_type = messagebus.compression_type
 
     def consumer(self):
-        return Consumer(self._location, self._topic, self._group, partition_id=None)
+        return DeprecatedConsumer(self._location, self._topic, self._group, partition_id=None)
 
     def producer(self):
-        return SimpleProducer(self._location, self._topic, self._codec)
+        return SimpleProducer(self._location, self._topic, self._compression_type)
 
 
 class MessageBus(BaseMessageBus):
@@ -206,9 +218,10 @@ def __init__(self, settings):
         self.spider_partition_id = settings.get('SPIDER_PARTITION_ID')
         self.max_next_requests = settings.MAX_NEXT_REQUESTS
         self.hostname_partitioning = settings.get('QUEUE_HOSTNAME_PARTITIONING')
-        codec = settings.get('KAFKA_CODEC')
-        self.codec = codec if codec else CODEC_NONE
+        self.compression_type = settings.get('KAFKA_COMPRESSION')
         self.kafka_location = settings.get('KAFKA_LOCATION')
+        self.spider_log_partitions = settings.get('SPIDER_LOG_PARTITIONS')
+        self.spider_feed_partitions = settings.get('SPIDER_FEED_PARTITIONS')
 
     def spider_log(self):
         return SpiderLogStream(self)

frontera/settings/default_settings.py

@@ -58,7 +58,7 @@
 #--------------------------------------------------------
 # Kafka
 #--------------------------------------------------------
-KAFKA_CODEC = None
+KAFKA_COMPRESSION = None
 OUTGOING_TOPIC = "frontier-todo"
 INCOMING_TOPIC = "frontier-done"
 SCORING_TOPIC = "frontier-score"

frontera/tests/test_message_bus.py

@@ -46,6 +46,7 @@ def spider_log_activity(self, messages):
                 self.sp_sl_p.send(sha1(str(randint(1, 1000))), 'http://helloworld.com/way/to/the/sun/' + str(0))
             else:
                 self.sp_sl_p.send(sha1(str(randint(1, 1000))), 'http://way.to.the.sun' + str(0))
+        self.sp_sl_p.flush()
 
     def spider_feed_activity(self):
         sf_c = 0
@@ -78,7 +79,7 @@ def db_activity(self, messages):
                 self.db_sf_p.send("newhost", "http://newhost/new/url/to/crawl")
             else:
                 self.db_sf_p.send("someotherhost", "http://newhost223/new/url/to/crawl")
-
+        self.db_sf_p.flush()
         return (sl_c, us_c)
 
 
@@ -108,9 +109,9 @@ def test_zmq_message_bus():
 class KafkaMessageBusTestCase(KafkaIntegrationTestCase):
     @classmethod
     def setUpClass(cls):
-        logging.basicConfig(level=logging.DEBUG)
-        logger = logging.getLogger("frontera.tests.kafka_utils.service")
-        logger.addHandler(logging.StreamHandler())
+        #logging.basicConfig(level=logging.DEBUG)
+        #logger = logging.getLogger("frontera.tests.kafka_utils.service")
+        #logger.addHandler(logging.StreamHandler())
         if not os.environ.get('KAFKA_VERSION'):
             return
 
@@ -127,7 +128,7 @@ def tearDownClass(cls):
         cls.server.close()
         cls.zk.close()
 
-    def test_kafka_message_bus(self):
+    def test_kafka_message_bus_integration(self):
         """
         Test MessageBus with default settings, IPv6 and Star as ZMQ_ADDRESS
         """
@@ -145,4 +146,4 @@ def test_kafka_message_bus(self):
         tester.spider_log_activity(64)
         assert tester.sw_activity() == 64
         assert tester.db_activity(128) == (64, 32)
-        assert tester.spider_feed_activity() == 128
+        assert tester.spider_feed_activity() == 128