WIP: Proper asyncio implementation using the c++ side

Author: nl-x-net

pulsar/asyncio.py

@@ -23,11 +23,13 @@
 
 import asyncio
 import functools
-from typing import Any
+from typing import Any, Iterable, Optional
 
 import _pulsar
+
 import pulsar
 
+
 class PulsarException(BaseException):
     """
     The exception that wraps the Pulsar error code
@@ -56,6 +58,7 @@ def __str__(self):
         """
         return f'{self._result.value} {self._result.name}'
 
+
 class Producer:
     """
     The Pulsar message producer, used to publish messages on a topic.
@@ -116,6 +119,152 @@ async def close(self) -> None:
         self._producer.close_async(functools.partial(_set_future, future, value=None))
         await future
 
+    async def flush(self):
+        """
+        Flush all the messages buffered in the client and wait until all messages have been successfully persisted.
+
+        Raises
+        ------
+        PulsarException
+        """
+        future = asyncio.get_running_loop().create_future()
+        self._producer.flush_async(functools.partial(_set_future, future, value=None))
+        await future
+
+    @property
+    def is_connected(self) -> bool:
+        """
+        Check if the producer is connected or not.
+        """
+
+        return self._producer.is_connected()
+
+    @property
+    def last_sequence_id(self) -> int:
+        """
+        Get the last sequence id
+        """
+        return self._producer.last_sequence_id()
+
+    @property
+    def name(self) -> str:
+        """
+        Get the name of the producer.
+        """
+        return self._producer.producer_name()
+
+    @property
+    def topic(self) -> str:
+        """
+        Get the topic name of the producer.
+        """
+        return self._producer.topic()
+
+
+class Consumer:
+    def __init__(self, consumer: _pulsar.Consumer) -> None:
+        self._consumer: _pulsar.Consumer = consumer
+
+    async def acknowledge(self, msg: pulsar.Message) -> None:
+        """
+        Acknowledge the reception of a single message.
+        """
+        future = asyncio.get_running_loop().create_future()
+        self._consumer.acknowledge_async(msg, functools.partial(_set_future, future))
+        await future
+
+    async def acknowledge_cumulative(self, msg: pulsar.Message) -> None:
+        """
+        Acknowledge the reception of all the messages in the stream up to (and including) the provided message.
+        """
+        future = asyncio.get_running_loop().create_future()
+        self._consumer.acknowledge_cumulative_async(msg, functools.partial(_set_future, future))
+        await future
+
+    async def negative_acknowledge(self, msg: pulsar.Message) -> None:
+        """
+        Acknowledge the failure to process a single message.
+        """
+        future = asyncio.get_running_loop().create_future()
+        self._consumer.negative_acknowledge_async(msg, functools.partial(_set_future, future))
+        await future
+
+    async def batch_receive(self) -> Iterable[pulsar.Message]:
+        """
+        Batch receiving messages.
+        """
+        future = asyncio.get_running_loop().create_future()
+        self._consumer.batch_receive_async(functools.partial(_set_future, future))
+        return await future
+
+    async def receive(self) -> pulsar.Message:
+        """
+        Receive a single message.
+        """
+        future = asyncio.get_running_loop().create_future()
+
+        self._consumer.receive_async(functools.partial(_set_future, future))
+        return await future
+
+    async def close(self):
+        """
+        Close the consumer.
+        """
+        future = asyncio.get_running_loop().create_future()
+        self._consumer.close_async(functools.partial(_set_future, future, value=None))
+        await future
+
+    async def seek(self, position: int):
+        """
+        Reset the subscription associated with this consumer to a specific message id or publish timestamp. The message id can either be a specific message or represent the first or last messages in the topic. ...
+        """
+        future = asyncio.get_running_loop().create_future()
+        self._consumer.seek_async(position, functools.partial(_set_future, future))
+        await future
+
+    async def unsubscribe(self):
+        """
+        Unsubscribe the current consumer from the topic.
+        """
+        future = asyncio.get_running_loop().create_future()
+        self._consumer.unsubscribe_async(functools.partial(_set_future, future))
+        await future
+
+    def pause_message_listener(self):
+        """
+        Pause receiving messages via the message_listener until resume_message_listener() is called.
+        """
+        self._consumer.pause_message_listener()
+
+    def resume_message_listener(self):
+        """
+        Resume receiving the messages via the message listener. Asynchronously receive all the messages enqueued from the time pause_message_listener() was called.
+        """
+        self._consumer.resume_message_listener()
+
+    def redeliver_unacknowledged_messages(self):
+        """
+        Redelivers all the unacknowledged messages. In failover mode, the request is ignored if the consumer is not active for the given topic. In shared mode, the consumer's messages to be redelivered are distributed across all the connected consumers...
+        """
+        self._consumer.redeliver_unacknowledged_messages()
+
+    @property
+    def last_message_id(self) -> int:
+        return self._consumer.last_message_id
+
+    @property
+    def is_connected(self) -> bool:
+        return self._consumer.is_connected()
+
+    @property
+    def subscription_name(self) -> str:
+        return self._consumer.subscription_name()
+
+    @property
+    def topic(self) -> str:
+        return self._consumer.topic()
+
+
 class Client:
     """
     The asynchronous version of `pulsar.Client`.
@@ -125,7 +274,18 @@ def __init__(self, service_url, **kwargs) -> None:
         """
         See `pulsar.Client.__init__`
         """
-        self._client: _pulsar.Client = pulsar.Client(service_url, **kwargs)._client
+        assert service_url.startswith('pulsar://'), "The service url must start with 'pulsar://'"
+        self._client = pulsar.Client(service_url, **kwargs)._client
+
+    async def subscribe(self, topics: str, subscription_name: str, consumer_type: _pulsar.ConsumerType,
+                        schema: Optional[_pulsar.SchemaInfo] = _pulsar.SchemaInfo(_pulsar.SchemaType.BYTES, "bytes", "")) -> Consumer:
+        conf = _pulsar.ConsumerConfiguration()
+        conf.consumer_type(consumer_type)
+        conf.schema(schema)
+
+        future = asyncio.get_running_loop().create_future()
+        self._client.subscribe_async(topics, subscription_name, conf, functools.partial(_set_future, future))
+        return Consumer(await future)
 
     async def create_producer(self, topic: str) -> Producer:
         """
@@ -163,10 +323,20 @@ async def close(self) -> None:
         self._client.close_async(functools.partial(_set_future, future, value=None))
         await future
 
+    async def get_topic_partitions(self, topic: str):
+        future = asyncio.get_running_loop().create_future()
+        self._client.get_partitions_for_topic_async(topic, functools.partial(_set_future, future))
+        return await future
+
+    def shutdown(self) -> None:
+        self._client.shutdown()
+
+
 def _set_future(future: asyncio.Future, result: _pulsar.Result, value: Any):
     def complete():
         if result == _pulsar.Result.Ok:
             future.set_result(value)
         else:
             future.set_exception(PulsarException(result))
+
     future.get_loop().call_soon_threadsafe(complete)

src/client.cc

@@ -41,6 +41,12 @@ Consumer Client_subscribe(Client& client, const std::string& topic, const std::s
         [&](SubscribeCallback callback) { client.subscribeAsync(topic, subscriptionName, conf, callback); });
 }
 
+void Client_subscribeAsync(Client& client, const std::string& topic, const std::string& subscriptionName,
+                          const ConsumerConfiguration& conf, SubscribeCallback callback) {
+    py::gil_scoped_release release;
+    client.subscribeAsync(topic, subscriptionName, conf, callback);
+}
+
 Consumer Client_subscribe_topics(Client& client, const std::vector<std::string>& topics,
                                  const std::string& subscriptionName, const ConsumerConfiguration& conf) {
     return waitForAsyncValue<Consumer>(
@@ -86,6 +92,7 @@ void export_client(py::module_& m) {
         .def("create_producer", &Client_createProducer)
         .def("create_producer_async", &Client_createProducerAsync)
         .def("subscribe", &Client_subscribe)
+        .def("subscribe_async", &Client_subscribeAsync)
         .def("subscribe_topics", &Client_subscribe_topics)
         .def("subscribe_pattern", &Client_subscribe_pattern)
         .def("create_reader", &Client_createReader)

src/consumer.cc

@@ -19,6 +19,7 @@
 #include "utils.h"
 
 #include <pulsar/Consumer.h>
+#include <pybind11/functional.h>
 #include <pybind11/pybind11.h>
 #include <pybind11/stl.h>
 
@@ -32,6 +33,11 @@ Message Consumer_receive(Consumer& consumer) {
     return waitForAsyncValue<Message>([&](ReceiveCallback callback) { consumer.receiveAsync(callback); });
 }
 
+void Consumer_receiveAsync(Consumer& consumer, ReceiveCallback callback) {
+      py::gil_scoped_acquire acquire;
+      consumer.receiveAsync(callback);
+}
+
 Message Consumer_receive_timeout(Consumer& consumer, int timeoutMs) {
     Message msg;
     Result res;
@@ -54,6 +60,33 @@ void Consumer_acknowledge(Consumer& consumer, const Message& msg) {
     waitForAsyncResult([&](ResultCallback callback) { consumer.acknowledgeAsync(msg, callback); });
 }
 
+void Consumer_acknowledgeAsync(Consumer& consumer, const Message& msg, py::object callback) {
+    // Capture the callback by value to ensure it's safe to use in the async context
+    auto py_callback = std::make_shared<py::object>(callback);
+
+    // Call the Pulsar asynchronous acknowledge method
+    consumer.acknowledgeAsync(msg, [py_callback](pulsar::Result result) {
+        // Acquire the Python GIL before calling any Python code
+        py::gil_scoped_acquire acquire;
+
+        try {
+            if (result == pulsar::ResultOk) {
+                py::print("result ok");
+                // Call the Python callback with None if the acknowledgment was successful
+                (*py_callback)(pulsar::ResultOk, py::none());
+                py::print("callback called");
+            } else {
+                py::print("error");
+                // Raise a Python exception for failure
+                PyErr_SetString(PyExc_Exception, "AcknowledgeAsync failed");
+                (*py_callback)(pulsar::ResultOk, py::none());  // Or pass some error object to indicate failure
+            }
+        } catch (const py::error_already_set& e) {
+            throw py::error_already_set();
+        }
+    });
+}
+
 void Consumer_acknowledge_message_id(Consumer& consumer, const MessageId& msgId) {
     waitForAsyncResult([&](ResultCallback callback) { consumer.acknowledgeAsync(msgId, callback); });
 }
@@ -81,6 +114,11 @@ void Consumer_close(Consumer& consumer) {
     waitForAsyncResult([&consumer](ResultCallback callback) { consumer.closeAsync(callback); });
 }
 
+void Consumer_closeAsync(Consumer& consumer, ResultCallback callback){
+    py::gil_scoped_acquire acquire;
+    consumer.closeAsync(callback);
+}
+
 void Consumer_pauseMessageListener(Consumer& consumer) { CHECK_RESULT(consumer.pauseMessageListener()); }
 
 void Consumer_resumeMessageListener(Consumer& consumer) { CHECK_RESULT(consumer.resumeMessageListener()); }
@@ -116,14 +154,17 @@ void export_consumer(py::module_& m) {
         .def("unsubscribe", &Consumer_unsubscribe)
         .def("receive", &Consumer_receive)
         .def("receive", &Consumer_receive_timeout)
+        .def("receive_async", &Consumer_receiveAsync)
         .def("batch_receive", &Consumer_batch_receive)
         .def("acknowledge", &Consumer_acknowledge)
         .def("acknowledge", &Consumer_acknowledge_message_id)
+        .def("acknowledge_async", &Consumer_acknowledgeAsync)
         .def("acknowledge_cumulative", &Consumer_acknowledge_cumulative)
         .def("acknowledge_cumulative", &Consumer_acknowledge_cumulative_message_id)
         .def("negative_acknowledge", &Consumer_negative_acknowledge)
         .def("negative_acknowledge", &Consumer_negative_acknowledge_message_id)
         .def("close", &Consumer_close)
+        .def("close_async", &Consumer_closeAsync)
         .def("pause_message_listener", &Consumer_pauseMessageListener)
         .def("resume_message_listener", &Consumer_resumeMessageListener)
         .def("redeliver_unacknowledged_messages", &Consumer::redeliverUnacknowledgedMessages)