kafka-python 3.0.0__py3-none-any.whl

kafka/net/wakeup_notifier.py

@@ -0,0 +1,72 @@
+import weakref
+
+from kafka.future import Future
+
+
+class WakeupNotifier:
+    """await wakeup(timeout_secs) when either ``timeout_secs`` elapses or
+    notify() is called -- whichever first. The notifier is safe to call
+    from any thread (it routes through call_soon_threadsafe).
+
+    Level-triggered: notify() arriving while no one is awaiting is latched
+    and consumed by the next ``__call__``. This closes a lost-wakeup race
+    where a caller's state mutation (e.g. ``cluster._need_update = True``)
+    and its ``notify()`` happen between another task's pre-await state
+    check and its ``await self._wakeup(...)``. Without latching, the
+    notification arrives at the IO thread before the task has registered
+    a future to signal, and the task would sleep for the full timeout
+    despite work being ready.
+
+    Used by the metadata refresh loop to sleep on its TTL while remaining
+    interruptible by external callers (e.g. KafkaProducer / KafkaConsumer
+    invoking cluster.request_update() from another thread).
+    """
+    def __init__(self, net):
+        self._net = weakref.proxy(net)
+        self._fut = None
+        # Set by ``_wakeup`` when no awaiter is registered; consumed by the
+        # next ``__call__``. All accesses run on the IO thread (notify
+        # routes through call_soon_threadsafe), so no lock is needed.
+        self._pending = False
+
+    def _wakeup(self):
+        if self._fut is not None and not self._fut.is_done:
+            self._fut.success(None)
+        else:
+            self._pending = True
+
+    async def __call__(self, timeout_secs=None):
+        if self._fut is not None:
+            raise RuntimeError('Concurrent access to WakeupNotifier!')
+        self._fut = Future()
+        if self._pending:
+            self._pending = False
+            self._fut.success(None)
+        if timeout_secs is not None:
+            try:
+                timer = self._net.call_later(timeout_secs, self._wakeup)
+            except ReferenceError:
+                return
+        else:
+            timer = None
+        try:
+            await self._fut
+        finally:
+            self._fut = None
+            if timer is not None and not timer.is_done:
+                try:
+                    self._net.unschedule(timer)
+                except (ValueError, RuntimeError, ReferenceError):
+                    pass
+
+    def notify(self):
+        # Always queue _wakeup on the IO thread. Skipping the queue when
+        # ``self._fut is None`` would re-introduce the lost-wakeup race:
+        # the check could pass before another thread enters ``__call__``
+        # and creates the future. Routing through the IO thread is one
+        # call_soon_threadsafe (~microseconds) and lets ``_wakeup`` decide
+        # under single-threaded semantics whether to signal or latch.
+        try:
+            self._net.call_soon_threadsafe(self._wakeup)
+        except ReferenceError:
+            pass

kafka/partitioner/__init__.py

@@ -0,0 +1,8 @@
+from .abc import Partitioner
+from .default import DefaultPartitioner, murmur2
+from .sticky import StickyPartitioner
+
+
+__all__ = [
+    'Partitioner', 'DefaultPartitioner', 'StickyPartitioner', 'murmur2'
+]

kafka/partitioner/abc.py

@@ -0,0 +1,8 @@
+from abc import ABC, abstractmethod
+
+
+class Partitioner(ABC):
+    @abstractmethod
+    def partition(self, topic, key, serialized_key, value, serialized_value, cluster):
+        pass
+

kafka/partitioner/default.py

@@ -0,0 +1,89 @@
+import random
+
+from .abc import Partitioner
+
+
+class DefaultPartitioner(Partitioner):
+    """Default partitioner.
+
+    Hashes key to partition using murmur2 hashing (from java client)
+    If key is None, selects partition randomly from available,
+    or from all partitions if none are currently available
+    """
+    def partition(self, topic, key, serialized_key, value, serialized_value, cluster):
+        if topic not in cluster.topics():
+            raise ValueError("Topic %s not found in ClusterMetadata" % (topic,))
+        all_partitions = sorted(cluster.partitions_for_topic(topic))
+        available = list(cluster.available_partitions_for_topic(topic))
+        if serialized_key is not None:
+            idx = murmur2(serialized_key)
+            idx &= 0x7fffffff
+            idx %= len(all_partitions)
+            return all_partitions[idx]
+        pool = available if available else all_partitions
+        return random.choice(pool)
+
+
+# https://github.com/apache/kafka/blob/0.8.2/clients/src/main/java/org/apache/kafka/common/utils/Utils.java#L244
+def murmur2(data):
+    """Pure-python Murmur2 implementation.
+
+    Based on java client, see org.apache.kafka.common.utils.Utils.murmur2
+
+    Args:
+        data (bytes): opaque bytes
+
+    Returns: MurmurHash2 of data
+    """
+    length = len(data)
+    seed = 0x9747b28c
+    # 'm' and 'r' are mixing constants generated offline.
+    # They're not really 'magic', they just happen to work well.
+    m = 0x5bd1e995
+    r = 24
+
+    # Initialize the hash to a random value
+    h = seed ^ length
+    length4 = length // 4
+
+    for i in range(length4):
+        i4 = i * 4
+        k = ((data[i4 + 0] & 0xff) +
+            ((data[i4 + 1] & 0xff) << 8) +
+            ((data[i4 + 2] & 0xff) << 16) +
+            ((data[i4 + 3] & 0xff) << 24))
+        k &= 0xffffffff
+        k *= m
+        k &= 0xffffffff
+        k ^= (k % 0x100000000) >> r # k ^= k >>> r
+        k &= 0xffffffff
+        k *= m
+        k &= 0xffffffff
+
+        h *= m
+        h &= 0xffffffff
+        h ^= k
+        h &= 0xffffffff
+
+    # Handle the last few bytes of the input array
+    extra_bytes = length % 4
+    if extra_bytes >= 3:
+        h ^= (data[(length & ~3) + 2] & 0xff) << 16
+        h &= 0xffffffff
+    if extra_bytes >= 2:
+        h ^= (data[(length & ~3) + 1] & 0xff) << 8
+        h &= 0xffffffff
+    if extra_bytes >= 1:
+        h ^= (data[length & ~3] & 0xff)
+        h &= 0xffffffff
+        h *= m
+        h &= 0xffffffff
+
+    h ^= (h % 0x100000000) >> 13 # h >>> 13;
+    h &= 0xffffffff
+    h *= m
+    h &= 0xffffffff
+    h ^= (h % 0x100000000) >> 15 # h >>> 15;
+    h &= 0xffffffff
+
+    return h

kafka/partitioner/sticky.py

@@ -0,0 +1,109 @@
+"""KIP-480 sticky partitioner.
+
+Records with a non-None key are hashed to a partition just like
+:class:`~kafka.partitioner.default.DefaultPartitioner`. Records with a
+None key go to a *sticky* partition - i.e. the same partition is reused
+for every null-key record on a topic until KafkaProducer signals that a
+batch has been completed (via :meth:`StickyPartitioner.on_new_batch`),
+at which point a different partition is picked.
+
+The goal is to give the RecordAccumulator larger, denser batches for
+null-key sends so per-batch overhead (CRC, compression, broker
+round-trip) is amortized across more records. Java's benchmark in
+KIP-480 reported substantial throughput/latency improvements over the
+default-random behavior, though kafka-python is unlikely to see similar
+improvements while predominantly CPU-bound on per-record overhead.
+"""
+
+import random
+import threading
+
+from .default import DefaultPartitioner
+
+
+class StickyPartitioner(DefaultPartitioner):
+    """Partitioner that sticks null-key records to one partition per
+    topic until ``on_new_batch`` rotates it.
+
+    Thread-safety: ``_sticky`` mutations are protected by ``_lock`` so
+    concurrent ``send()`` callers can't observe a torn read-modify-write.
+    """
+
+    def __init__(self):
+        self._sticky = {}  # topic -> partition_id
+        self._lock = threading.Lock()
+
+    def partition(self, topic, key, serialized_key, value, serialized_value, cluster):
+        """Choose a partition for the next record.
+
+        Arguments:
+            topic (str): topic to partition on.
+            key (any): Unserialized key.
+            serialized_key (bytes or None): partitioning key.
+            value (any): Unserialized value.
+            serialized_value (bytes or None): serialized value.
+            cluster (ClusterMetadata): metadata for cluster; provides
+                all and available partitions for topic.
+
+        Raises:
+            ValueError: if topic is not in ClusterMetadata
+
+        Returns:
+            int: chosen partition ID.
+        """
+        if topic not in cluster.topics():
+            raise ValueError("Topic %s not found in ClusterMetadata" % (topic,))
+        if serialized_key is not None:
+            return super().partition(topic, key, serialized_key, value, serialized_value, cluster)
+        # Null key: reuse the sticky partition if still valid.
+        with self._lock:
+            partition = self._sticky.get(topic)
+            if partition is not None:
+                all_partitions = sorted(cluster.partitions_for_topic(topic))
+                available = list(cluster.available_partitions_for_topic(topic))
+                if available:
+                    if partition in available:
+                        return partition
+                elif partition in all_partitions:
+                    return partition
+                # Stale (leader unavailable, topic shrunk); fall through to re-pick.
+            return self._pick_sticky_locked(topic, cluster)
+
+    def on_new_batch(self, topic, cluster, prev_partition):
+        """Hook called by ``KafkaProducer`` on the abort-for-new-batch
+        retry path: rotate the sticky for ``topic`` so the next
+        null-key record lands on a different partition.
+
+        Stale events (where another thread already rotated us off
+        ``prev_partition``) are no-ops.
+        """
+        with self._lock:
+            if self._sticky.get(topic) != prev_partition:
+                # Another caller already rotated us; don't override.
+                return
+            self._pick_sticky_locked(topic, cluster, avoid=prev_partition)
+
+    def _pick_sticky_locked(self, topic, cluster, avoid=None):
+        """Pick a new sticky partition for ``topic``. Must be called with
+        ``self._lock`` held. Returns None when the topic is no longer in
+        cluster metadata (caller is expected to no-op in that case)."""
+        all_partitions = cluster.partitions_for_topic(topic)
+        if not all_partitions:
+            return None
+        all_partitions = sorted(all_partitions)
+        available = list(cluster.available_partitions_for_topic(topic) or ())
+        if available:
+            if len(available) == 1:
+                partition = available[0]
+            else:
+                # >= 2 available: pick uniformly, avoiding ``avoid`` if set.
+                candidates = [p for p in available if p != avoid] if avoid is not None else available
+                if not candidates:
+                    candidates = available
+                partition = random.choice(candidates)
+        else:
+            # No partitions are currently available - pick from the full
+            # set without enforcing ``!= avoid``
+            partition = random.choice(all_partitions)
+        self._sticky[topic] = partition
+        return partition

kafka/producer/__init__.py

@@ -0,0 +1,5 @@
+from kafka.producer.kafka import KafkaProducer
+
+__all__ = [
+    'KafkaProducer'
+]

kafka/producer/__main__.py

@@ -0,0 +1,5 @@
+import sys
+
+from kafka.cli.producer import run_cli
+
+sys.exit(run_cli())

kafka/producer/future.py

@@ -0,0 +1,101 @@
+import collections
+import threading
+
+from kafka import errors as Errors
+from kafka.future import Future
+from kafka.util import Timer
+
+
+class FutureProduceResult(Future):
+    __slots__ = ('topic_partition', '_latch')
+
+    def __init__(self, topic_partition):
+        super().__init__()
+        self.topic_partition = topic_partition
+        self._latch = threading.Event()
+
+    def success(self, value):
+        ret = super().success(value)
+        self._latch.set()
+        return ret
+
+    def failure(self, error):
+        ret = super().failure(error)
+        self._latch.set()
+        return ret
+
+    def wait(self, timeout=None):
+        # wait() on python2.6 returns None instead of the flag value
+        return self._latch.wait(timeout) or self._latch.is_set()
+
+
+class FutureRecordMetadata(Future):
+    __slots__ = ('_produce_future', 'args')
+    def __init__(self, produce_future, batch_index, timestamp_ms, checksum, serialized_key_size, serialized_value_size, serialized_header_size):
+        super().__init__()
+        self._produce_future = produce_future
+        # packing args as a tuple is a minor speed optimization
+        self.args = (batch_index, timestamp_ms, checksum, serialized_key_size, serialized_value_size, serialized_header_size)
+        produce_future._add_cb_eb(self._produce_success, self.failure)
+
+    def _produce_success(self, result):
+        offset, produce_timestamp_ms, record_exceptions_fn = result
+
+        # Unpacking from args tuple is minor speed optimization
+        (batch_index, timestamp_ms, checksum,
+         serialized_key_size, serialized_value_size, serialized_header_size) = self.args
+
+        if record_exceptions_fn is not None:
+            self.failure(record_exceptions_fn(batch_index))
+        else:
+            # None is when Broker does not support the API (<0.10) and
+            # -1 is when the broker is configured for CREATE_TIME timestamps
+            if produce_timestamp_ms is not None and produce_timestamp_ms != -1:
+                timestamp_ms = produce_timestamp_ms
+            if offset != -1 and batch_index is not None:
+                offset += batch_index
+            tp = self._produce_future.topic_partition
+            metadata = RecordMetadata(tp[0], tp[1], tp, offset, timestamp_ms,
+                                      checksum, serialized_key_size,
+                                      serialized_value_size, serialized_header_size)
+            self.success(metadata)
+
+    def rebind(self, new_produce_future, new_batch_index):
+        """Rebind this future to a new produce future with a new batch index.
+
+        Used when a batch is split due to MESSAGE_TOO_LARGE. The original
+        FutureRecordMetadata is rebound to the new (smaller) batch's future.
+
+        This must be called from the sender thread while the old produce_future
+        has not been completed. Any user thread blocked in get() on the old
+        produce_future's latch will be woken and will re-wait on the new one.
+        """
+        old_produce_future = self._produce_future
+        self._produce_future = new_produce_future
+        _, timestamp_ms, checksum, sk, sv, sh = self.args
+        self.args = (new_batch_index, timestamp_ms, checksum, sk, sv, sh)
+        new_produce_future._add_cb_eb(self._produce_success, self.failure)
+        # Wake any thread blocked in get() so it re-waits on the new future.
+        # The old produce_future is never completed, so its stale callbacks
+        # (registered in __init__) will never fire.
+        old_produce_future._latch.set()
+
+    def get(self, timeout=None):
+        """Wait for up to timeout seconds for future to complete."""
+        # Loop because rebind() may wake us from the old produce_future's
+        # latch before the record is actually done. A batch may be split
+        # multiple times, so each rebind wakes us and we re-wait on the
+        # (possibly new) _produce_future.
+        timer = Timer(timeout * 1000 if timeout is not None else None)
+        while not self.is_done and not timer.expired:
+            if not self._produce_future.wait(timer.timeout_secs):
+                raise Errors.KafkaTimeoutError(
+                    "Timeout after waiting for %s secs." % (timeout,))
+        if self.failed():
+            raise self.exception # pylint: disable-msg=raising-bad-type
+        return self.value
+
+
+RecordMetadata = collections.namedtuple(
+    'RecordMetadata', ['topic', 'partition', 'topic_partition', 'offset', 'timestamp',
+                       'checksum', 'serialized_key_size', 'serialized_value_size', 'serialized_header_size'])