Repository: Parsely/pykafka
Pull request #685 (2017-05-20 10:48:55)

Set and get the message timestamp, if the broker allows it.

emmett9001: @maparent This looks great at first glance. I'll have to test it out, but this is a very welcome pull request. Thanks for your contribution!

Related to #526 
maparent: Thank you!
BTW, my comment was a bit of a misnomer; I can get existing timestamps, and set the protocol number so librdkafka can set them itself, but I do not yet set it in non-rdkafka situation. I'm actually working on this, but I'd be curious to have an opinion: should I set the timestamp whenever the kafka version allows it?

(What I like about this option is the ability to set a timestamp different from now(), which is unfortunately impossible with librdkafka.)

An issue with that is that the protocol `X.pack_into` methods are not aware of the Producer, and hence of the library version. I'm thinking of adding a "protocol_version" argument to the message constructor (would default to 0).  Another option is to have a default value of None for the timestamp, and the producer would have to explicitly set `timestamp=utcnow()` if it wants to use protocol v1, but that seems more ad hoc. (Final option is to be introduce awareness of producer at the pack_into level, but that is a very heavy-handed change.)
Thoughts, preferences?

emmett9001: @maparent I think adding `protocol_version` to `Message.__init__` is the cleanest way to handle this.
maparent: Good. Do you want me to do this in this PR, or as a separate one? (Might not be today either way, but would be within a week at the most.)
emmett9001: Doing it in this PR makes sense. Once that's in, I'll give this a detailed review and leave comments, if any.
emmett9001: These changes look great so far - once this functionality has been added to the pure-python `Producer`, we should be able to get this merged.

Patch:
diff --git a/pykafka/cluster.py b/pykafka/cluster.py
index 1d4cc11..6c0aa7f 100644
--- a/pykafka/cluster.py
+++ b/pykafka/cluster.py
@@ -199,6 +199,7 @@ class Cluster(object):
         self._ssl_config = ssl_config
         self._zookeeper_connect = zookeeper_hosts
         self._max_connection_retries = 3
+        self._max_connection_retries_offset_mgr = 8
         self._broker_version = broker_version
         if ':' in self._source_address:
             self._source_port = int(self._source_address.split(':')[1])
@@ -402,11 +403,12 @@ class Cluster(object):
         """
         log.info("Attempting to discover offset manager for consumer group '%s'",
                  consumer_group)
-        for i in range(self._max_connection_retries):
+        max_connection_retries = self._max_connection_retries_offset_mgr
+        for i in range(max_connection_retries):
+            if i > 0:
+                log.debug("Retrying offset manager discovery")
+            time.sleep(i * 2)
             for broker in itervalues(self.brokers):
-                if i > 0:
-                    log.debug("Retrying offset manager discovery")
-                time.sleep(i * 2)
 
                 req = GroupCoordinatorRequest(consumer_group)
                 future = broker.handler.request(req)
@@ -414,11 +416,11 @@ class Cluster(object):
                     res = future.get(GroupCoordinatorResponse)
                 except GroupCoordinatorNotAvailable:
                     log.error('Error discovering offset manager.')
-                    if i == self._max_connection_retries - 1:
+                    if i == max_connection_retries - 1:
                         raise
                 except SocketDisconnectedError:
                     log.warning("Socket disconnected during offset manager discovery")
-                    if i == self._max_connection_retries - 1:
+                    if i == max_connection_retries - 1:
                         raise
                     self.update()
                 else:
diff --git a/pykafka/producer.py b/pykafka/producer.py
index 50b3cd9..781a89e 100644
--- a/pykafka/producer.py
+++ b/pykafka/producer.py
@@ -25,6 +25,7 @@ import struct
 import sys
 import threading
 import weakref
+from pkg_resources import parse_version
 
 from six import reraise
 
@@ -44,6 +45,7 @@ from .partitioners import random_partitioner
 from .protocol import Message, ProduceRequest
 from .utils.compat import iteritems, itervalues, Empty
 from .utils.error_handlers import valid_int
+from .utils import msg_protocol_version
 
 log = logging.getLogger(__name__)
 
@@ -156,6 +158,7 @@ class Producer(object):
         :type auto_start: bool
         """
         self._cluster = cluster
+        self._protocol_version = msg_protocol_version(cluster._broker_version)
         self._topic = topic
         self._partitioner = partitioner
         self._compression = compression
@@ -300,7 +303,7 @@ class Producer(object):
                 for queue_reader in queue_readers:
                     queue_reader.join()
 
-    def produce(self, message, partition_key=None):
+    def produce(self, message, partition_key=None, timestamp=None):
         """Produce a message.
 
         :param message: The message to produce (use None to send null)
@@ -325,6 +328,8 @@ class Producer(object):
         msg = Message(value=message,
                       partition_key=partition_key,
                       partition_id=partition_id,
+                      timestamp=timestamp,
+                      protocol_version=self._protocol_version,
                       # We must pass our thread-local Queue instance directly,
                       # as results will be written to it in a worker thread
                       delivery_report_q=self._delivery_reports.queue)
diff --git a/pykafka/protocol.py b/pykafka/protocol.py
index 313a587..e0ad8ae 100644
--- a/pykafka/protocol.py
+++ b/pykafka/protocol.py
@@ -59,6 +59,9 @@ import logging
 import struct
 from collections import defaultdict, namedtuple
 from zlib import crc32
+from datetime import datetime
+import types
+from six import integer_types
 
 from .common import CompressionType, Message
 from .exceptions import ERROR_CODES, MessageSizeTooLarge
@@ -155,7 +158,6 @@ class Message(Message, Serializable):
     :ivar partition_id: The id of the partition to which this message belongs
     :ivar delivery_report_q: For use by :class:`pykafka.producer.Producer`
     """
-    MAGIC = 0
 
     __slots__ = [
         "compression_type",
@@ -165,7 +167,9 @@ class Message(Message, Serializable):
         "partition_id",
         "partition",
         "produce_attempt",
-        "delivery_report_q"
+        "delivery_report_q",
+        "protocol_version",
+        "timestamp"
     ]
 
     def __init__(self,
@@ -175,11 +179,16 @@ class Message(Message, Serializable):
                  offset=-1,
                  partition_id=-1,
                  produce_attempt=0,
+                 protocol_version=0,
+                 timestamp=None,
                  delivery_report_q=None):
         self.compression_type = compression_type
         self.partition_key = partition_key
         self.value = value
         self.offset = offset
+        if timestamp is None and protocol_version > 0:
+            timestamp = datetime.utcnow()
+        self.set_timestamp(timestamp)
         # this is set on decode to expose it to clients that use the protocol
         # implementation but not the consumer
         self.partition_id = partition_id
@@ -188,6 +197,8 @@ class Message(Message, Serializable):
         self.produce_attempt = produce_attempt
         # delivery_report_q is used by the producer
         self.delivery_report_q = delivery_report_q
+        assert protocol_version in (0, 1)
+        self.protocol_version = protocol_version
 
     def __len__(self):
         size = 4 + 1 + 1 + 4 + 4
@@ -195,18 +206,26 @@ class Message(Message, Serializable):
             size += len(self.value)
         if self.partition_key is not None:
             size += len(self.partition_key)
+        if self.protocol_version > 0 and self.timestamp:
+            size += 8
         return size
 
     @classmethod
     def decode(self, buff, msg_offset=-1, partition_id=-1):
-        fmt = 'iBBYY'
-        response = struct_helpers.unpack_from(fmt, buff, 0)
-        crc, _, attr, key, val = response
+        (crc, protocol_version) = struct_helpers.unpack_from('iB', buff, 0)
+        offset = 5
+        timestamp = 0
+        if protocol_version > 0:
+            (timestamp,) = struct_helpers.unpack_from('Q', buff, offset)
+            offset += 8
+        (attr, key, val) = struct_helpers.unpack_from('BYY', buff, offset)
         # TODO: Handle CRC failure
         return Message(val,
                        partition_key=key,
                        compression_type=attr,
                        offset=msg_offset,
+                       protocol_version=protocol_version,
+                       timestamp=timestamp,
                        partition_id=partition_id)
 
     def pack_into(self, buff, offset):
@@ -218,21 +237,51 @@ class Message(Message, Serializable):
         :param offset: The offset to start the write at
         """
         # NB a length of 0 means an empty string, whereas -1 means null
+        # Assuming a CreateTime timestamp, not a LogAppendTime.
         len_key = -1 if self.partition_key is None else len(self.partition_key)
         len_value = -1 if self.value is None else len(self.value)
-        fmt = '!BBi%dsi%ds' % (max(len_key, 0), max(len_value, 0))
-        args = (self.MAGIC,
+        protocol_version = self.protocol_version
+        # Only actually use protocol 1 if timestamp is defined.
+        if self.protocol_version == 1 and self.timestamp:
+            fmt = '!BBQi%dsi%ds' % (max(len_key, 0), max(len_value, 0))
+        else:
+            protocol_version = 0
+            fmt = '!BBi%dsi%ds' % (max(len_key, 0), max(len_value, 0))
+        args = [protocol_version,
                 self.compression_type,
                 len_key,
                 self.partition_key or b"",
                 len_value,
-                self.value or b"")
+                self.value or b""]
+        if protocol_version > 0:
+            args.insert(2, int(self.timestamp))
         struct.pack_into(fmt, buff, offset + 4, *args)
         fmt_size = struct.calcsize(fmt)
         data = buffer(buff[(offset + 4):(offset + 4 + fmt_size)])
         crc = crc32(data) & 0xffffffff
         struct.pack_into('!I', buff, offset, crc)
 
+    @property
+    def timestamp_dt(self):
+        """Get the timestamp as a datetime, if valid"""
+        if self.timestamp > 0:
+            # Assuming a unix epoch
+            return datetime.utcfromtimestamp(self.timestamp / 1000.0)
+
+    @timestamp_dt.setter
+    def timestamp_dt(self, dt):
+        """Set the timestamp from a datetime object"""
+        self.timestamp = int(
+            1000 * (dt - datetime(1970, 1, 1)).total_seconds())
+
+    def set_timestamp(self, ts):
+        if isinstance(ts, integer_types + (float, type(None))):
+            self.timestamp = ts
+        elif isinstance(ts, datetime):
+            self.timestamp_dt = ts
+        else:
+            raise RuntimeError()
+
 
 class MessageSet(Serializable):
     """Representation of a set of messages in Kafka
@@ -301,7 +350,9 @@ class MessageSet(Serializable):
             compressed = compression.encode_snappy(buffer(uncompressed))
         else:
             raise TypeError("Unknown compression: %s" % self.compression_type)
-        return Message(compressed, compression_type=self.compression_type)
+        protocol_version = max((m.protocol_version for m in self._messages))
+        return Message(compressed, compression_type=self.compression_type,
+                       protocol_version=protocol_version)
 
     @classmethod
     def decode(cls, buff, partition_id=-1):
diff --git a/pykafka/rdkafka/_rd_kafkamodule.c b/pykafka/rdkafka/_rd_kafkamodule.c
index 5d7b051..97255a4 100644
--- a/pykafka/rdkafka/_rd_kafkamodule.c
+++ b/pykafka/rdkafka/_rd_kafkamodule.c
@@ -563,6 +563,14 @@ Producer_delivery_report_callback(rd_kafka_t *rk,
             logger, "exception", "s", "Failure in delivery callback");
         Py_XDECREF(res);
         PyErr_Clear();
+    } else {
+        // Set timestamp to the one created by librdkafka
+        rd_kafka_timestamp_type_t timestamp_type;
+        int64_t timestamp = rd_kafka_message_timestamp(rkmessage, &timestamp_type);
+        if (timestamp_type != RD_KAFKA_TIMESTAMP_NOT_AVAILABLE) {
+            PyObject* timestamp_o = PyLong_FromUnsignedLongLong(timestamp);
+            PyObject_SetAttrString(message, "timestamp", timestamp_o);
+        }
     }
 
     Py_DECREF(message);  /* We INCREF'd this in Producer_produce() */
@@ -966,17 +974,23 @@ Consumer_consume(RdkHandle *self, PyObject *args)
 
     if (rkmessage->err == RD_KAFKA_RESP_ERR_NO_ERROR) {
         /* Build a pykafka.protocol.Message */
+        rd_kafka_timestamp_type_t timestamp_type;
+        int64_t timestamp = rd_kafka_message_timestamp(rkmessage, &timestamp_type);
+        if (timestamp_type == RD_KAFKA_TIMESTAMP_NOT_AVAILABLE) {
+            timestamp = 0;
+        }
 #if PY_MAJOR_VERSION >= 3
-        const char *format = "{s:y#,s:y#,s:l,s:L}";
+        const char *format = "{s:y#,s:y#,s:l,s:L,s:L}";
 #else
-        const char *format = "{s:s#,s:s#,s:l,s:L}";
+        const char *format = "{s:s#,s:s#,s:l,s:L,s:L}";
 #endif
         kwargs = Py_BuildValue(
             format,
             "value", rkmessage->payload, rkmessage->len,
             "partition_key", rkmessage->key, rkmessage->key_len,
             "partition_id", (long)rkmessage->partition,
-            "offset", (PY_LONG_LONG)rkmessage->offset);
+            "offset", (PY_LONG_LONG)rkmessage->offset,
+            "timestamp", (PY_LONG_LONG)timestamp);
         if (! kwargs) goto cleanup;
         empty_args = PyTuple_New(0);
         if (! empty_args) goto cleanup;
diff --git a/pykafka/rdkafka/producer.py b/pykafka/rdkafka/producer.py
index bc9352f..11f6adf 100644
--- a/pykafka/rdkafka/producer.py
+++ b/pykafka/rdkafka/producer.py
@@ -109,7 +109,7 @@ class RdKafkaProducer(Producer):
         # marked 'P' or '*') that appear in librdkafka/CONFIGURATION.md should
         # be listed below, in either `conf` or `topic_conf`, even if we do not
         # set them and they are commented out.
-
+        ver10 = parse_version(self._broker_version) >= parse_version("0.10.0")
         conf = {  # destination: rd_kafka_conf_set
             "client.id": "pykafka.rdkafka",
             # Handled via rd_kafka_brokers_add instead:
@@ -163,11 +163,12 @@ class RdKafkaProducer(Producer):
 
             # Report successful and failed messages so we know to dealloc them
             "delivery.report.only.error": "false",
+            "api.version.request": ver10,
             # "dr_cb"
             # "dr_msg_cb"  # gets set in _rd_kafka module
         }
         # broker.version.fallback is incompatible with >-0.10
-        if parse_version(self._broker_version) < parse_version("0.10.0"):
+        if not ver10:
             conf["broker.version.fallback"] = self._broker_version
         conf.update(helpers.rdk_ssl_config(self._cluster))
 
diff --git a/pykafka/rdkafka/simple_consumer.py b/pykafka/rdkafka/simple_consumer.py
index f7f5161..3483038 100644
--- a/pykafka/rdkafka/simple_consumer.py
+++ b/pykafka/rdkafka/simple_consumer.py
@@ -189,7 +189,7 @@ class RdKafkaSimpleConsumer(SimpleConsumer):
         # marked 'C' or '*') that appear in librdkafka/CONFIGURATION.md should
         # be listed below, in either `conf` or `topic_conf`, even if we do not
         # set them and they are commented out.
-
+        ver10 = parse_version(self._broker_version) >= parse_version("0.10.0")
         conf = {  # destination: rd_kafka_conf_set
             "client.id": "pykafka.rdkafka",
             # Handled via rd_kafka_brokers_add instead:
@@ -245,6 +245,7 @@ class RdKafkaSimpleConsumer(SimpleConsumer):
             "fetch.message.max.bytes": self._fetch_message_max_bytes,
             "fetch.min.bytes": self._fetch_min_bytes,
             "fetch.error.backoff.ms": self._fetch_error_backoff_ms,
+            "api.version.request": ver10,
 
             # We're outsourcing message fetching, but not offset management or
             # consumer rebalancing to librdkafka.  Thus, consumer-group id
@@ -253,7 +254,7 @@ class RdKafkaSimpleConsumer(SimpleConsumer):
             ##"group.id"
             }
         # broker.version.fallback is incompatible with >-0.10
-        if parse_version(self._broker_version) < parse_version("0.10.0"):
+        if not ver10:
             conf["broker.version.fallback"] = self._broker_version
         conf.update(helpers.rdk_ssl_config(self._cluster))
 
diff --git a/pykafka/utils/__init__.py b/pykafka/utils/__init__.py
index f1a8384..441d599 100644
--- a/pykafka/utils/__init__.py
+++ b/pykafka/utils/__init__.py
@@ -13,6 +13,7 @@ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License.
 """
+from pkg_resources import parse_version
 
 
 class Serializable(object):
@@ -25,3 +26,9 @@ class Serializable(object):
     def pack_into(self, buff, offset):
         """Pack serialized bytes into buff starting at offset ``offset``"""
         raise NotImplementedError()
+
+
+def msg_protocol_version(broker_version):
+    if parse_version(broker_version) >= parse_version("0.10.0"):
+        return 1
+    return 0


Test patch:
diff --git a/tests/pykafka/test_protocol.py b/tests/pykafka/test_protocol.py
index b54916a..8ac3e19 100644
--- a/tests/pykafka/test_protocol.py
+++ b/tests/pykafka/test_protocol.py
@@ -226,6 +226,8 @@ class TestFetchAPI(unittest2.TestCase):
             'partition_id': 0,
             'produce_attempt': 0,
             'delivery_report_q': None,
+            'timestamp': 0,
+            'protocol_version': 0,
             'partition': None
         }, {
             'partition_key': b'test_key',
@@ -235,6 +237,8 @@ class TestFetchAPI(unittest2.TestCase):
             'partition_id': 0,
             'produce_attempt': 0,
             'delivery_report_q': None,
+            'timestamp': 0,
+            'protocol_version': 0,
             'partition': None
         }, {
             'partition_key': None,
@@ -244,6 +248,8 @@ class TestFetchAPI(unittest2.TestCase):
             'partition_id': 0,
             'produce_attempt': 0,
             'delivery_report_q': None,
+            'timestamp': 0,
+            'protocol_version': 0,
             'partition': None
         }]