[KIP-932] Add share consumer commit workflow (#2241)

Author: k-raina

src/confluent_kafka/cimpl.pyi

@@ -284,6 +284,7 @@ class Message:
         timestamp: Optional[Tuple[int, int]] = ...,
         latency: Optional[float] = ...,
         leader_epoch: Optional[int] = ...,
+        delivery_count: Optional[int] = ...,
     ) -> None: ...
     def topic(self) -> Optional[str]: ...
     def partition(self) -> Optional[int]: ...
@@ -295,6 +296,7 @@ class Message:
     def timestamp(self) -> Tuple[int, int]: ...  # (timestamp_type, timestamp)
     def latency(self) -> Optional[float]: ...
     def leader_epoch(self) -> Optional[int]: ...
+    def delivery_count(self) -> Optional[int]: ...
     def set_headers(self, headers: HeadersType) -> None: ...
     def set_key(self, key: Any) -> None: ...
     def set_value(self, value: Any) -> None: ...
@@ -343,13 +345,9 @@ class Producer:
             Producer({'bootstrap.servers': 'localhost:9092'})
         """
         ...
+
     @overload
-    def __init__(
-        self,
-        config: Dict[str, Any],
-        /,
-        **kwargs: Any
-    ) -> None:
+    def __init__(self, config: Dict[str, Any], /, **kwargs: Any) -> None:
         """
         Create Producer with configuration dict and additional keyword arguments.
         Keyword arguments override values in the config dict.
@@ -363,6 +361,7 @@ class Producer:
             Producer({'bootstrap.servers': 'localhost'}, enable_idempotence=True)
         """
         ...
+
     @overload
     def __init__(self, **config: Any) -> None:
         """
@@ -376,6 +375,7 @@ class Producer:
             Producer(bootstrap_servers='localhost:9092')
         """
         ...
+
     def produce(
         self,
         topic: str,
@@ -426,13 +426,9 @@ class Consumer:
             Consumer({'bootstrap.servers': 'localhost', 'group.id': 'mygroup'})
         """
         ...
+
     @overload
-    def __init__(
-        self,
-        config: dict[str, Any],
-        /,
-        **kwargs: Any
-    ) -> None:
+    def __init__(self, config: dict[str, Any], /, **kwargs: Any) -> None:
         """
         Create Consumer with configuration dict and additional keyword arguments.
         Keyword arguments override values in the config dict.
@@ -446,6 +442,7 @@ class Consumer:
             Consumer({'bootstrap.servers': 'localhost'}, group_id='mygroup')
         """
         ...
+
     @overload
     def __init__(self, **config: Any) -> None:
         """
@@ -459,6 +456,7 @@ class Consumer:
             Consumer(bootstrap_servers='localhost', group_id='mygroup')
         """
         ...
+
     def subscribe(
         self,
         topics: List[str],
@@ -482,6 +480,7 @@ class Consumer:
         Message and offsets omitted, asynchronous.
         """
         ...
+
     @overload
     def commit(
         self,
@@ -492,6 +491,7 @@ class Consumer:
         Message and offsets omitted, synchronous.
         """
         ...
+
     @overload
     def commit(
         self,
@@ -503,6 +503,7 @@ class Consumer:
         Message specified, asynchronous.
         """
         ...
+
     @overload
     def commit(
         self,
@@ -514,17 +515,19 @@ class Consumer:
         Message specified, synchronous.
         """
         ...
+
     @overload
     def commit(
-            self,
-            *,
-            offsets: List[TopicPartition],
-            asynchronous: Literal[True] = ...,
+        self,
+        *,
+        offsets: List[TopicPartition],
+        asynchronous: Literal[True] = ...,
     ) -> None:
         """
         Offsets specified, asynchronous.
         """
         ...
+
     @overload
     def commit(
         self,
@@ -536,6 +539,7 @@ class Consumer:
         Offsets specified, synchronous
         """
         ...
+
     def get_watermark_offsets(
         self, partition: TopicPartition, timeout: float = -1, cached: bool = False
     ) -> Tuple[int, int]: ...
@@ -560,6 +564,7 @@ class Consumer:
 
 class ShareConsumer:
     """Share Consumer for queue-like message consumption (KIP-932)."""
+
     @overload
     def __init__(self, config: Dict[str, Any]) -> None: ...
     @overload
@@ -579,6 +584,12 @@ class ShareConsumer:
     def acknowledge_offset(
         self, topic: str, partition: int, offset: int, ack_type: AcknowledgeType = ...
     ) -> None: ...
+    # TODO KIP-932: Java's share-consumer commit returns a map keyed by
+    # TopicIdPartition (topic name + topic UUID + partition). Python uses
+    # the existing TopicPartition (no UUID) for now. Add a TopicIdPartition
+    # class once the interface is finalized.
+    def commit_sync(self, timeout: float = 60) -> Dict[TopicPartition, Optional[KafkaError]]: ...
+    def commit_async(self) -> None: ...
     def close(self) -> None: ...
     def __enter__(self) -> "ShareConsumer": ...
     def __exit__(self, exc_type: Any, exc_value: Any, exc_traceback: Any) -> Optional[bool]: ...

src/confluent_kafka/src/ShareConsumer.c

@@ -494,6 +494,102 @@ static PyObject *ShareConsumer_acknowledge_offset(ShareConsumerHandle *self,
 }
 
 
+/**
+ * @brief Synchronously commit pending acknowledgements.
+ *
+ * Implicit mode auto-converts the records returned by the latest poll() to
+ * ACCEPT inside librdkafka before sending. Blocks until all broker replies
+ * arrive or the timeout expires.
+ *
+ * @returns dict mapping TopicPartition -> None on success or KafkaError on
+ *          per-partition failure. Empty dict when no acknowledgements are
+ *          pending.
+ */
+static PyObject *ShareConsumer_commit_sync(ShareConsumerHandle *self,
+                                           PyObject *args,
+                                           PyObject *kwargs) {
+        /* TODO KIP-932: check if kwargs is needed */
+        /* TODO KIP-932: pick the right default timeout */
+        double tmout                             = 60.0;
+        rd_kafka_error_t *error                  = NULL;
+        rd_kafka_topic_partition_list_t *c_parts = NULL;
+        PyObject *result                         = NULL;
+        CallState cs;
+        static char *kws[] = {"timeout", NULL};
+
+        if (!self->rkshare) {
+                PyErr_SetString(PyExc_RuntimeError,
+                                ERR_MSG_SHARE_CONSUMER_CLOSED);
+                goto err;
+        }
+
+        if (!PyArg_ParseTupleAndKeywords(args, kwargs, "|d", kws, &tmout))
+                goto err;
+
+        CallState_begin(&self->base, &cs);
+        error = rd_kafka_share_commit_sync(self->rkshare, cfl_timeout_ms(tmout),
+                                           &c_parts);
+        if (!CallState_end(&self->base, &cs))
+                goto err;
+
+        if (error) {
+                cfl_PyErr_from_error_destroy(error);
+                error = NULL;
+                goto err;
+        }
+
+        /* TODO KIP-932: c_parts shouldn't be NULL here, drop once librdkafka
+         * guarantees it */
+        if (!c_parts)
+                return PyDict_New();
+
+        result = c_parts_to_dict_topic_partition_to_error(c_parts);
+        rd_kafka_topic_partition_list_destroy(c_parts);
+        return result;
+
+err:
+        if (c_parts)
+                rd_kafka_topic_partition_list_destroy(c_parts);
+        if (error)
+                rd_kafka_error_destroy(error);
+        return NULL;
+}
+
+
+/**
+ * @brief Asynchronously commit pending acknowledgements.
+ *
+ * Returns immediately; broker results are delivered via the
+ * share_acknowledgement_commit_cb (when configured).
+ */
+static PyObject *ShareConsumer_commit_async(ShareConsumerHandle *self,
+                                            PyObject *ignore) {
+        rd_kafka_error_t *error;
+        CallState cs;
+
+        if (!self->rkshare) {
+                PyErr_SetString(PyExc_RuntimeError,
+                                ERR_MSG_SHARE_CONSUMER_CLOSED);
+                return NULL;
+        }
+
+        CallState_begin(&self->base, &cs);
+        error = rd_kafka_share_commit_async(self->rkshare);
+        if (!CallState_end(&self->base, &cs)) {
+                if (error)
+                        rd_kafka_error_destroy(error);
+                return NULL;
+        }
+
+        if (error) {
+                cfl_PyErr_from_error_destroy(error);
+                return NULL;
+        }
+
+        Py_RETURN_NONE;
+}
+
+
 /**
  * @brief Close the share consumer.
  */
@@ -674,6 +770,47 @@ static PyMethodDef ShareConsumer_methods[] = {
      "  :raises RuntimeError: if called on a closed share consumer.\n"
      "\n"},
 
+    {"commit_sync", (PyCFunction)ShareConsumer_commit_sync,
+     METH_VARARGS | METH_KEYWORDS,
+     ".. py:function:: commit_sync([timeout=60])\n"
+     "\n"
+     "  Synchronously commit pending acknowledgements and block until the\n"
+     "  broker responds or the timeout expires.\n"
+     "\n"
+     "  In implicit acknowledgement mode "
+     "(``share.acknowledgement.mode=implicit``,\n"
+     "  the default), all records returned by the previous :py:func:`poll` "
+     "call\n"
+     "  are auto-converted to ACCEPT before being sent. In explicit mode, "
+     "only\n"
+     "  records previously passed to :py:func:`acknowledge` /\n"
+     "  :py:func:`acknowledge_offset` are sent.\n"
+     "\n"
+     "  :param float timeout: Maximum time to block (seconds). Default: 60.\n"
+     "                        Pass -1 for infinite.\n"
+     "  :returns: Dict mapping TopicPartition to None on success or "
+     "KafkaError\n"
+     "           on per-partition failure. Empty dict when no\n"
+     "           acknowledgements are pending.\n"
+     "  :rtype: dict(TopicPartition, KafkaError | None)\n"
+     "  :raises KafkaException: on error\n"
+     "  :raises RuntimeError: if called on a closed share consumer\n"
+     "  :raises TypeError: if timeout is not a float\n"
+     "\n"},
+
+    {"commit_async", (PyCFunction)ShareConsumer_commit_async, METH_NOARGS,
+     ".. py:function:: commit_async()\n"
+     "\n"
+     "  Asynchronously commit pending acknowledgements. Returns immediately;\n"
+     "  broker results are delivered via the configured\n"
+     "  ``share_acknowledgement_commit_cb`` callback.\n"
+     "\n"
+     "  :returns: None\n"
+     "  :raises KafkaException: on error\n"
+     "  :raises RuntimeError: if called on a closed share consumer\n"
+     "  :raises TypeError: if any arguments are passed\n"
+     "\n"},
+
     {"close", (PyCFunction)ShareConsumer_close, METH_NOARGS,
      ".. py:function:: close()\n"
      "\n"