Merge pull request confluentinc#254 from johnistan/jdennison/multiple-close-calls-cause-segfault

edenhill · web-flow · commit 3257fea51c4a · 2017-10-09T18:09:07.000+02:00
Makes calling consumer.close() multiple times safe
diff --git a/confluent_kafka/src/Consumer.c b/confluent_kafka/src/Consumer.c
@@ -550,19 +550,25 @@ static PyObject *Consumer_poll (Handle *self, PyObject *args,
 
 
 static PyObject *Consumer_close (Handle *self, PyObject *ignore) {
-	CallState cs;
+        CallState cs;
+
+        if (!self->rk) {
+                PyErr_SetString(PyExc_RuntimeError,
+                                "Consumer already closed");
+                return NULL;
+        }
 
-	CallState_begin(self, &cs);
+        CallState_begin(self, &cs);
 
-	rd_kafka_consumer_close(self->rk);
+        rd_kafka_consumer_close(self->rk);
 
-	rd_kafka_destroy(self->rk);
-	self->rk = NULL;
+        rd_kafka_destroy(self->rk);
+        self->rk = NULL;
 
         if (!CallState_end(self, &cs))
                 return NULL;
 
-	Py_RETURN_NONE;
+        Py_RETURN_NONE;
 }
 
 
@@ -736,6 +742,7 @@ static PyMethodDef Consumer_methods[] = {
 	  "see :py:func::`poll()` for more info.\n"
 	  "\n"
 	  "  :rtype: None\n"
+      "  :raises: RuntimeError if called on a closed consumer\n"
 	  "\n"
 	},
 	{ NULL }
diff --git a/tests/test_Consumer.py b/tests/test_Consumer.py
@@ -158,3 +158,22 @@ def poll(self, somearg):
     sc = SubConsumer({"group.id": "test", "session.timeout.ms": "90"})
     sc.poll("astring")
     sc.close()
+
+
+def test_multiple_close_throw_exception():
+    """ Calling Consumer.close() multiple times should throw Runtime Exception
+    """
+    c = Consumer({'group.id': 'test',
+                  'enable.auto.commit': True,
+                  'enable.auto.offset.store': False,
+                  'socket.timeout.ms': 50,
+                  'session.timeout.ms': 100})
+
+    c.subscribe(["test"])
+
+    c.unsubscribe()
+    c.close()
+
+    with pytest.raises(RuntimeError) as ex:
+        c.close()
+        assert 'Consumer already closed' == str(ex.value)
