Handle delivery.report.only.error in Python (confluentinc#84)

edenhill · edenhill · commit 9eaefa57082a · 2017-02-28T21:33:01.000+01:00
diff --git a/confluent_kafka/src/Producer.c b/confluent_kafka/src/Producer.c
@@ -152,6 +152,10 @@ static void dr_msg_cb (rd_kafka_t *rk, const rd_kafka_message_t *rkm,
 		goto done;
 	}
 
+        /* Skip callback if delivery.report.only.error=true */
+        if (self->u.Producer.dr_only_error && !rkm->err)
+                goto done;
+
 	msgobj = Message_new0(self, rkm);
 	
 	args = Py_BuildValue("(OO)",
diff --git a/confluent_kafka/src/confluent_kafka.c b/confluent_kafka/src/confluent_kafka.c
@@ -1124,7 +1124,26 @@ static int producer_conf_set_special (Handle *self, rd_kafka_conf_t *conf,
 		}
 
 		return 1;
-	}
+
+        } else if (!strcmp(name, "delivery.report.only.error")) {
+                /* Since we allocate msgstate for each produced message
+                 * with a callback we can't use delivery.report.only.error
+                 * as-is, as we wouldn't be able to ever free those msgstates.
+                 * Instead we shortcut this setting in the Python client,
+                 * providing the same functionality from dr_msg_cb trampoline.
+                 */
+
+                if (!PyBool_Check(valobj)) {
+                        cfl_PyErr_Format(
+                                RD_KAFKA_RESP_ERR__INVALID_ARG,
+                                "%s requires bool", name);
+                        return -1;
+                }
+
+                self->u.Producer.dr_only_error = valobj == Py_True;
+
+                return 1;
+        }
 
 	return 0; /* Not handled */
 }
diff --git a/confluent_kafka/src/confluent_kafka.h b/confluent_kafka/src/confluent_kafka.h
@@ -138,6 +138,8 @@ typedef struct {
 				const rd_kafka_topic_t *,
 				const void *, size_t, int32_t,
 				void *, void *);  /**< Fallback C partitioner*/
+
+                        int dr_only_error; /**< delivery.report.only.error */
 		} Producer;
 
 		/**
diff --git a/examples/integration_test.py b/examples/integration_test.py
@@ -134,6 +134,62 @@ def verify_producer():
     # Block until all messages are delivered/failed
     p.flush()
 
+    #
+    # Additional isolated tests
+    #
+    test_producer_dr_only_error()
+
+
+
+def test_producer_dr_only_error():
+    """
+    The C delivery.report.only.error configuration property
+    can't be used with the Python client since the Python client
+    allocates a msgstate for each produced message that has a callback,
+    and on success (with delivery.report.only.error=true) the delivery report
+    will not be called and the msgstate will thus never be freed.
+
+    Since a proper broker is required for messages to be succesfully sent
+    this test must be run from the integration tests rather than
+    the unit tests.
+    """
+    p = confluent_kafka.Producer({"bootstrap.servers": bootstrap_servers,
+                                  'broker.address.family':'v4',
+                                  "delivery.report.only.error":True})
+
+    class DrOnlyTest (object):
+        def __init__ (self):
+            self.remaining = 1
+
+        def handle_err (self, err, msg):
+            """ This delivery handler should only get called for errored msgs """
+            assert "BAD:" in msg.value().decode('utf-8')
+            assert err is not None
+            self.remaining -= 1
+
+        def handle_success (self, err, msg):
+            """ This delivery handler should never get called """
+            # FIXME: Can we verify that it is actually garbage collected?
+            assert "GOOD:" in msg.value().decode('utf-8')
+            assert err is None
+            assert False, "should never come here"
+
+    state = DrOnlyTest()
+
+    print('only.error: Verifying delivery.report.only.error')
+    p.produce(topic, "BAD: This message will make not make it".encode('utf-8'),
+              partition=99, on_delivery=state.handle_err)
+
+    p.produce(topic, "GOOD: This message will make make it".encode('utf-8'),
+              on_delivery=state.handle_success)
+
+    print('only.error: Waiting for flush')
+    p.flush(10000)
+
+    print('only.error: Remaining messages now %d' % state.remaining)
+    assert state.remaining == 0
+
+
 
 def verify_avro():
     from confluent_kafka import avro

Original file line number	Diff line number	Diff line change
`@@ -152,6 +152,10 @@ static void dr_msg_cb (rd_kafka_t rk, const rd_kafka_message_t rkm,`
`152`	`152`	`goto done;`
`153`	`153`	`}`
`154`	`154`
	`155`	`+ /* Skip callback if delivery.report.only.error=true */`
	`156`	`+ if (self->u.Producer.dr_only_error && !rkm->err)`
	`157`	`+ goto done;`
	`158`	`+`
`155`	`159`	`msgobj = Message_new0(self, rkm);`
`156`	`160`
`157`	`161`	`args = Py_BuildValue("(OO)",`