[KIP-932] Implement acknowledgement RPC (#2239)

Author: k-raina

src/confluent_kafka/__init__.py

@@ -20,6 +20,7 @@
 
 from ._model import Node  # noqa: F401
 from ._model import (
+    AcknowledgeType,
     ConsumerGroupState,
     ConsumerGroupTopicPartitions,
     ConsumerGroupType,
@@ -86,6 +87,7 @@
     "TopicCollection",
     "TopicPartitionInfo",
     "ElectionType",
+    "AcknowledgeType",
 ]
 
 

src/confluent_kafka/_model/__init__.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from enum import Enum
+from enum import Enum, IntEnum
 from typing import List, Optional
 
 from .. import cimpl
@@ -193,3 +193,20 @@ def __lt__(self, other: object) -> bool:
         if not isinstance(other, ElectionType):
             return NotImplemented
         return self.value < other.value
+
+
+class AcknowledgeType(IntEnum):
+    """
+    Share Consumer acknowledgement type used to tell the broker how to
+    handle a polled message in explicit acknowledgement mode.
+
+    Values:
+    -------
+    """
+
+    #: Record was processed successfully — broker will not redeliver it.
+    ACCEPT = cimpl.SHARE_ACKNOWLEDGE_TYPE_ACCEPT
+    #: Could not process — Release it for another delivery attempt
+    RELEASE = cimpl.SHARE_ACKNOWLEDGE_TYPE_RELEASE
+    #: Could not process - Do not release for another delivery attempt
+    REJECT = cimpl.SHARE_ACKNOWLEDGE_TYPE_REJECT

src/confluent_kafka/cimpl.pyi

@@ -45,6 +45,7 @@ except ImportError:
 
 from confluent_kafka.admin._metadata import ClusterMetadata, GroupMetadata
 
+from ._model import AcknowledgeType
 from ._types import HeadersType
 
 # Callback types with proper class references (defined locally to avoid circular imports)
@@ -83,6 +84,7 @@ class KafkaError:
     ILLEGAL_GENERATION: int
     ILLEGAL_SASL_STATE: int
     INCONSISTENT_GROUP_PROTOCOL: int
+    INCONSISTENT_TOPIC_ID: int
     INCONSISTENT_VOTER_SET: int
     INVALID_COMMIT_OFFSET_SIZE: int
     INVALID_CONFIG: int
@@ -567,7 +569,16 @@ class ShareConsumer:
     def subscribe(self, topics: List[str]) -> None: ...
     def unsubscribe(self) -> None: ...
     def subscription(self) -> List[str]: ...
+    # TODO KIP-932: poll() returns List[Message] today. Java returns a
+    # dedicated container object (ConsumerRecords) instead of a list so it
+    # can carry extra metadata alongside the records. Replace List[Message]
+    # with a Messages container class once we have a clear use for that
+    # metadata in Python.
     def poll(self, timeout: float = -1) -> List[Message]: ...
+    def acknowledge(self, message: Message, ack_type: AcknowledgeType = ...) -> None: ...
+    def acknowledge_offset(
+        self, topic: str, partition: int, offset: int, ack_type: AcknowledgeType = ...
+    ) -> None: ...
     def close(self) -> None: ...
     def __enter__(self) -> "ShareConsumer": ...
     def __exit__(self, exc_type: Any, exc_value: Any, exc_traceback: Any) -> Optional[bool]: ...
@@ -786,6 +797,9 @@ OFFSET_SPEC_EARLIEST: int
 OFFSET_SPEC_LATEST: int
 OFFSET_SPEC_MAX_TIMESTAMP: int
 OFFSET_STORED: int
+SHARE_ACKNOWLEDGE_TYPE_ACCEPT: int
+SHARE_ACKNOWLEDGE_TYPE_RELEASE: int
+SHARE_ACKNOWLEDGE_TYPE_REJECT: int
 RESOURCE_ANY: int
 RESOURCE_BROKER: int
 RESOURCE_GROUP: int

src/confluent_kafka/src/AdminTypes.c

@@ -658,6 +658,16 @@ static void AdminTypes_AddObjectsElectionType(PyObject *m) {
                                 RD_KAFKA_ELECTION_TYPE_UNCLEAN);
 }
 
+static void AdminTypes_AddObjectsShareAcknowledgeType(PyObject *m) {
+        /* rd_kafka_share_AcknowledgeType_t */
+        PyModule_AddIntConstant(m, "SHARE_ACKNOWLEDGE_TYPE_ACCEPT",
+                                RD_KAFKA_SHARE_ACKNOWLEDGE_TYPE_ACCEPT);
+        PyModule_AddIntConstant(m, "SHARE_ACKNOWLEDGE_TYPE_RELEASE",
+                                RD_KAFKA_SHARE_ACKNOWLEDGE_TYPE_RELEASE);
+        PyModule_AddIntConstant(m, "SHARE_ACKNOWLEDGE_TYPE_REJECT",
+                                RD_KAFKA_SHARE_ACKNOWLEDGE_TYPE_REJECT);
+}
+
 /**
  * @brief Add Admin types to module
  */
@@ -679,4 +689,5 @@ void AdminTypes_AddObjects(PyObject *m) {
         AdminTypes_AddObjectsIsolationLevel(m);
         AdminTypes_AddObjectsOffsetSpec(m);
         AdminTypes_AddObjectsElectionType(m);
+        AdminTypes_AddObjectsShareAcknowledgeType(m);
 }

src/confluent_kafka/src/ShareConsumer.c

@@ -393,6 +393,107 @@ static PyObject *ShareConsumer_poll(ShareConsumerHandle *self,
 }
 
 
+/**
+ * @brief Acknowledge previously polled message.
+ *
+ * Internally delegates to rd_kafka_share_acknowledge_offset() because the
+ * Python Message object does not retain the underlying rd_kafka_message_t
+ * pointer (Message_new0 copies fields out and destroys the rkm)
+ *
+ * TODO KIP-932: Java splits ack APIs by message kind — successful records
+ * go through acknowledge(message), error/GAP records through
+ * acknowledge_offset(topic, partition, offset), and crossing the wires
+ * throws. NJC clients (this one included) accept either. Revisit once the
+ * Java-vs-NJC alignment is settled.
+ */
+static PyObject *ShareConsumer_acknowledge(ShareConsumerHandle *self,
+                                           PyObject *args,
+                                           PyObject *kwargs) {
+        Message *msg      = NULL;
+        int ack_type      = (int)RD_KAFKA_SHARE_ACKNOWLEDGE_TYPE_ACCEPT;
+        PyObject *uo8     = NULL;
+        const char *topic = NULL;
+        rd_kafka_resp_err_t err;
+        static char *kws[] = {"message", "ack_type", NULL};
+
+        if (!self->rkshare) {
+                PyErr_SetString(PyExc_RuntimeError,
+                                ERR_MSG_SHARE_CONSUMER_CLOSED);
+                return NULL;
+        }
+
+        if (!PyArg_ParseTupleAndKeywords(args, kwargs, "O!|i", kws,
+                                         &MessageType, &msg, &ack_type))
+                return NULL;
+
+        /* Validation (ack_type range, topic, partition, offset) is left to
+         * librdkafka so every failure surfaces as KafkaException. Pass NULL
+         * topic through for the None case rather than raising ValueError. */
+
+        if (msg->topic && msg->topic != Py_None) {
+                topic = cfl_PyUnistr_AsUTF8(msg->topic, &uo8);
+                if (!topic) {
+                        Py_XDECREF(uo8);
+                        return NULL;
+                }
+        }
+
+        err = rd_kafka_share_acknowledge_offset(
+            self->rkshare, topic, msg->partition, msg->offset,
+            (rd_kafka_share_AcknowledgeType_t)ack_type);
+
+        Py_XDECREF(uo8);
+
+        if (err) {
+                cfl_PyErr_Format(err, "Failed to acknowledge message: %s",
+                                 rd_kafka_err2str(err));
+                return NULL;
+        }
+
+        Py_RETURN_NONE;
+}
+
+
+/**
+ * @brief Acknowledge a message by topic/partition/offset directly.
+ */
+static PyObject *ShareConsumer_acknowledge_offset(ShareConsumerHandle *self,
+                                                  PyObject *args,
+                                                  PyObject *kwargs) {
+        const char *topic;
+        int partition;
+        long long offset;
+        int ack_type = (int)RD_KAFKA_SHARE_ACKNOWLEDGE_TYPE_ACCEPT;
+        rd_kafka_resp_err_t err;
+        static char *kws[] = {"topic", "partition", "offset", "ack_type", NULL};
+
+        if (!self->rkshare) {
+                PyErr_SetString(PyExc_RuntimeError,
+                                ERR_MSG_SHARE_CONSUMER_CLOSED);
+                return NULL;
+        }
+
+        if (!PyArg_ParseTupleAndKeywords(args, kwargs, "siL|i", kws, &topic,
+                                         &partition, &offset, &ack_type))
+                return NULL;
+
+        /* ack_type, partition and offset are all validated by
+         * rd_kafka_share_acknowledge_offset()*/
+
+        err = rd_kafka_share_acknowledge_offset(
+            self->rkshare, topic, (int32_t)partition, (int64_t)offset,
+            (rd_kafka_share_AcknowledgeType_t)ack_type);
+
+        if (err) {
+                cfl_PyErr_Format(err, "Failed to acknowledge offset: %s",
+                                 rd_kafka_err2str(err));
+                return NULL;
+        }
+
+        Py_RETURN_NONE;
+}
+
+
 /**
  * @brief Close the share consumer.
  */
@@ -514,6 +615,65 @@ static PyMethodDef ShareConsumer_methods[] = {
      "  :raises KeyboardInterrupt: if Ctrl+C pressed during consumption\n"
      "\n"},
 
+    /* TODO KIP-932: librdkafka error code → Python exception mapping is
+     * provisional. Today the share consumer translates every librdkafka
+     * error code into KafkaException via cfl_PyErr_Format(). Longer term we
+     * want each code to map to the Python exception a user porting from
+     * Java would expect, e.g.
+     *   _INVALID_ARG → ValueError    (matches Java IllegalArgumentException)
+     *   _STATE       → RuntimeError  (matches Java IllegalStateException)
+     * Open question: per-partition broker errors in commit_sync's result
+     * dict (mirrors Java's Map<TopicIdPartition, Optional<...>>) — keep as
+     * KafkaError, or translate as well?
+     *
+     * Revisit holistically once:
+     *   - librdkafka's share-consumer error surface is stable (some codes
+     *     may be redefined as work progresses), and
+     *   - the equivalent translation lands on commit_sync / commit_async /
+     *     ack-callback paths (currently TODO'd separately).
+     */
+    {"acknowledge", (PyCFunction)ShareConsumer_acknowledge,
+     METH_VARARGS | METH_KEYWORDS,
+     ".. py:function:: acknowledge(message, "
+     "[ack_type=AcknowledgeType.ACCEPT])\n"
+     "\n"
+     "  Acknowledge a previously polled message in explicit acknowledgement\n"
+     "  mode. Tells the broker how to handle the record.\n"
+     "\n"
+     "  :param Message message: A message returned by poll().\n"
+     "  :param AcknowledgeType ack_type: ACCEPT (default), RELEASE, or "
+     "REJECT.\n"
+     "  :raises TypeError: if message is not a Message instance or ack_type "
+     "is not an integer.\n"
+     "  :raises KafkaException: if the consumer is not in explicit\n"
+     "                          acknowledgement mode, the message is no "
+     "longer\n"
+     "                          in-flight, ack_type is invalid, or "
+     "message.topic() is None.\n"
+     "  :raises RuntimeError: if called on a closed share consumer.\n"
+     "\n"},
+
+    {"acknowledge_offset", (PyCFunction)ShareConsumer_acknowledge_offset,
+     METH_VARARGS | METH_KEYWORDS,
+     ".. py:function:: acknowledge_offset(topic, partition, offset, "
+     "[ack_type=AcknowledgeType.ACCEPT])\n"
+     "\n"
+     "  Acknowledge a message by topic/partition/offset.\n"
+     "\n"
+     "  :param str topic: Topic name.\n"
+     "  :param int partition: Partition id.\n"
+     "  :param int offset: Offset to acknowledge.\n"
+     "  :param AcknowledgeType ack_type: ACCEPT (default), RELEASE, or "
+     "REJECT.\n"
+     "  :raises TypeError: if topic is not a str, partition/offset are not "
+     "integers, or ack_type is not an integer.\n"
+     "  :raises KafkaException: if the consumer is not in explicit\n"
+     "                          acknowledgement mode, the offset is not\n"
+     "                          in-flight, the offset is a GAP record,\n"
+     "                          or ack_type is invalid.\n"
+     "  :raises RuntimeError: if called on a closed share consumer.\n"
+     "\n"},
+
     {"close", (PyCFunction)ShareConsumer_close, METH_NOARGS,
      ".. py:function:: close()\n"
      "\n"

tests/common/__init__.py

@@ -22,6 +22,7 @@
 import uuid
 
 from confluent_kafka import Consumer, ShareConsumer
+from confluent_kafka.admin import AlterConfigOpType, ConfigEntry, ConfigResource
 
 _GROUP_PROTOCOL_ENV = 'TEST_CONSUMER_GROUP_PROTOCOL'
 
@@ -146,42 +147,71 @@ def unique_id(prefix):
     return f'{prefix}-{uuid.uuid4().hex[:10]}'
 
 
-def drain_share_consumers(consumers, n_expected, timeout_s=20.0, poll_timeout_s=0.5):
-    """Round-robin poll until total non-error messages reach n_expected.
+def set_group_config(kafka_cluster, group_id, name, value):
+    """Set one dynamic group config via incremental_alter_configs."""
+    res = ConfigResource(
+        ConfigResource.Type.GROUP,
+        group_id,
+        incremental_configs=[
+            ConfigEntry(name, str(value), incremental_operation=AlterConfigOpType.SET),
+        ],
+    )
+    for f in kafka_cluster.admin().incremental_alter_configs([res]).values():
+        f.result()
 
-    Returns a list of message lists, one per input consumer, in the same order.
-    Stops early once the expected total is reached, or when timeout_s elapses.
 
-    Tests with N>2 consumers under the suite-wide
-    group.share.record.lock.duration.ms=1000
-    should lower poll_timeout_s so a full round-robin round completes within
-    the 1s lock window — otherwise locks expire before the same consumer's
-    next poll can implicit-ack, and records get redelivered to other
-    consumers (i.e. apparent duplicate delivery).
+def poll_first_batch(consumer, timeout_s=20.0, poll_timeout_s=0.5):
+    """Return the first non-empty batch of non-error messages from a share consumer.
 
-    IMPORTANT: implicit-ack only. This helper assumes share consumers are in
-    implicit-ack mode (the only mode the Python wrapper currently exposes).
-    In implicit mode, the second poll() automatically acknowledges records
-    delivered by the first, so the loop here is safe.
+    Polls repeatedly until at least one non-error message arrives, or until
+    timeout_s elapses. Returns the contents of a single poll() call — does not
+    accumulate across polls and does not acknowledge anything.
 
-    TODO KIP-932: when explicit-ack mode lands in the Python wrapper
-    (ShareConsumer.acknowledge()), update this helper to ack each message as
-    it's drained, otherwise the broker will return INFLIGHT-records errors
-    on the second poll. Same caveat exists in the librdkafka tests.
+    Use this when a test needs to inspect the first records the broker delivers
+    *before* any implicit- or explicit-ack happens. Unlike drain_share_consumers,
+    which loops poll() calls and therefore implicit-acks each prior batch on
+    every iteration, poll_first_batch returns whatever the first non-empty
+    poll() returned and stops, leaving acknowledgement to the caller.
 
-    TODO KIP-932: after the final poll() in the loop below, the last batch of
-    records is never implicitly acknowledged (no subsequent poll to piggyback
-    on). Some tests assume the tail batch is ack'd. Fix once explicit-ack is
-    exposed: emit an explicit ack for the last drained batch before returning.
+    Returns an empty list if timeout_s elapses without any non-error message.
+    """
+    deadline = time.time() + timeout_s
+    while time.time() < deadline:
+        batch = [msg for msg in consumer.poll(timeout=poll_timeout_s) if msg.error() is None]
+        if batch:
+            return batch
+    return []
+
+
+def drain_share_consumers(consumers, total_messages, timeout_s=20.0, poll_timeout_s=0.5, *, ack_type=None):
+    """Round-robin poll consumers until total_messages non-error messages arrive (across all consumers combined).
+
+    Returns one message list per consumer, in input order. Stops at
+    total_messages or timeout_s.
+
+    With N>2 consumers and the suite-wide 1s lock duration, drop poll_timeout_s
+    so each round finishes within the lock window — otherwise records leak to
+    other consumers and look like duplicate deliveries.
+
+    ack_type=None drains in implicit-ack mode (next poll() auto-acks the prior
+    batch; the tail batch is left unacked). Pass an AcknowledgeType to ack each
+    message inline — required when share.acknowledgement.mode=explicit.
+
+    TODO KIP-932: once commit_sync() / commit_async() are exposed on the
+    ShareConsumer binding, the implicit-ack drain (ack_type=None) should
+    still emit a commit on the tail batch so callers don't have to chase
+    leaked records. Add a commit step here once those APIs land.
     """
     received = [[] for _ in consumers]
     deadline = time.time() + timeout_s
     while time.time() < deadline:
         for sc, bucket in zip(consumers, received):
-            for m in sc.poll(timeout=poll_timeout_s):
-                if m.error() is None:
-                    bucket.append(m)
-        if sum(len(b) for b in received) >= n_expected:
+            for msg in sc.poll(timeout=poll_timeout_s):
+                if msg.error() is None:
+                    bucket.append(msg)
+                    if ack_type is not None:
+                        sc.acknowledge(msg, ack_type)
+        if sum(len(bucket) for bucket in received) >= total_messages:
             break
     return received