Add producer purge method with optional blocking argument (@peteryin21, #548)

Author: peteryin21

confluent_kafka/src/Producer.c

@@ -145,7 +145,7 @@ static void dr_msg_cb (rd_kafka_t *rk, const rd_kafka_message_t *rkm,
                 goto done;
 
 	msgobj = Message_new0(self, rkm);
-	
+
         args = Py_BuildValue("(OO)", ((Message *)msgobj)->error, msgobj);
 
 	Py_DECREF(msgobj);
@@ -532,6 +532,36 @@ static PyObject *Producer_abort_transaction(Handle *self, PyObject *args) {
         Py_RETURN_NONE;
 }
 
+static void *Producer_purge (Handle *self, PyObject *args,
+                                 PyObject *kwargs) {
+        int in_queue = 1;
+        int in_flight = 1;
+        int blocking = 1;
+        int purge_strategy = 0;
+
+        rd_kafka_resp_err_t err;
+        static char *kws[] = { "in_queue", "in_flight", "blocking", NULL};
+
+        if (!PyArg_ParseTupleAndKeywords(args, kwargs, "|bbb", kws, &in_queue, &in_flight, &blocking))
+                return NULL;
+        if (in_queue)
+                purge_strategy = RD_KAFKA_PURGE_F_QUEUE;
+        if (in_flight) 
+                purge_strategy |= RD_KAFKA_PURGE_F_INFLIGHT;
+        if (blocking)
+                purge_strategy |= RD_KAFKA_PURGE_F_NON_BLOCKING;
+
+        err = rd_kafka_purge(self->rk, purge_strategy);
+
+        if (err) {
+                cfl_PyErr_Format(err, "Purge failed: %s", rd_kafka_err2str(err));
+                return NULL;
+        }
+
+        Py_RETURN_NONE;
+}
+
+
 static PyMethodDef Producer_methods[] = {
 	{ "produce", (PyCFunction)Producer_produce,
 	  METH_VARARGS|METH_KEYWORDS,
@@ -597,6 +627,18 @@ static PyMethodDef Producer_methods[] = {
 	  "callbacks may be triggered.\n"
 	  "\n"
 	},
+	{ "purge", (PyCFunction)Producer_purge, METH_VARARGS|METH_KEYWORDS,
+          ".. py:function:: purge([in_queue=True], [in_flight=True], [blocking=True])\n"
+          "\n"
+	  "   Purge messages currently handled by the producer instance.\n"
+	  "   The application will need to call poll() or flush() "
+	  "afterwards to serve the delivery report callbacks of the purged messages."
+	  "\n"
+	  "  :param: bool in_queue: Purge messages from internal queues. By default, true.\n"
+	  "  :param: bool in_flight: Purge messages in flight to or from the broker. By default, true.\n"
+	  "  :param: bool blocking: If set to False, will not wait on background thread queue\n"
+	  "purging to finish. By default, true."
+	},
         { "list_topics", (PyCFunction)list_topics, METH_VARARGS|METH_KEYWORDS,
           list_topics_doc
         },

tests/test_Producer.py

@@ -235,3 +235,38 @@ def test_transaction_api():
     assert ex.value.args[0].retriable() is False
     assert ex.value.args[0].fatal() is False
     assert ex.value.args[0].txn_requires_abort() is False
+
+
+def test_purge():
+    """
+    Verify that when we have a higher message.timeout.ms timeout, we can use purge()
+    to stop waiting for messages and get delivery reports
+    """
+    p = Producer(
+        {"socket.timeout.ms": 10, "error_cb": error_cb, "message.timeout.ms": 30000}
+    )  # 30 seconds
+
+    # Hack to detect on_delivery was called because inner functions can modify nonlocal objects.
+    # When python2 support is dropped, we can use the "nonlocal" keyword instead
+    cb_detector = {"on_delivery_called": False}
+
+    def on_delivery(err, msg):
+        cb_detector["on_delivery_called"] = True
+        # Because we are purging messages, we should see a PURGE_QUEUE kafka error
+        assert err.code() == KafkaError._PURGE_QUEUE
+
+    # Our message won't be delivered, but also won't timeout yet because our timeout is 30s.
+    p.produce(topic="some_topic", value="testing", partition=9, callback=on_delivery)
+    p.flush(0.002)
+    assert not cb_detector["on_delivery_called"]
+
+    # When in_queue set to false, we won't purge the message and get delivery callback
+    p.purge(in_queue=False)
+    p.flush(0.002)
+    assert not cb_detector["on_delivery_called"]
+
+    # When we purge including the queue, the message should have delivered a delivery report
+    # with a PURGE_QUEUE error
+    p.purge()
+    p.flush(0.002)
+    assert cb_detector["on_delivery_called"]