From a769b2b9a405ceba8bbb46dbcfdf6dfce998fb26 Mon Sep 17 00:00:00 2001 From: Ojasva Jain Date: Tue, 21 Jul 2026 19:07:22 +0530 Subject: [PATCH 01/10] Fix Producer.close() races with concurrent calls and with itself Producer.close() previously raced with concurrent produce()/poll()/ flush()/produce_batch()/transaction calls and with itself when called from multiple threads, both leading to use-after-free/double-free on the underlying rd_kafka_t handle. Adds an active_calls/closing guard (Handle_enter_rk_use/Handle_exit_rk_use) so every method that touches self->rk registers itself before use, and close() drains in-flight calls before tearing down; a CAS on `closing` ensures only one concurrent close() call performs the actual teardown, with losing callers waiting for it to finish rather than racing it. Adds tests/parallel/test_producer_close_race.py covering each affected method racing close(), close() racing itself, and close()'s blocking behavior. Uses pytest-forked (POSIX only) so a regression segfault fails only that test. Integration tests against a real broker are still pending. --- requirements/requirements-tests.txt | 1 + src/confluent_kafka/src/Producer.c | 129 ++++++--- src/confluent_kafka/src/confluent_kafka.c | 44 ++- src/confluent_kafka/src/confluent_kafka.h | 68 ++++- tests/parallel/__init__.py | 0 tests/parallel/test_producer_close_race.py | 309 +++++++++++++++++++++ 6 files changed, 514 insertions(+), 37 deletions(-) create mode 100644 tests/parallel/__init__.py create mode 100644 tests/parallel/test_producer_close_race.py diff --git a/requirements/requirements-tests.txt b/requirements/requirements-tests.txt index a597b70aa..17d623bd4 100644 --- a/requirements/requirements-tests.txt +++ b/requirements/requirements-tests.txt @@ -13,6 +13,7 @@ pytest_cov pluggy<1.6.0 pytest-asyncio async-timeout +pytest-forked; sys_platform != "win32" # Formatting tools black>=24.0.0 diff --git a/src/confluent_kafka/src/Producer.c b/src/confluent_kafka/src/Producer.c index 0bfc810b8..8a104ee51 100644 --- a/src/confluent_kafka/src/Producer.c +++ b/src/confluent_kafka/src/Producer.c @@ -27,6 +27,12 @@ #include "confluent_kafka.h" +#ifdef _WIN32 +#include +#else +#include +#endif + /** * @brief KNOWN ISSUES @@ -296,12 +302,11 @@ Producer_produce(Handle *self, PyObject *args, PyObject *kwargs) { if (!dr_cb || dr_cb == Py_None) dr_cb = self->u.Producer.default_dr_cb; - if (!self->rk) { + if (!Handle_enter_rk_use(self)) { #ifdef RD_KAFKA_V_HEADERS if (rd_headers) rd_kafka_headers_destroy(rd_headers); #endif - PyErr_SetString(PyExc_RuntimeError, ERR_MSG_PRODUCER_CLOSED); return NULL; } @@ -323,6 +328,8 @@ Producer_produce(Handle *self, PyObject *args, PyObject *kwargs) { key_len, msgstate); #endif + Handle_exit_rk_use(self); + if (err) { if (msgstate) Producer_msgstate_destroy(msgstate); @@ -421,12 +428,13 @@ static PyObject *Producer_poll(Handle *self, PyObject *args, PyObject *kwargs) { if (!PyArg_ParseTupleAndKeywords(args, kwargs, "|d", kws, &tmout)) return NULL; - if (!self->rk) { - PyErr_SetString(PyExc_RuntimeError, ERR_MSG_PRODUCER_CLOSED); + if (!Handle_enter_rk_use(self)) return NULL; - } r = Producer_poll0(self, cfl_timeout_ms(tmout)); + + Handle_exit_rk_use(self); + if (r == -1) return NULL; @@ -469,10 +477,8 @@ Producer_flush(Handle *self, PyObject *args, PyObject *kwargs) { if (!PyArg_ParseTupleAndKeywords(args, kwargs, "|d", kws, &tmout)) return NULL; - if (!self->rk) { - PyErr_SetString(PyExc_RuntimeError, ERR_MSG_PRODUCER_CLOSED); + if (!Handle_enter_rk_use(self)) return NULL; - } total_timeout_ms = cfl_timeout_ms(tmout); CallState_begin(self, &cs); @@ -507,6 +513,7 @@ Producer_flush(Handle *self, PyObject *args, PyObject *kwargs) { * interruptibility) */ chunk_count++; if (check_signals_between_chunks(self, &cs)) { + Handle_exit_rk_use(self); return NULL; /* Signal detected */ } @@ -524,12 +531,16 @@ Producer_flush(Handle *self, PyObject *args, PyObject *kwargs) { } } - if (!CallState_end(self, &cs)) + if (!CallState_end(self, &cs)) { + Handle_exit_rk_use(self); return NULL; + } if (err) /* Get the queue length on error (timeout) */ qlen = rd_kafka_outq_len(self->rk); + Handle_exit_rk_use(self); + return cfl_PyInt_FromInt(qlen); } @@ -542,6 +553,37 @@ Producer_close(Handle *self, PyObject *args, PyObject *kwargs) { if (!self->rk) Py_RETURN_TRUE; + /* Only one concurrent close() can destroy rk, otherwise, + * two threads could both reach rd_kafka_destroy() on the + * same handle (a double-free). The losing thread(s) wait for the + * winner to finish and then return True, same as a normal close(), + * rather than racing it. */ + if (!atomic_int_cas(&self->closing, 0, 1)) { + while (self->rk) { + CallState_begin(self, &cs); +#ifdef _WIN32 + Sleep(100); +#else + usleep(100000); +#endif + CallState_end(self, &cs); + } + Py_RETURN_TRUE; + } + + /* Signal in-flight calls to stop, and wait for them to finish + * using self->rk before destroying it -- see Handle_enter_rk_use(). + * New calls will see `closing` and fail with ERR_MSG_PRODUCER_CLOSED. */ + while (atomic_int_get(&self->active_calls) > 0) { + CallState_begin(self, &cs); +#ifdef _WIN32 + Sleep(100); +#else + usleep(100000); +#endif + CallState_end(self, &cs); + } + CallState_begin(self, &cs); /* Flush any pending messages (wait indefinitely to ensure delivery) */ @@ -817,10 +859,8 @@ Producer_produce_batch(Handle *self, PyObject *args, PyObject *kwargs) { return cfl_PyInt_FromInt(0); } - if (!self->rk) { - PyErr_SetString(PyExc_RuntimeError, ERR_MSG_PRODUCER_CLOSED); + if (!Handle_enter_rk_use(self)) return NULL; - } /* Allocate arrays for librdkafka messages and msgstates */ rkmessages = calloc(message_cnt, sizeof(*rkmessages)); @@ -849,6 +889,8 @@ Producer_produce_batch(Handle *self, PyObject *args, PyObject *kwargs) { messages_list, rkt, partition, rkmessages, msgstates, message_cnt); cleanup: + Handle_exit_rk_use(self); + /* Cleanup resources */ if (rkt) rd_kafka_topic_destroy(rkt); @@ -871,21 +913,22 @@ static PyObject *Producer_init_transactions(Handle *self, PyObject *args) { if (!PyArg_ParseTuple(args, "|d", &tmout)) return NULL; - if (!self->rk) { - PyErr_SetString(PyExc_RuntimeError, ERR_MSG_PRODUCER_CLOSED); + if (!Handle_enter_rk_use(self)) return NULL; - } CallState_begin(self, &cs); error = rd_kafka_init_transactions(self->rk, cfl_timeout_ms(tmout)); if (!CallState_end(self, &cs)) { + Handle_exit_rk_use(self); if (error) /* Ignore error in favour of callstate exception */ rd_kafka_error_destroy(error); return NULL; } + Handle_exit_rk_use(self); + if (error) { cfl_PyErr_from_error_destroy(error); return NULL; @@ -897,13 +940,13 @@ static PyObject *Producer_init_transactions(Handle *self, PyObject *args) { static PyObject *Producer_begin_transaction(Handle *self) { rd_kafka_error_t *error; - if (!self->rk) { - PyErr_SetString(PyExc_RuntimeError, ERR_MSG_PRODUCER_CLOSED); + if (!Handle_enter_rk_use(self)) return NULL; - } error = rd_kafka_begin_transaction(self->rk); + Handle_exit_rk_use(self); + if (error) { cfl_PyErr_from_error_destroy(error); return NULL; @@ -924,16 +967,17 @@ static PyObject *Producer_send_offsets_to_transaction(Handle *self, if (!PyArg_ParseTuple(args, "OO|d", &offsets, &metadata, &tmout)) return NULL; - if (!self->rk) { - PyErr_SetString(PyExc_RuntimeError, ERR_MSG_PRODUCER_CLOSED); + if (!Handle_enter_rk_use(self)) return NULL; - } - if (!(c_offsets = py_to_c_parts(offsets))) + if (!(c_offsets = py_to_c_parts(offsets))) { + Handle_exit_rk_use(self); return NULL; + } if (!(cgmd = py_to_c_cgmd(metadata))) { rd_kafka_topic_partition_list_destroy(c_offsets); + Handle_exit_rk_use(self); return NULL; } @@ -946,11 +990,14 @@ static PyObject *Producer_send_offsets_to_transaction(Handle *self, rd_kafka_topic_partition_list_destroy(c_offsets); if (!CallState_end(self, &cs)) { + Handle_exit_rk_use(self); if (error) /* Ignore error in favour of callstate exception */ rd_kafka_error_destroy(error); return NULL; } + Handle_exit_rk_use(self); + if (error) { cfl_PyErr_from_error_destroy(error); return NULL; @@ -967,21 +1014,22 @@ static PyObject *Producer_commit_transaction(Handle *self, PyObject *args) { if (!PyArg_ParseTuple(args, "|d", &tmout)) return NULL; - if (!self->rk) { - PyErr_SetString(PyExc_RuntimeError, ERR_MSG_PRODUCER_CLOSED); + if (!Handle_enter_rk_use(self)) return NULL; - } CallState_begin(self, &cs); error = rd_kafka_commit_transaction(self->rk, cfl_timeout_ms(tmout)); if (!CallState_end(self, &cs)) { + Handle_exit_rk_use(self); if (error) /* Ignore error in favour of callstate exception */ rd_kafka_error_destroy(error); return NULL; } + Handle_exit_rk_use(self); + if (error) { cfl_PyErr_from_error_destroy(error); return NULL; @@ -998,21 +1046,22 @@ static PyObject *Producer_abort_transaction(Handle *self, PyObject *args) { if (!PyArg_ParseTuple(args, "|d", &tmout)) return NULL; - if (!self->rk) { - PyErr_SetString(PyExc_RuntimeError, ERR_MSG_PRODUCER_CLOSED); + if (!Handle_enter_rk_use(self)) return NULL; - } CallState_begin(self, &cs); error = rd_kafka_abort_transaction(self->rk, cfl_timeout_ms(tmout)); if (!CallState_end(self, &cs)) { + Handle_exit_rk_use(self); if (error) /* Ignore error in favour of callstate exception */ rd_kafka_error_destroy(error); return NULL; } + Handle_exit_rk_use(self); + if (error) { cfl_PyErr_from_error_destroy(error); return NULL; @@ -1034,10 +1083,8 @@ static void *Producer_purge(Handle *self, PyObject *args, PyObject *kwargs) { &in_flight, &blocking)) return NULL; - if (!self->rk) { - PyErr_SetString(PyExc_RuntimeError, ERR_MSG_PRODUCER_CLOSED); + if (!Handle_enter_rk_use(self)) return NULL; - } if (in_queue) purge_strategy = RD_KAFKA_PURGE_F_QUEUE; @@ -1048,6 +1095,8 @@ static void *Producer_purge(Handle *self, PyObject *args, PyObject *kwargs) { err = rd_kafka_purge(self->rk, purge_strategy); + Handle_exit_rk_use(self); + if (err) { cfl_PyErr_Format(err, "Purge failed: %s", rd_kafka_err2str(err)); @@ -1403,9 +1452,23 @@ static PyMethodDef Producer_methods[] = { static Py_ssize_t Producer__len__(Handle *self) { - if (!self->rk) + Py_ssize_t len; + + /* __len__ must never raise, so we can't use Handle_enter_rk_use() + * (which sets an exception on failure) -- fall back to returning 0, + * , if the Handle is closed/closing. */ + if (atomic_int_get(&self->closing) || !self->rk) + return 0; + atomic_int_inc(&self->active_calls); + if (atomic_int_get(&self->closing) || !self->rk) { + atomic_int_dec(&self->active_calls); return 0; - return rd_kafka_outq_len(self->rk); + } + + len = rd_kafka_outq_len(self->rk); + + atomic_int_dec(&self->active_calls); + return len; } diff --git a/src/confluent_kafka/src/confluent_kafka.c b/src/confluent_kafka/src/confluent_kafka.c index 7ec5dd781..29d5dd9ab 100644 --- a/src/confluent_kafka/src/confluent_kafka.c +++ b/src/confluent_kafka/src/confluent_kafka.c @@ -2431,7 +2431,8 @@ int wait_for_oauth_token_set(Handle *h) { int max_wait_sec = 10; int retry_interval_sec = 1; /* Check every 1 sec */ int elapsed_sec = 0; - while (!h->oauth_token_set && elapsed_sec < max_wait_sec) { + while (!atomic_int_get(&h->oauth_token_set) && + elapsed_sec < max_wait_sec) { CallState cs; CallState_begin(h, &cs); #ifdef _WIN32 @@ -2443,7 +2444,7 @@ int wait_for_oauth_token_set(Handle *h) { elapsed_sec += retry_interval_sec; } - if (!h->oauth_token_set) { + if (!atomic_int_get(&h->oauth_token_set)) { /* Token timeout. Don't tear down here — each _init knows * whether to call rd_kafka_destroy() or * rd_kafka_share_destroy() for what it allocated. */ @@ -2529,7 +2530,7 @@ oauth_cb(rd_kafka_t *rk, const char *oauthbearer_config, void *opaque) { PyErr_Format(PyExc_ValueError, "%s", err_msg); goto fail; } - h->oauth_token_set = 1; + atomic_int_set(&h->oauth_token_set, 1); goto done; fail: @@ -3314,6 +3315,43 @@ int CallState_end(Handle *h, CallState *cs) { } +/** + * @brief Mark self->rk as in-use by the calling thread, so that close() + * (running concurrently on another thread) will wait for us before + * destroying it. Must be called before CallState_begin(). + * + * @returns 1 if self->rk is safe to use (active_calls has been + * incremented; caller must call Handle_exit_rk_use() on every + * return path), or 0 with ERR_MSG_PRODUCER_CLOSED set if the + * Handle is closed/closing (nothing to undo). + * + * @warning Not re-entrant: don't call from a method that's already + * between its own Handle_enter_rk_use()/Handle_exit_rk_use(). + */ +int Handle_enter_rk_use(Handle *h) { + if (atomic_int_get(&h->closing) || !h->rk) { + PyErr_SetString(PyExc_RuntimeError, ERR_MSG_PRODUCER_CLOSED); + return 0; + } + atomic_int_inc(&h->active_calls); + /* close() may have started between our check above and the + * increment; re-check now that we're counted. */ + if (atomic_int_get(&h->closing) || !h->rk) { + atomic_int_dec(&h->active_calls); + PyErr_SetString(PyExc_RuntimeError, ERR_MSG_PRODUCER_CLOSED); + return 0; + } + return 1; +} + +/** + * @brief Counterpart to Handle_enter_rk_use(): call on every return path + * after a successful Handle_enter_rk_use(). + */ +void Handle_exit_rk_use(Handle *h) { + atomic_int_dec(&h->active_calls); +} + /** * @brief Get the current thread's CallState and re-locks the GIL. */ diff --git a/src/confluent_kafka/src/confluent_kafka.h b/src/confluent_kafka/src/confluent_kafka.h index fc204ce02..34746d845 100644 --- a/src/confluent_kafka/src/confluent_kafka.h +++ b/src/confluent_kafka/src/confluent_kafka.h @@ -35,6 +35,55 @@ #endif +/** + * @brief Minimal portable atomic-int primitives. + * + * Not : this project's build does not pin a C standard + * version, and MSVC's C11 support is version- and flag-gated, + * so it cannot be relied on across this project's actual build matrix. + * + * Modeled on librdkafka's own rdatomic.h: GCC/Clang __atomic_* builtins + * on non-Windows, Interlocked* on Windows. + */ +#if defined(_MSC_VER) +typedef volatile LONG atomic_int_t; + +#define atomic_int_init(p, v) (*(p) = (v)) +#define atomic_int_inc(p) InterlockedIncrement((p)) +#define atomic_int_dec(p) InterlockedDecrement((p)) +#define atomic_int_get(p) InterlockedCompareExchange((p), 0, 0) +#define atomic_int_set(p, v) InterlockedExchange((p), (v)) + +/** + * @brief Atomic compare-and-swap: if *p == expected, set *p = desired and + * return 1; otherwise leave *p unchanged and return 0. + */ +static __inline int atomic_int_cas(atomic_int_t *p, LONG expected, + LONG desired) { + return InterlockedCompareExchange(p, desired, expected) == expected; +} + +#else /* gcc / clang */ +typedef int atomic_int_t; + +#define atomic_int_init(p, v) __atomic_store_n((p), (v), __ATOMIC_SEQ_CST) +#define atomic_int_inc(p) __atomic_add_fetch((p), 1, __ATOMIC_SEQ_CST) +#define atomic_int_dec(p) __atomic_sub_fetch((p), 1, __ATOMIC_SEQ_CST) +#define atomic_int_get(p) __atomic_load_n((p), __ATOMIC_SEQ_CST) +#define atomic_int_set(p, v) __atomic_store_n((p), (v), __ATOMIC_SEQ_CST) + +/** + * @brief Atomic compare-and-swap: if *p == expected, set *p = desired and + * return 1; otherwise leave *p unchanged and return 0. + */ +static inline int atomic_int_cas(atomic_int_t *p, int expected, int desired) { + return __atomic_compare_exchange_n(p, &expected, desired, + 0 /* strong */, __ATOMIC_SEQ_CST, + __ATOMIC_SEQ_CST); +} +#endif + + /** * @brief confluent-kafka-python version, must match that of pyproject.toml. */ @@ -241,7 +290,13 @@ typedef struct { PyObject *logger; PyObject *oauth_cb; - int oauth_token_set; + atomic_int_t oauth_token_set; + + /* Protects self->rk from being freed by close() while another method + * is still using it. See Handle_enter_rk_use()/Handle_exit_rk_use() + * in confluent_kafka.c. */ + atomic_int_t active_calls; + atomic_int_t closing; union { /** @@ -350,6 +405,17 @@ void CallState_begin(Handle *h, CallState *cs); */ int CallState_end(Handle *h, CallState *cs); +/** + * @brief Mark self->rk as in-use by the calling thread. + * @returns 1 if safe to use (caller must call Handle_exit_rk_use() on every + * return path), 0 with ERR_MSG_PRODUCER_CLOSED set otherwise. + */ +int Handle_enter_rk_use(Handle *h); +/** + * @brief Counterpart to Handle_enter_rk_use(). + */ +void Handle_exit_rk_use(Handle *h); + /** * @brief Get the current thread's CallState and re-locks the GIL. */ diff --git a/tests/parallel/__init__.py b/tests/parallel/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/parallel/test_producer_close_race.py b/tests/parallel/test_producer_close_race.py new file mode 100644 index 000000000..ccb084113 --- /dev/null +++ b/tests/parallel/test_producer_close_race.py @@ -0,0 +1,309 @@ +#!/usr/bin/env python +# +# Copyright 2026 Confluent Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import sys +import threading +import time + +import pytest + +from confluent_kafka import Consumer, Producer, TopicPartition + +# pytest-forked runs each marked test in its own forked child process (via +# os.fork(), POSIX only), so a segfault in one test only fails that test +# instead of taking down the whole pytest run. Not available on Windows. +forked = pytest.mark.forked +skip_on_windows = pytest.mark.skipif( + sys.platform == "win32", reason="pytest-forked requires os.fork(), not available on Windows" +) + +############################################################################### +# Tests for races between Producer.close() and concurrent calls to +# other methods on the same Producer instance. +############################################################################### + +_PRODUCER_CONF = {'bootstrap.servers': 'localhost:9092', 'socket.timeout.ms': 10, 'message.timeout.ms': 10} +_TXN_PRODUCER_CONF = dict(_PRODUCER_CONF, **{'transactional.id': 'test-producer-close-race-txn'}) +ITERATIONS = 10 + + +def _race_close_against(worker, conf=None, num_workers=1): + """ + Run `worker(producer, stop_event)` on `num_workers` other threads while + calling close() from the main thread, repeated `iterations` times + against a fresh Producer each time. Every close() call (the main + thread's, once per iteration) must return True. + """ + errors = [] + close_results = [] + + for i in range(ITERATIONS): + producer = Producer(conf or _PRODUCER_CONF) + stop_event = threading.Event() + start_barrier = threading.Barrier(num_workers + 1) + + def run_worker(): + try: + start_barrier.wait() + worker(producer, stop_event) + except Exception as e: # noqa: BLE001 - want to see any exception, not just crashes + errors.append((i, e)) + + threads = [threading.Thread(target=run_worker) for _ in range(num_workers)] + for t in threads: + t.start() + + start_barrier.wait() + close_results.append(producer.close()) + + stop_event.set() + for t in threads: + t.join(timeout=10) + assert all(not t.is_alive() for t in threads), f"iteration {i}: a worker thread did not finish after close()" + + assert not errors, f"unexpected exceptions from worker threads: {errors}" + assert all(close_results), f"not every close() call returned True: {close_results}" + + +def _worker_produce(producer, stop_event): + while not stop_event.is_set(): + try: + producer.produce('mytopic', value=b'x') + except RuntimeError: + # Expected once close() has fully completed on this thread's + # view of self->rk; anything other than a clean RuntimeError + # (e.g. a segfault) is the bug this test is trying to catch. + break + + +def _worker_poll(producer, stop_event): + while not stop_event.is_set(): + try: + producer.poll(0) + except RuntimeError: + break + + +@forked +@skip_on_windows +def test_close_races_produce(): + """close() concurrent with produce() on another thread.""" + _race_close_against(_worker_produce) + + +@forked +@skip_on_windows +def test_close_races_multiple_producers_and_pollers(): + """ + close() concurrent with several threads calling produce()/poll() at + once, not just one. + """ + num_workers = 8 + _race_close_against(_worker_produce, num_workers=num_workers) + _race_close_against(_worker_poll, num_workers=num_workers) + + +@forked +@skip_on_windows +def test_close_races_poll(): + """close() concurrent with poll() on another thread.""" + _race_close_against(_worker_poll) + + +@forked +@skip_on_windows +def test_close_races_flush(): + """close() concurrent with flush() on another thread.""" + + def worker(producer, stop_event): + while not stop_event.is_set(): + try: + producer.flush(0.01) + except RuntimeError: + break + + _race_close_against(worker) + + +@forked +@skip_on_windows +def test_close_races_produce_batch(): + """close() concurrent with produce_batch() on another thread.""" + + def worker(producer, stop_event): + messages = [{'value': b'x'}, {'value': b'y'}] + while not stop_event.is_set(): + try: + producer.produce_batch('mytopic', messages) + except RuntimeError: + break + + _race_close_against(worker) + + +@forked +@skip_on_windows +def test_close_races_init_transactions(): + """close() concurrent with init_transactions() on another thread.""" + + def worker(producer, stop_event): + while not stop_event.is_set(): + try: + producer.init_transactions(0.05) + except RuntimeError: + break + except Exception: # noqa: BLE001 - librdkafka state/timeout errors are expected without a broker + pass + + _race_close_against(worker, conf=_TXN_PRODUCER_CONF) + + +@forked +@skip_on_windows +def test_close_races_begin_transaction(): + """close() concurrent with begin_transaction() on another thread.""" + + def worker(producer, stop_event): + while not stop_event.is_set(): + try: + producer.begin_transaction() + except RuntimeError: + break + except Exception: # noqa: BLE001 - librdkafka state errors are expected without a broker + pass + + _race_close_against(worker, conf=_TXN_PRODUCER_CONF) + + +@forked +@skip_on_windows +def test_close_races_commit_transaction(): + """close() concurrent with commit_transaction() on another thread.""" + + def worker(producer, stop_event): + while not stop_event.is_set(): + try: + producer.commit_transaction(0.05) + except RuntimeError: + break + except Exception: # noqa: BLE001 - librdkafka state/timeout errors are expected without a broker + pass + + _race_close_against(worker, conf=_TXN_PRODUCER_CONF) + + +@forked +@skip_on_windows +def test_close_races_abort_transaction(): + """close() concurrent with abort_transaction() on another thread.""" + + def worker(producer, stop_event): + while not stop_event.is_set(): + try: + producer.abort_transaction(0.05) + except RuntimeError: + break + except Exception: # noqa: BLE001 - librdkafka state/timeout errors are expected without a broker + pass + + _race_close_against(worker, conf=_TXN_PRODUCER_CONF) + + +@forked +@skip_on_windows +def test_close_races_send_offsets_to_transaction(): + """close() concurrent with send_offsets_to_transaction() on another thread.""" + + def worker(producer, stop_event): + # consumer_group_metadata() doesn't need a live broker connection. + consumer = Consumer({'group.id': 'test-producer-close-race', 'socket.timeout.ms': 10}) + metadata = consumer.consumer_group_metadata() + consumer.close() + + offsets = [TopicPartition('mytopic', 0, 1)] + while not stop_event.is_set(): + try: + producer.send_offsets_to_transaction(offsets, metadata, 0.05) + except RuntimeError: + break + except Exception: # noqa: BLE001 - librdkafka state/timeout errors are expected without a broker + pass + + _race_close_against(worker, conf=_TXN_PRODUCER_CONF) + + +@forked +@skip_on_windows +def test_close_races_purge(): + """close() concurrent with purge() on another thread.""" + + def worker(producer, stop_event): + while not stop_event.is_set(): + try: + producer.purge() + except RuntimeError: + break + + _race_close_against(worker) + + +@forked +@skip_on_windows +def test_close_races_close(): + """Multiple threads calling close() on the same Producer at once.""" + worker_close_results = [] + num_workers = 7 + + def worker(producer, stop_event): + worker_close_results.append(producer.close()) + + _race_close_against(worker, num_workers=num_workers) + + assert all(worker_close_results), f"not every worker close() call returned True: {worker_close_results}" + assert len(worker_close_results) == num_workers * ITERATIONS + + +############################################################################### +# close() blocks until an in-flight call finishes, +# rather than just not crashing while one is running. +############################################################################### + + +def test_close_waits_for_in_flight_call(): + """close() blocks until an in-flight poll() call finishes.""" + producer = Producer(_PRODUCER_CONF) + poll_started = threading.Event() + poll_duration = 10 + + def run_poll(): + poll_started.set() + producer.poll(poll_duration) + + t = threading.Thread(target=run_poll) + t.start() + poll_started.wait() + time.sleep(2) + + close_start = time.monotonic() + producer.close() + close_duration = time.monotonic() - close_start + + t.join(timeout=10) + assert not t.is_alive(), "poll() thread did not finish after close()" + assert close_duration >= 7, ( + f"close() took only {close_duration:.2f}s -- expected it to block " + f"for close to the in-flight poll({poll_duration}s) call" + ) From 0d6a9eb50a3de4a82ddfa799f4b55391ff1e2535 Mon Sep 17 00:00:00 2001 From: Ojasva Jain Date: Tue, 21 Jul 2026 21:55:39 +0530 Subject: [PATCH 02/10] Replace pytest-forked with subprocess-based test isolation --- requirements/requirements-tests.txt | 1 - tests/parallel/test_producer_close_race.py | 70 ++++++++++------------ 2 files changed, 31 insertions(+), 40 deletions(-) diff --git a/requirements/requirements-tests.txt b/requirements/requirements-tests.txt index 17d623bd4..a597b70aa 100644 --- a/requirements/requirements-tests.txt +++ b/requirements/requirements-tests.txt @@ -13,7 +13,6 @@ pytest_cov pluggy<1.6.0 pytest-asyncio async-timeout -pytest-forked; sys_platform != "win32" # Formatting tools black>=24.0.0 diff --git a/tests/parallel/test_producer_close_race.py b/tests/parallel/test_producer_close_race.py index ccb084113..2265836ec 100644 --- a/tests/parallel/test_producer_close_race.py +++ b/tests/parallel/test_producer_close_race.py @@ -14,21 +14,11 @@ # See the License for the specific language governing permissions and # limitations under the License. -import sys import threading import time -import pytest - from confluent_kafka import Consumer, Producer, TopicPartition - -# pytest-forked runs each marked test in its own forked child process (via -# os.fork(), POSIX only), so a segfault in one test only fails that test -# instead of taking down the whole pytest run. Not available on Windows. -forked = pytest.mark.forked -skip_on_windows = pytest.mark.skipif( - sys.platform == "win32", reason="pytest-forked requires os.fork(), not available on Windows" -) +from tests.parallel.conftest import subprocess_isolated ############################################################################### # Tests for races between Producer.close() and concurrent calls to @@ -97,15 +87,13 @@ def _worker_poll(producer, stop_event): break -@forked -@skip_on_windows +@subprocess_isolated def test_close_races_produce(): """close() concurrent with produce() on another thread.""" _race_close_against(_worker_produce) -@forked -@skip_on_windows +@subprocess_isolated def test_close_races_multiple_producers_and_pollers(): """ close() concurrent with several threads calling produce()/poll() at @@ -116,15 +104,13 @@ def test_close_races_multiple_producers_and_pollers(): _race_close_against(_worker_poll, num_workers=num_workers) -@forked -@skip_on_windows +@subprocess_isolated def test_close_races_poll(): """close() concurrent with poll() on another thread.""" _race_close_against(_worker_poll) -@forked -@skip_on_windows +@subprocess_isolated def test_close_races_flush(): """close() concurrent with flush() on another thread.""" @@ -138,8 +124,7 @@ def worker(producer, stop_event): _race_close_against(worker) -@forked -@skip_on_windows +@subprocess_isolated def test_close_races_produce_batch(): """close() concurrent with produce_batch() on another thread.""" @@ -154,8 +139,7 @@ def worker(producer, stop_event): _race_close_against(worker) -@forked -@skip_on_windows +@subprocess_isolated def test_close_races_init_transactions(): """close() concurrent with init_transactions() on another thread.""" @@ -171,8 +155,7 @@ def worker(producer, stop_event): _race_close_against(worker, conf=_TXN_PRODUCER_CONF) -@forked -@skip_on_windows +@subprocess_isolated def test_close_races_begin_transaction(): """close() concurrent with begin_transaction() on another thread.""" @@ -188,8 +171,7 @@ def worker(producer, stop_event): _race_close_against(worker, conf=_TXN_PRODUCER_CONF) -@forked -@skip_on_windows +@subprocess_isolated def test_close_races_commit_transaction(): """close() concurrent with commit_transaction() on another thread.""" @@ -205,8 +187,7 @@ def worker(producer, stop_event): _race_close_against(worker, conf=_TXN_PRODUCER_CONF) -@forked -@skip_on_windows +@subprocess_isolated def test_close_races_abort_transaction(): """close() concurrent with abort_transaction() on another thread.""" @@ -222,8 +203,7 @@ def worker(producer, stop_event): _race_close_against(worker, conf=_TXN_PRODUCER_CONF) -@forked -@skip_on_windows +@subprocess_isolated def test_close_races_send_offsets_to_transaction(): """close() concurrent with send_offsets_to_transaction() on another thread.""" @@ -245,8 +225,7 @@ def worker(producer, stop_event): _race_close_against(worker, conf=_TXN_PRODUCER_CONF) -@forked -@skip_on_windows +@subprocess_isolated def test_close_races_purge(): """close() concurrent with purge() on another thread.""" @@ -260,8 +239,7 @@ def worker(producer, stop_event): _race_close_against(worker) -@forked -@skip_on_windows +@subprocess_isolated def test_close_races_close(): """Multiple threads calling close() on the same Producer at once.""" worker_close_results = [] @@ -287,10 +265,13 @@ def test_close_waits_for_in_flight_call(): producer = Producer(_PRODUCER_CONF) poll_started = threading.Event() poll_duration = 10 + poll_finished_at = None def run_poll(): + nonlocal poll_finished_at poll_started.set() producer.poll(poll_duration) + poll_finished_at = time.monotonic() t = threading.Thread(target=run_poll) t.start() @@ -299,11 +280,22 @@ def run_poll(): close_start = time.monotonic() producer.close() - close_duration = time.monotonic() - close_start + close_end = time.monotonic() t.join(timeout=10) assert not t.is_alive(), "poll() thread did not finish after close()" - assert close_duration >= 7, ( - f"close() took only {close_duration:.2f}s -- expected it to block " - f"for close to the in-flight poll({poll_duration}s) call" + assert poll_finished_at is not None, "poll() never finished" + + # close() should have waited for close to the actual remaining poll() + # duration, not a fixed guess -- a slower/loaded machine can take longer + # to reach close_start, shrinking how much of poll_duration is actually + # left to wait for, so we compare against what really happened instead + # of a hardcoded constant. + close_duration = close_end - close_start + remaining_poll_duration = poll_finished_at - close_start + assert close_duration >= remaining_poll_duration * 0.9, ( + f"close() took {close_duration:.2f}s but the in-flight poll() call " + f"was still going to run for {remaining_poll_duration:.2f}s more " + f"-- close() returned too early relative to what it should have " + f"waited for" ) From 75d909705139cdaf32b854ddd86f72bb89163462 Mon Sep 17 00:00:00 2001 From: Ojasva Jain Date: Tue, 21 Jul 2026 23:06:19 +0530 Subject: [PATCH 03/10] Fix subprocess_isolated import and flaky close() timing assertion --- tests/parallel/_subprocess_isolation.py | 75 ++++++++++++++++++++++ tests/parallel/test_producer_close_race.py | 2 +- 2 files changed, 76 insertions(+), 1 deletion(-) create mode 100644 tests/parallel/_subprocess_isolation.py diff --git a/tests/parallel/_subprocess_isolation.py b/tests/parallel/_subprocess_isolation.py new file mode 100644 index 000000000..46191d735 --- /dev/null +++ b/tests/parallel/_subprocess_isolation.py @@ -0,0 +1,75 @@ +#!/usr/bin/env python +# +# Copyright 2026 Confluent Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Shared test infrastructure for tests/parallel/. + +Tests here deliberately race Producer/Consumer methods against each other, +so a real regression can segfault the process instead of just failing an +assertion. subprocess_isolated() below runs a decorated test in a fresh +`python -m pytest` subprocess, so a crash only fails that one test instead +of taking down the whole suite. + +This is a plain module, not conftest.py: conftest.py is auto-loaded by +pytest's own plugin machinery and isn't meant to be imported as a regular +module -- doing so (`from tests.parallel.conftest import ...`) is not +guaranteed to resolve the same way from a re-invoked subprocess as it does +under pytest's own collection, and failed with +"ModuleNotFoundError: No module named 'tests.parallel.conftest'" in +exactly that scenario. +""" + +import functools +import os +import subprocess +import sys + +_SUBPROCESS_MARKER_ENV = "_PARALLEL_TESTS_SUBPROCESS" +_SUBPROCESS_TIMEOUT_SECONDS = 120 + + +def subprocess_isolated(test_func): + """ + Run `test_func` in a fresh `python -m pytest` subprocess instead of + in-process. A clean run passes normally; a crash (e.g. segfault) shows + up as a non-zero/negative subprocess return code, which is turned into + a normal assertion failure here -- so it fails only this test rather + than taking down the whole run. + """ + + @functools.wraps(test_func) + def wrapper(*args, **kwargs): + if os.environ.get(_SUBPROCESS_MARKER_ENV): + # Already inside the re-invoked subprocess: run the real body. + return test_func(*args, **kwargs) + + test_file = sys.modules[test_func.__module__].__file__ + node_id = f"{os.path.relpath(test_file)}::{test_func.__name__}" + env = dict(os.environ, **{_SUBPROCESS_MARKER_ENV: "1"}) + + result = subprocess.run( + [sys.executable, "-m", "pytest", "-p", "no:cacheprovider", "-q", node_id], + capture_output=True, + text=True, + timeout=_SUBPROCESS_TIMEOUT_SECONDS, + env=env, + ) + assert result.returncode == 0, ( + f"{test_func.__name__} crashed/failed in subprocess " + f"(returncode={result.returncode}):\nstdout:\n{result.stdout}\nstderr:\n{result.stderr}" + ) + + return wrapper diff --git a/tests/parallel/test_producer_close_race.py b/tests/parallel/test_producer_close_race.py index 2265836ec..bff4bdc3e 100644 --- a/tests/parallel/test_producer_close_race.py +++ b/tests/parallel/test_producer_close_race.py @@ -18,7 +18,7 @@ import time from confluent_kafka import Consumer, Producer, TopicPartition -from tests.parallel.conftest import subprocess_isolated +from tests.parallel._subprocess_isolation import subprocess_isolated ############################################################################### # Tests for races between Producer.close() and concurrent calls to From a7453ccd3389355393956ebd1b86c83627d60a53 Mon Sep 17 00:00:00 2001 From: Ojasva Jain Date: Wed, 22 Jul 2026 20:13:35 +0530 Subject: [PATCH 04/10] Rename tests/parallel to tests/concurrency and add integration tests for Producer close()/transaction races --- tests/{parallel => concurrency}/__init__.py | 0 .../_subprocess_isolation.py | 12 +- .../test_producer_close_race.py | 7 +- .../integration/producer/test_concurrency.py | 355 ++++++++++++++++++ 4 files changed, 359 insertions(+), 15 deletions(-) rename tests/{parallel => concurrency}/__init__.py (100%) rename tests/{parallel => concurrency}/_subprocess_isolation.py (81%) rename tests/{parallel => concurrency}/test_producer_close_race.py (96%) create mode 100644 tests/integration/producer/test_concurrency.py diff --git a/tests/parallel/__init__.py b/tests/concurrency/__init__.py similarity index 100% rename from tests/parallel/__init__.py rename to tests/concurrency/__init__.py diff --git a/tests/parallel/_subprocess_isolation.py b/tests/concurrency/_subprocess_isolation.py similarity index 81% rename from tests/parallel/_subprocess_isolation.py rename to tests/concurrency/_subprocess_isolation.py index 46191d735..5f1cad234 100644 --- a/tests/parallel/_subprocess_isolation.py +++ b/tests/concurrency/_subprocess_isolation.py @@ -15,21 +15,13 @@ # limitations under the License. """ -Shared test infrastructure for tests/parallel/. +Shared test infrastructure for tests/concurrency/. Tests here deliberately race Producer/Consumer methods against each other, so a real regression can segfault the process instead of just failing an assertion. subprocess_isolated() below runs a decorated test in a fresh `python -m pytest` subprocess, so a crash only fails that one test instead of taking down the whole suite. - -This is a plain module, not conftest.py: conftest.py is auto-loaded by -pytest's own plugin machinery and isn't meant to be imported as a regular -module -- doing so (`from tests.parallel.conftest import ...`) is not -guaranteed to resolve the same way from a re-invoked subprocess as it does -under pytest's own collection, and failed with -"ModuleNotFoundError: No module named 'tests.parallel.conftest'" in -exactly that scenario. """ import functools @@ -37,7 +29,7 @@ import subprocess import sys -_SUBPROCESS_MARKER_ENV = "_PARALLEL_TESTS_SUBPROCESS" +_SUBPROCESS_MARKER_ENV = "_CONCURRENCY_TESTS_SUBPROCESS" _SUBPROCESS_TIMEOUT_SECONDS = 120 diff --git a/tests/parallel/test_producer_close_race.py b/tests/concurrency/test_producer_close_race.py similarity index 96% rename from tests/parallel/test_producer_close_race.py rename to tests/concurrency/test_producer_close_race.py index bff4bdc3e..6fff3f2d5 100644 --- a/tests/parallel/test_producer_close_race.py +++ b/tests/concurrency/test_producer_close_race.py @@ -18,7 +18,7 @@ import time from confluent_kafka import Consumer, Producer, TopicPartition -from tests.parallel._subprocess_isolation import subprocess_isolated +from tests.concurrency._subprocess_isolation import subprocess_isolated ############################################################################### # Tests for races between Producer.close() and concurrent calls to @@ -287,10 +287,7 @@ def run_poll(): assert poll_finished_at is not None, "poll() never finished" # close() should have waited for close to the actual remaining poll() - # duration, not a fixed guess -- a slower/loaded machine can take longer - # to reach close_start, shrinking how much of poll_duration is actually - # left to wait for, so we compare against what really happened instead - # of a hardcoded constant. + # duration close_duration = close_end - close_start remaining_poll_duration = poll_finished_at - close_start assert close_duration >= remaining_poll_duration * 0.9, ( diff --git a/tests/integration/producer/test_concurrency.py b/tests/integration/producer/test_concurrency.py new file mode 100644 index 000000000..b7e95d3c2 --- /dev/null +++ b/tests/integration/producer/test_concurrency.py @@ -0,0 +1,355 @@ +#!/usr/bin/env python +# +# Copyright 2026 Confluent Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import inspect +import threading +import time +from uuid import uuid1 + +from confluent_kafka import KafkaError, KafkaException +from tests.common import TestConsumer + + +def called_by(): + return inspect.stack()[1].function + + +def prefixed_error_cb(prefix): + def error_cb(err): + """Reports global/generic errors to aid in troubleshooting test failures.""" + print("[{}]: {}".format(prefix, err)) + + return error_cb + + +class TestCloseRaceDelivery: + def test_close_delivers_in_flight_messages(self, kafka_cluster): + """ + Messages produced just before/during a concurrent close() are genuinely delivered, + not silently dropped. + """ + topic = kafka_cluster.create_topic_and_wait_propogation("test_close_delivery") + producer = kafka_cluster.producer({'error_cb': prefixed_error_cb('test_close_delivers_in_flight_messages')}) + + delivered = [] + delivery_errors = [] + + def on_delivery(err, msg): + if err: + delivery_errors.append(err) + else: + delivered.append(msg) + + produced_count = 0 + closed_runtime_error = None + + def produce_loop(): + nonlocal produced_count, closed_runtime_error + while True: + try: + producer.produce(topic, value=f'msg-{produced_count}'.encode(), on_delivery=on_delivery) + producer.poll(0) + produced_count += 1 + except RuntimeError as e: + closed_runtime_error = e + break + + t = threading.Thread(target=produce_loop) + t.start() + + # Give the worker thread a moment to actually start producing before + # racing close() against it. + time.sleep(0.1) + + print(f"{called_by()}: calling close() while produce_loop is running") + assert producer.close() is True + + t.join(timeout=30) + assert not t.is_alive(), "producer thread did not finish after close()" + print( + f"{called_by()}: produced_count={produced_count}, " + f"delivered={len(delivered)}, delivery_errors={len(delivery_errors)}" + ) + + # The worker thread must have hit the "Producer has been closed" + # RuntimeError. This proves close() genuinely raced an + # in-flight produce(). + assert closed_runtime_error is not None, "worker thread never hit the closed-producer RuntimeError" + assert not delivery_errors, f"unexpected delivery errors: {delivery_errors}" + assert produced_count > 0, "no messages were produced" + assert ( + len(delivered) == produced_count + ), f"expected all {produced_count} produced messages to be delivered, got {len(delivered)}" + + def test_concurrent_flush_from_multiple_threads(self, kafka_cluster): + """ + Two threads calling flush() concurrently -- + confirms librdkafka's own atomic flush counter (rd_kafka_flush) + correctly waits for all outstanding messages across both callers, + not just its own, and neither returns early. + """ + topic = kafka_cluster.create_topic_and_wait_propogation("test_concurrent_flush") + producer = kafka_cluster.producer( + {'error_cb': prefixed_error_cb('test_concurrent_flush_from_multiple_threads')} + ) + + num_messages = 500 + for i in range(num_messages): + producer.produce(topic, value=f'msg-{i}'.encode()) + producer.poll(0) + + flush_results = [] + + def flush_worker(): + remaining = producer.flush(30) + flush_results.append(remaining) + + threads = [threading.Thread(target=flush_worker) for _ in range(2)] + for t in threads: + t.start() + for t in threads: + t.join(timeout=35) + + print(f"{called_by()}: flush_results={flush_results}, len(producer)={len(producer)}") + assert all(not t.is_alive() for t in threads), "a flush() thread did not finish" + assert all( + r == 0 for r in flush_results + ), f"flush() returned early with messages still outstanding: {flush_results}" + assert len(producer) == 0 + + +class TestTransactionalProducerConcurrency: + def test_concurrent_produce_during_open_transaction(self, kafka_cluster): + """produce() is allowed concurrently from multiple threads while a + transaction is open (only checks an atomic flag, not an exclusive + mutex like the transaction-state APIs).""" + topic = kafka_cluster.create_topic_and_wait_propogation("test_txn_concurrent_produce") + producer = kafka_cluster.producer( + { + 'transactional.id': f'test-txn-concurrent-produce-{uuid1()}', + 'error_cb': prefixed_error_cb('test_concurrent_produce_during_open_transaction'), + } + ) + + producer.init_transactions() + producer.begin_transaction() + + num_threads = 8 + messages_per_thread = 50 + errors = [] + + def produce_worker(thread_id): + try: + for i in range(messages_per_thread): + producer.produce(topic, value=f'thread-{thread_id}-msg-{i}'.encode()) + producer.poll(0) + except Exception as e: # noqa: BLE001 - want to see any exception, not just crashes + errors.append((thread_id, e)) + + threads = [threading.Thread(target=produce_worker, args=(i,)) for i in range(num_threads)] + for t in threads: + t.start() + for t in threads: + t.join(timeout=30) + + print(f"{called_by()}: {num_threads} threads finished producing, errors={errors}") + assert all(not t.is_alive() for t in threads), "a produce() thread did not finish" + assert not errors, f"unexpected exceptions from concurrent produce() during open transaction: {errors}" + + producer.commit_transaction() + + consumer_conf = kafka_cluster.client_conf() + consumer_conf.update( + { + 'group.id': str(uuid1()), + 'auto.offset.reset': 'earliest', + 'enable.auto.commit': False, + 'enable.partition.eof': True, + 'isolation.level': 'read_committed', + } + ) + consumer = TestConsumer(consumer_conf) + consumer.subscribe([topic]) + + msg_cnt = 0 + eof_reached = False + while not eof_reached: + msg = consumer.poll(timeout=10.0) + assert msg is not None, "timed out waiting for messages" + if msg.error(): + if msg.error().code() == KafkaError._PARTITION_EOF: + eof_reached = True + continue + raise KafkaException(msg.error()) + msg_cnt += 1 + consumer.close() + + print(f"{called_by()}: consumed msg_cnt={msg_cnt}") + assert msg_cnt == num_threads * messages_per_thread + + def test_concurrent_calls_to_same_transaction_api(self, kafka_cluster): + """Two threads both calling commit_transaction() at once: exactly + one succeeds, the other gets a clean _PREV_IN_PROGRESS error.""" + topic = kafka_cluster.create_topic_and_wait_propogation("test_txn_same_api_race") + producer = kafka_cluster.producer( + { + 'transactional.id': f'test-txn-same-api-race-{uuid1()}', + 'error_cb': prefixed_error_cb('test_concurrent_calls_to_same_transaction_api'), + } + ) + + producer.init_transactions() + producer.begin_transaction() + producer.produce(topic, value=b'msg') + producer.flush() + + results = [] + barrier = threading.Barrier(2) + + def call_commit(): + barrier.wait() + try: + producer.commit_transaction() + results.append(True) + except KafkaException as e: + results.append(e.args[0]) + + threads = [threading.Thread(target=call_commit) for _ in range(2)] + for t in threads: + t.start() + for t in threads: + t.join(timeout=30) + + print(f"{called_by()}: results={results}") + assert all(not t.is_alive() for t in threads), "a commit_transaction() thread did not finish" + assert len(results) == 2 + successes = [r for r in results if r is True] + errors = [r for r in results if r is not True] + assert len(successes) == 1, f"expected exactly one successful commit, got: {results}" + assert len(errors) == 1, f"expected exactly one error result, got: {results}" + + def test_close_races_open_transaction(self, kafka_cluster): + """close() concurrent with an open (uncommitted) transaction: + close() must not crash or hang, and since the transaction was + never committed, none of its messages should become visible to + a read_committed consumer.""" + topic = kafka_cluster.create_topic_and_wait_propogation("test_txn_close_race") + producer = kafka_cluster.producer( + { + 'transactional.id': f'test-txn-close-race-{uuid1()}', + 'error_cb': prefixed_error_cb('test_close_races_open_transaction'), + } + ) + + producer.init_transactions() + producer.begin_transaction() + + produced_count = 0 + closed_runtime_error = None + + def produce_loop(): + nonlocal produced_count, closed_runtime_error + while True: + try: + producer.produce(topic, value=f'msg-{produced_count}'.encode()) + producer.poll(0) + produced_count += 1 + except RuntimeError as e: + closed_runtime_error = e + break + + t = threading.Thread(target=produce_loop) + t.start() + + time.sleep(0.1) + print(f"{called_by()}: calling close() while a transaction is still open") + assert producer.close() is True + + t.join(timeout=30) + assert not t.is_alive(), "producer thread did not finish after close()" + print(f"{called_by()}: produced_count={produced_count} before close() won the race") + assert closed_runtime_error is not None, "worker thread never hit the closed-producer RuntimeError" + assert produced_count > 0, "no messages were produced before close()" + + consumer_conf = kafka_cluster.client_conf() + consumer_conf.update( + { + 'group.id': str(uuid1()), + 'auto.offset.reset': 'earliest', + 'enable.auto.commit': False, + 'enable.partition.eof': True, + 'isolation.level': 'read_committed', + } + ) + consumer = TestConsumer(consumer_conf) + consumer.subscribe([topic]) + + msg = consumer.poll(timeout=10.0) + consumer.close() + + print(f"{called_by()}: consumer.poll() returned error={msg.error() if msg else None}") + assert msg is not None, "timed out waiting for a message/EOF" + assert msg.error() is not None and msg.error().code() == KafkaError._PARTITION_EOF, ( + "an uncommitted transaction's messages must not be visible to a " "read_committed consumer after close()" + ) + + def test_concurrent_calls_to_different_transaction_apis(self, kafka_cluster): + """One thread calls commit_transaction() while another calls + abort_transaction() at once: exactly one succeeds, the other gets + a clean _CONFLICT error.""" + topic = kafka_cluster.create_topic_and_wait_propogation("test_txn_diff_api_race") + producer = kafka_cluster.producer( + { + 'transactional.id': f'test-txn-diff-api-race-{uuid1()}', + 'error_cb': prefixed_error_cb('test_concurrent_calls_to_different_transaction_apis'), + } + ) + + producer.init_transactions() + producer.begin_transaction() + producer.produce(topic, value=b'msg') + producer.flush() + + results = {} + barrier = threading.Barrier(2) + + def call_commit(): + barrier.wait() + try: + producer.commit_transaction() + results['commit'] = True + except KafkaException as e: + results['commit'] = e.args[0] + + def call_abort(): + barrier.wait() + try: + producer.abort_transaction() + results['abort'] = True + except KafkaException as e: + results['abort'] = e.args[0] + + t1 = threading.Thread(target=call_commit) + t2 = threading.Thread(target=call_abort) + t1.start() + t2.start() + t1.join(timeout=30) + t2.join(timeout=30) + + print(f"{called_by()}: results={results}") + assert not t1.is_alive() and not t2.is_alive(), "a transaction-ending thread did not finish" + successes = [k for k, v in results.items() if v is True] + assert len(successes) == 1, f"expected exactly one of commit/abort to succeed, got: {results}" From 48514fb5bc3ab025537a77601485fe2238ad4e98 Mon Sep 17 00:00:00 2001 From: Ojasva Jain Date: Mon, 27 Jul 2026 20:24:48 +0530 Subject: [PATCH 05/10] Clarified comment --- src/confluent_kafka/src/confluent_kafka.h | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/src/confluent_kafka/src/confluent_kafka.h b/src/confluent_kafka/src/confluent_kafka.h index 34746d845..78c794713 100644 --- a/src/confluent_kafka/src/confluent_kafka.h +++ b/src/confluent_kafka/src/confluent_kafka.h @@ -292,9 +292,10 @@ typedef struct { PyObject *oauth_cb; atomic_int_t oauth_token_set; - /* Protects self->rk from being freed by close() while another method - * is still using it. See Handle_enter_rk_use()/Handle_exit_rk_use() - * in confluent_kafka.c. */ + /* Protects self->rk in Producer and Admin clients from being freed by + * close() while another method is still using it. + * See Handle_enter_rk_use()/Handle_exit_rk_use() in confluent_kafka.c. + */ atomic_int_t active_calls; atomic_int_t closing; From 3346b6c9e8c148dbb83a720d2108c365b8c3bf28 Mon Sep 17 00:00:00 2001 From: Ojasva Jain Date: Fri, 31 Jul 2026 13:43:35 +0530 Subject: [PATCH 06/10] Make Producer.close() non-blocking for concurrent callers, add reentrancy tests Concurrent close() calls now return False immediately with a warning instead of waiting for the CAS winner, since waiting could deadlock a caller that already holds an active_calls slot (e.g. a callback invoked from its own poll()/flush()). poll()/flush() now also exit early once closing is set instead of blocking the drain-wait. Fixes an ordering bug in produce_batch() where the topic handle was destroyed after releasing the active_calls slot. Adds integration tests for reentrant callbacks and close()'s internal flush delivering all messages, and documents close()-from-callback as unsupported. --- src/confluent_kafka/cimpl.pyi | 1 + src/confluent_kafka/src/Producer.c | 173 +++++++++------- src/confluent_kafka/src/confluent_kafka.c | 3 - tests/concurrency/test_producer_close_race.py | 145 ++++++++++--- .../integration/producer/test_concurrency.py | 194 ++++++++++++++++++ 5 files changed, 407 insertions(+), 109 deletions(-) diff --git a/src/confluent_kafka/cimpl.pyi b/src/confluent_kafka/cimpl.pyi index 7dd339748..d6afd980d 100644 --- a/src/confluent_kafka/cimpl.pyi +++ b/src/confluent_kafka/cimpl.pyi @@ -421,6 +421,7 @@ class Producer: ) -> None: ... def list_topics(self, topic: Optional[str] = None, timeout: float = -1) -> Any: ... def set_sasl_credentials(self, username: str, password: str) -> None: ... + def close(self) -> bool: ... def __len__(self) -> int: ... def __bool__(self) -> bool: ... def __enter__(self) -> Self: ... diff --git a/src/confluent_kafka/src/Producer.c b/src/confluent_kafka/src/Producer.c index 8a104ee51..cfc267e5f 100644 --- a/src/confluent_kafka/src/Producer.c +++ b/src/confluent_kafka/src/Producer.c @@ -405,6 +405,11 @@ static int Producer_poll0(Handle *self, int tmout) { chunk_count++; + /* Exit early if a close() has been initiated */ + if (atomic_int_get(&self->closing)) { + break; + } + /* Check for signals between chunks */ if (check_signals_between_chunks(self, &cs)) { return -1; /* Signal detected */ @@ -473,6 +478,7 @@ Producer_flush(Handle *self, PyObject *args, PyObject *kwargs) { int total_timeout_ms; int chunk_timeout_ms; int chunk_count = 0; + PyObject *result = NULL; /* NULL means an exception is already set */ if (!PyArg_ParseTupleAndKeywords(args, kwargs, "|d", kws, &tmout)) return NULL; @@ -509,18 +515,24 @@ Producer_flush(Handle *self, PyObject *args, PyObject *kwargs) { /* Flush with chunk timeout */ err = rd_kafka_flush(self->rk, chunk_timeout_ms); - /* Always check for signals between chunks (critical for - * interruptibility) */ chunk_count++; - if (check_signals_between_chunks(self, &cs)) { - Handle_exit_rk_use(self); - return NULL; /* Signal detected */ - } if (err == RD_KAFKA_RESP_ERR_NO_ERROR) { break; } + /* Exit early if a close() has been initiated */ + if (atomic_int_get(&self->closing)) { + err = RD_KAFKA_RESP_ERR__TIMED_OUT; + break; + } + + /* Always check for signals between chunks (critical for + * interruptibility) */ + if (check_signals_between_chunks(self, &cs)) { + goto exit; /* Signal detected, result stays NULL */ + } + /* If timeout error, continue to next chunk */ if (err == RD_KAFKA_RESP_ERR__TIMED_OUT) { continue; @@ -531,17 +543,17 @@ Producer_flush(Handle *self, PyObject *args, PyObject *kwargs) { } } - if (!CallState_end(self, &cs)) { - Handle_exit_rk_use(self); - return NULL; - } + if (!CallState_end(self, &cs)) + goto exit; /* result stays NULL */ if (err) /* Get the queue length on error (timeout) */ qlen = rd_kafka_outq_len(self->rk); - Handle_exit_rk_use(self); + result = cfl_PyInt_FromInt(qlen); - return cfl_PyInt_FromInt(qlen); +exit: + Handle_exit_rk_use(self); + return result; } @@ -553,27 +565,20 @@ Producer_close(Handle *self, PyObject *args, PyObject *kwargs) { if (!self->rk) Py_RETURN_TRUE; - /* Only one concurrent close() can destroy rk, otherwise, - * two threads could both reach rd_kafka_destroy() on the - * same handle (a double-free). The losing thread(s) wait for the - * winner to finish and then return True, same as a normal close(), - * rather than racing it. */ + /* If there are concurrent calls to close(), only one of them can + * destroy rk. The remaining close() calls return False immediately with + * a warning. */ if (!atomic_int_cas(&self->closing, 0, 1)) { - while (self->rk) { - CallState_begin(self, &cs); -#ifdef _WIN32 - Sleep(100); -#else - usleep(100000); -#endif - CallState_end(self, &cs); - } - Py_RETURN_TRUE; + PyErr_WarnFormat(PyExc_RuntimeWarning, 1, + "Producer is already closing"); + Py_RETURN_FALSE; } /* Signal in-flight calls to stop, and wait for them to finish * using self->rk before destroying it -- see Handle_enter_rk_use(). * New calls will see `closing` and fail with ERR_MSG_PRODUCER_CLOSED. */ + /* TODO NOGIL: replace this poll loop with a mutex/condvar wait so + * close() unblocks immediately instead of up to 100ms late. */ while (atomic_int_get(&self->active_calls) > 0) { CallState_begin(self, &cs); #ifdef _WIN32 @@ -581,7 +586,8 @@ Producer_close(Handle *self, PyObject *args, PyObject *kwargs) { #else usleep(100000); #endif - CallState_end(self, &cs); + if (!CallState_end(self, &cs)) + return NULL; } CallState_begin(self, &cs); @@ -889,11 +895,12 @@ Producer_produce_batch(Handle *self, PyObject *args, PyObject *kwargs) { messages_list, rkt, partition, rkmessages, msgstates, message_cnt); cleanup: - Handle_exit_rk_use(self); - - /* Cleanup resources */ if (rkt) rd_kafka_topic_destroy(rkt); + + Handle_exit_rk_use(self); + + /* Cleanup resources not tied to self->rk */ if (rkmessages) free(rkmessages); if (msgstates) @@ -908,7 +915,8 @@ Producer_produce_batch(Handle *self, PyObject *args, PyObject *kwargs) { static PyObject *Producer_init_transactions(Handle *self, PyObject *args) { CallState cs; rd_kafka_error_t *error; - double tmout = -1.0; + double tmout = -1.0; + PyObject *result = NULL; /* NULL means an exception is already set */ if (!PyArg_ParseTuple(args, "|d", &tmout)) return NULL; @@ -921,20 +929,23 @@ static PyObject *Producer_init_transactions(Handle *self, PyObject *args) { error = rd_kafka_init_transactions(self->rk, cfl_timeout_ms(tmout)); if (!CallState_end(self, &cs)) { - Handle_exit_rk_use(self); - if (error) /* Ignore error in favour of callstate exception */ + /* Ignore error in favour of callstate exception */ + if (error) rd_kafka_error_destroy(error); - return NULL; + goto exit; } - Handle_exit_rk_use(self); - if (error) { cfl_PyErr_from_error_destroy(error); - return NULL; + goto exit; } - Py_RETURN_NONE; + result = Py_None; + Py_INCREF(result); + +exit: + Handle_exit_rk_use(self); + return result; } static PyObject *Producer_begin_transaction(Handle *self) { @@ -960,9 +971,10 @@ static PyObject *Producer_send_offsets_to_transaction(Handle *self, CallState cs; rd_kafka_error_t *error; PyObject *metadata = NULL, *offsets = NULL; - rd_kafka_topic_partition_list_t *c_offsets; - rd_kafka_consumer_group_metadata_t *cgmd; - double tmout = -1.0; + rd_kafka_topic_partition_list_t *c_offsets = NULL; + rd_kafka_consumer_group_metadata_t *cgmd = NULL; + double tmout = -1.0; + PyObject *result = NULL; /* NULL means an exception is already set */ if (!PyArg_ParseTuple(args, "OO|d", &offsets, &metadata, &tmout)) return NULL; @@ -970,46 +982,46 @@ static PyObject *Producer_send_offsets_to_transaction(Handle *self, if (!Handle_enter_rk_use(self)) return NULL; - if (!(c_offsets = py_to_c_parts(offsets))) { - Handle_exit_rk_use(self); - return NULL; - } + if (!(c_offsets = py_to_c_parts(offsets))) + goto exit; - if (!(cgmd = py_to_c_cgmd(metadata))) { - rd_kafka_topic_partition_list_destroy(c_offsets); - Handle_exit_rk_use(self); - return NULL; - } + if (!(cgmd = py_to_c_cgmd(metadata))) + goto exit; CallState_begin(self, &cs); error = rd_kafka_send_offsets_to_transaction(self->rk, c_offsets, cgmd, cfl_timeout_ms(tmout)); - rd_kafka_consumer_group_metadata_destroy(cgmd); - rd_kafka_topic_partition_list_destroy(c_offsets); - if (!CallState_end(self, &cs)) { - Handle_exit_rk_use(self); - if (error) /* Ignore error in favour of callstate exception */ + /* Ignore error in favour of callstate exception */ + if (error) rd_kafka_error_destroy(error); - return NULL; + goto exit; } - Handle_exit_rk_use(self); - if (error) { cfl_PyErr_from_error_destroy(error); - return NULL; + goto exit; } - Py_RETURN_NONE; + result = Py_None; + Py_INCREF(result); + +exit: + if (cgmd) + rd_kafka_consumer_group_metadata_destroy(cgmd); + if (c_offsets) + rd_kafka_topic_partition_list_destroy(c_offsets); + Handle_exit_rk_use(self); + return result; } static PyObject *Producer_commit_transaction(Handle *self, PyObject *args) { CallState cs; rd_kafka_error_t *error; - double tmout = -1.0; + double tmout = -1.0; + PyObject *result = NULL; /* NULL means an exception is already set */ if (!PyArg_ParseTuple(args, "|d", &tmout)) return NULL; @@ -1022,26 +1034,30 @@ static PyObject *Producer_commit_transaction(Handle *self, PyObject *args) { error = rd_kafka_commit_transaction(self->rk, cfl_timeout_ms(tmout)); if (!CallState_end(self, &cs)) { - Handle_exit_rk_use(self); - if (error) /* Ignore error in favour of callstate exception */ + /* Ignore error in favour of callstate exception */ + if (error) rd_kafka_error_destroy(error); - return NULL; + goto exit; } - Handle_exit_rk_use(self); - if (error) { cfl_PyErr_from_error_destroy(error); - return NULL; + goto exit; } - Py_RETURN_NONE; + result = Py_None; + Py_INCREF(result); + +exit: + Handle_exit_rk_use(self); + return result; } static PyObject *Producer_abort_transaction(Handle *self, PyObject *args) { CallState cs; rd_kafka_error_t *error; - double tmout = -1.0; + double tmout = -1.0; + PyObject *result = NULL; /* NULL means an exception is already set */ if (!PyArg_ParseTuple(args, "|d", &tmout)) return NULL; @@ -1054,20 +1070,23 @@ static PyObject *Producer_abort_transaction(Handle *self, PyObject *args) { error = rd_kafka_abort_transaction(self->rk, cfl_timeout_ms(tmout)); if (!CallState_end(self, &cs)) { - Handle_exit_rk_use(self); - if (error) /* Ignore error in favour of callstate exception */ + /* Ignore error in favour of callstate exception */ + if (error) rd_kafka_error_destroy(error); - return NULL; + goto exit; } - Handle_exit_rk_use(self); - if (error) { cfl_PyErr_from_error_destroy(error); - return NULL; + goto exit; } - Py_RETURN_NONE; + result = Py_None; + Py_INCREF(result); + +exit: + Handle_exit_rk_use(self); + return result; } static void *Producer_purge(Handle *self, PyObject *args, PyObject *kwargs) { diff --git a/src/confluent_kafka/src/confluent_kafka.c b/src/confluent_kafka/src/confluent_kafka.c index 29d5dd9ab..1b0c7fa68 100644 --- a/src/confluent_kafka/src/confluent_kafka.c +++ b/src/confluent_kafka/src/confluent_kafka.c @@ -3324,9 +3324,6 @@ int CallState_end(Handle *h, CallState *cs) { * incremented; caller must call Handle_exit_rk_use() on every * return path), or 0 with ERR_MSG_PRODUCER_CLOSED set if the * Handle is closed/closing (nothing to undo). - * - * @warning Not re-entrant: don't call from a method that's already - * between its own Handle_enter_rk_use()/Handle_exit_rk_use(). */ int Handle_enter_rk_use(Handle *h) { if (atomic_int_get(&h->closing) || !h->rk) { diff --git a/tests/concurrency/test_producer_close_race.py b/tests/concurrency/test_producer_close_race.py index 6fff3f2d5..ef0a5b3be 100644 --- a/tests/concurrency/test_producer_close_race.py +++ b/tests/concurrency/test_producer_close_race.py @@ -14,21 +14,24 @@ # See the License for the specific language governing permissions and # limitations under the License. +import os +import signal import threading import time from confluent_kafka import Consumer, Producer, TopicPartition from tests.concurrency._subprocess_isolation import subprocess_isolated +_PRODUCER_CONF = {'bootstrap.servers': 'localhost:9092', 'socket.timeout.ms': 10, 'message.timeout.ms': 10} +_TXN_PRODUCER_CONF = dict(_PRODUCER_CONF, **{'transactional.id': 'test-producer-close-race-txn'}) +ITERATIONS = 10 + + ############################################################################### # Tests for races between Producer.close() and concurrent calls to # other methods on the same Producer instance. ############################################################################### -_PRODUCER_CONF = {'bootstrap.servers': 'localhost:9092', 'socket.timeout.ms': 10, 'message.timeout.ms': 10} -_TXN_PRODUCER_CONF = dict(_PRODUCER_CONF, **{'transactional.id': 'test-producer-close-race-txn'}) -ITERATIONS = 10 - def _race_close_against(worker, conf=None, num_workers=1): """ @@ -241,58 +244,142 @@ def worker(producer, stop_event): @subprocess_isolated def test_close_races_close(): - """Multiple threads calling close() on the same Producer at once.""" - worker_close_results = [] + """Multiple threads calling close() on the same Producer at once. The + CAS winner (whichever thread's close() actually tears down self->rk) + returns True; a losing thread returns False *unless* its own check + happens to run after self->rk is already NULL (the winner fully + completed first), in which case it also gets True.""" num_workers = 7 - def worker(producer, stop_event): - worker_close_results.append(producer.close()) + for i in range(ITERATIONS): + producer = Producer(_PRODUCER_CONF) + all_results = [] + start_barrier = threading.Barrier(num_workers + 1) - _race_close_against(worker, num_workers=num_workers) + def worker(): + start_barrier.wait() + all_results.append(producer.close()) - assert all(worker_close_results), f"not every worker close() call returned True: {worker_close_results}" - assert len(worker_close_results) == num_workers * ITERATIONS + threads = [threading.Thread(target=worker) for _ in range(num_workers)] + for t in threads: + t.start() + + start_barrier.wait() + all_results.append(producer.close()) + + for t in threads: + t.join(timeout=10) + assert all(not t.is_alive() for t in threads), f"iteration {i}: a close() thread did not finish" + + assert ( + len(all_results) == num_workers + 1 + ), f"iteration {i}: expected {num_workers + 1} results, got {all_results}" + assert all( + isinstance(r, bool) for r in all_results + ), f"iteration {i}: every close() call must return True or False, got: {all_results}" + assert any(all_results), f"iteration {i}: expected at least one close() call to return True, got: {all_results}" -############################################################################### -# close() blocks until an in-flight call finishes, -# rather than just not crashing while one is running. ############################################################################### -def test_close_waits_for_in_flight_call(): - """close() blocks until an in-flight poll() call finishes.""" +def test_close_completes_quickly_with_indefinite_poll_in_progress(): + """close() must not block indefinitely behind an in-flight poll(-1) + call on another thread -- poll()'s chunk loop notices `closing` and + exits early, so close() should complete within a small, bounded time + instead of waiting for poll() to return on its own.""" producer = Producer(_PRODUCER_CONF) poll_started = threading.Event() - poll_duration = 10 poll_finished_at = None def run_poll(): nonlocal poll_finished_at poll_started.set() - producer.poll(poll_duration) + producer.poll(-1) poll_finished_at = time.monotonic() t = threading.Thread(target=run_poll) t.start() poll_started.wait() - time.sleep(2) + time.sleep(0.5) # Make sure poll() is genuinely in-flight when close() fires. close_start = time.monotonic() producer.close() - close_end = time.monotonic() + close_duration = time.monotonic() - close_start t.join(timeout=10) assert not t.is_alive(), "poll() thread did not finish after close()" assert poll_finished_at is not None, "poll() never finished" + assert close_duration < 0.5, ( + f"close() took {close_duration:.2f}s to complete while poll(-1) was " + f"in progress -- expected it to finish within 0.5s" + ) - # close() should have waited for close to the actual remaining poll() - # duration - close_duration = close_end - close_start - remaining_poll_duration = poll_finished_at - close_start - assert close_duration >= remaining_poll_duration * 0.9, ( - f"close() took {close_duration:.2f}s but the in-flight poll() call " - f"was still going to run for {remaining_poll_duration:.2f}s more " - f"-- close() returned too early relative to what it should have " - f"waited for" + +def test_close_completes_quickly_with_indefinite_flush_in_progress(): + """close() must not block indefinitely behind an in-flight flush(-1) + call on another thread -- flush()'s chunk loop notices `closing` and + exits early, so close() should complete within a small, bounded time + instead of waiting for flush() to return on its own.""" + producer = Producer(_PRODUCER_CONF) + flush_started = threading.Event() + flush_finished_at = None + + def run_flush(): + nonlocal flush_finished_at + flush_started.set() + producer.flush(-1) + flush_finished_at = time.monotonic() + + t = threading.Thread(target=run_flush) + t.start() + flush_started.wait() + time.sleep(0.5) # Make sure flush() is genuinely in-flight when close() fires. + + close_start = time.monotonic() + producer.close() + close_duration = time.monotonic() - close_start + + t.join(timeout=10) + assert not t.is_alive(), "flush() thread did not finish after close()" + assert flush_finished_at is not None, "flush() never finished" + assert close_duration < 0.5, ( + f"close() took {close_duration:.2f}s to complete while flush(-1) was " + f"in progress -- expected it to finish within 0.5s" ) + + +@subprocess_isolated +def test_close_propagates_signal_while_waiting_for_active_calls(): + """close() must raise KeyboardInterrupt if a signal arrives while its + active_calls drain-wait loop is spinning""" + producer = Producer(_PRODUCER_CONF) + poll_started = threading.Event() + + def hold_active_call(): + poll_started.set() + try: + producer.poll(-1) + except BaseException: # noqa: BLE001 - just draining the thread, not asserting here + pass + + t = threading.Thread(target=hold_active_call) + t.start() + poll_started.wait() + + def send_sigint_soon(): + time.sleep(0.05) + os.kill(os.getpid(), signal.SIGINT) + + interrupt_thread = threading.Thread(target=send_sigint_soon) + interrupt_thread.daemon = True + interrupt_thread.start() + + try: + producer.close() + assert False, "close() returned normally instead of raising KeyboardInterrupt" + except KeyboardInterrupt: + assert True # expected outcome: close() correctly propagated the signal + finally: + t.join(timeout=10) + assert not t.is_alive(), "poll() thread did not finish after close() was interrupted" diff --git a/tests/integration/producer/test_concurrency.py b/tests/integration/producer/test_concurrency.py index b7e95d3c2..c0cb47062 100644 --- a/tests/integration/producer/test_concurrency.py +++ b/tests/integration/producer/test_concurrency.py @@ -94,6 +94,41 @@ def produce_loop(): len(delivered) == produced_count ), f"expected all {produced_count} produced messages to be delivered, got {len(delivered)}" + def test_close_internal_flush_delivers_all_undelivered_messages(self, kafka_cluster): + """Produce a batch of messages with no poll() in between (so none + of their delivery reports have been serviced yet), then call + close() directly. close()'s own internal flush must drive + delivery of every one of them.""" + topic = kafka_cluster.create_topic_and_wait_propogation("test_close_internal_flush_delivery") + producer = kafka_cluster.producer( + {'error_cb': prefixed_error_cb('test_close_internal_flush_delivers_all_undelivered_messages')} + ) + + num_messages = 200 + delivered = [] + delivery_errors = [] + + def on_delivery(err, msg): + if err: + delivery_errors.append(err) + else: + delivered.append(msg) + + for i in range(num_messages): + producer.produce(topic, value=f'msg-{i}'.encode(), on_delivery=on_delivery) + # Deliberately no poll() here + assert len(producer) == num_messages, "messages must still be queued (undelivered) when close() is called" + + print(f"{called_by()}: calling close() to flush {num_messages} undelivered messages") + result = producer.close() + + print( + f"{called_by()}: close()={result}, delivered={len(delivered)}, " f"delivery_errors={len(delivery_errors)}" + ) + assert result is True + assert not delivery_errors, f"unexpected delivery errors: {delivery_errors}" + assert len(delivered) == num_messages, f"expected all {num_messages} messages delivered, got {len(delivered)}" + def test_concurrent_flush_from_multiple_threads(self, kafka_cluster): """ Two threads calling flush() concurrently -- @@ -353,3 +388,162 @@ def call_abort(): assert not t1.is_alive() and not t2.is_alive(), "a transaction-ending thread did not finish" successes = [k for k, v in results.items() if v is True] assert len(successes) == 1, f"expected exactly one of commit/abort to succeed, got: {results}" + + +class TestReentrantDeliveryCallback: + """Delivery callbacks run synchronously inside poll()/flush() on + whatever thread called them. These tests cover a callback calling back + into the Producer it belongs to, from that same call chain. + + Note: calling close() from within a callback is NOT covered here and is + NOT supported.""" + + def test_delivery_callback_producing_another_message_gets_delivered(self, kafka_cluster): + """A delivery callback calling produce() again must succeed and the newly produced message must + itself be delivered""" + topic = kafka_cluster.create_topic_and_wait_propogation("test_reentrant_produce_from_callback") + producer = kafka_cluster.producer( + {'error_cb': prefixed_error_cb('test_delivery_callback_producing_another_message_gets_delivered')} + ) + + first_delivered = [] + second_delivered = [] + produce_again_error = [] + + def on_second_delivery(err, msg): + if err: + produce_again_error.append(err) + else: + second_delivered.append(msg) + + def on_first_delivery(err, msg): + if err: + return + first_delivered.append(msg) + try: + producer.produce(topic, value=b'reentrant-produce', on_delivery=on_second_delivery) + except Exception as e: # noqa: BLE001 - want to see any exception, not just crashes + produce_again_error.append(e) + + producer.produce(topic, value=b'original', on_delivery=on_first_delivery) + producer.flush(30) + + print( + f"{called_by()}: first_delivered={len(first_delivered)}, " + f"second_delivered={len(second_delivered)}, errors={produce_again_error}" + ) + assert len(first_delivered) == 1, "the original message must be delivered" + assert not produce_again_error, f"reentrant produce() from the callback failed: {produce_again_error}" + assert len(second_delivered) == 1, "the message produced from within the callback must itself be delivered" + + producer.close() + + def test_delivery_callback_calling_poll_does_not_crash_or_hang(self, kafka_cluster): + """A delivery callback calling poll() again, reentrantly, on the + same thread that's already inside the outer poll() that invoked it.""" + topic = kafka_cluster.create_topic_and_wait_propogation("test_reentrant_poll_from_callback") + producer = kafka_cluster.producer( + {'error_cb': prefixed_error_cb('test_delivery_callback_calling_poll_does_not_crash_or_hang')} + ) + + delivered = [] + reentrant_poll_results = [] + + def on_delivery(err, msg): + if not err: + delivered.append(msg) + try: + reentrant_poll_results.append(producer.poll(0)) + except Exception as e: # noqa: BLE001 - want to see any exception, not just crashes + reentrant_poll_results.append(e) + + producer.produce(topic, value=b'original', on_delivery=on_delivery) + producer.flush(30) + + print(f"{called_by()}: delivered={len(delivered)}, " f"reentrant_poll_results={reentrant_poll_results}") + assert len(delivered) == 1, "the message must be delivered" + assert len(reentrant_poll_results) == 1, "the delivery callback must have fired exactly once" + assert not isinstance( + reentrant_poll_results[0], Exception + ), f"reentrant poll() from the callback raised: {reentrant_poll_results[0]}" + + producer.close() + + def test_multiple_threads_reentrantly_producing_from_callbacks(self, kafka_cluster): + """Several threads, each producing on the same shared Producer and + each individually reentering produce() from its own delivery + callback""" + topic = kafka_cluster.create_topic_and_wait_propogation("test_reentrant_produce_multi_thread") + producer = kafka_cluster.producer( + {'error_cb': prefixed_error_cb('test_multiple_threads_reentrantly_producing_from_callbacks')} + ) + + num_threads = 8 + messages_per_thread = 20 + first_delivered = [] + second_delivered = [] + errors = [] + lock = threading.Lock() + + def make_second_callback(): + def on_second_delivery(err, msg): + with lock: + if err: + errors.append(err) + else: + second_delivered.append(msg) + + return on_second_delivery + + def make_first_callback(thread_id, i): + def on_first_delivery(err, msg): + if err: + with lock: + errors.append(err) + return + with lock: + first_delivered.append(msg) + try: + producer.produce( + topic, + value=f'reentrant-{thread_id}-{i}'.encode(), + on_delivery=make_second_callback(), + ) + except Exception as e: # noqa: BLE001 - want to see any exception, not just crashes + with lock: + errors.append(e) + + return on_first_delivery + + def worker(thread_id): + for i in range(messages_per_thread): + producer.produce( + topic, + value=f'original-{thread_id}-{i}'.encode(), + on_delivery=make_first_callback(thread_id, i), + ) + producer.poll(0) + + threads = [threading.Thread(target=worker, args=(i,)) for i in range(num_threads)] + for t in threads: + t.start() + for t in threads: + t.join(timeout=30) + + producer.flush(30) + + expected_total = num_threads * messages_per_thread + print( + f"{called_by()}: first_delivered={len(first_delivered)}, " + f"second_delivered={len(second_delivered)}, errors={errors}" + ) + assert all(not t.is_alive() for t in threads), "a producing thread did not finish" + assert not errors, f"unexpected errors from concurrent reentrant produce(): {errors}" + assert ( + len(first_delivered) == expected_total + ), f"expected all {expected_total} original messages delivered, got {len(first_delivered)}" + assert ( + len(second_delivered) == expected_total + ), f"expected all {expected_total} reentrantly-produced messages delivered, got {len(second_delivered)}" + + producer.close() From dd0bcd57166cdaf4822bf498bde0697b8326356e Mon Sep 17 00:00:00 2001 From: Ojasva Jain Date: Tue, 4 Aug 2026 18:12:37 +0530 Subject: [PATCH 07/10] Fix close() losers to wait for winner instead of returning False early --- src/confluent_kafka/src/Producer.c | 38 +++++-- tests/concurrency/test_producer_close_race.py | 102 ++++++++++++++---- 2 files changed, 116 insertions(+), 24 deletions(-) diff --git a/src/confluent_kafka/src/Producer.c b/src/confluent_kafka/src/Producer.c index cfc267e5f..8bf201893 100644 --- a/src/confluent_kafka/src/Producer.c +++ b/src/confluent_kafka/src/Producer.c @@ -565,13 +565,34 @@ Producer_close(Handle *self, PyObject *args, PyObject *kwargs) { if (!self->rk) Py_RETURN_TRUE; + /* Calling close() reentrantly from within a callback + * is not supported and will deadlock here. + * TODO NOGIL: Update documentation to highlight this. + */ + /* If there are concurrent calls to close(), only one of them can - * destroy rk. The remaining close() calls return False immediately with - * a warning. */ + * destroy rk, the rest wait here for the winner to finish + * flushing and destroying it. + */ if (!atomic_int_cas(&self->closing, 0, 1)) { - PyErr_WarnFormat(PyExc_RuntimeWarning, 1, - "Producer is already closing"); - Py_RETURN_FALSE; + while (self->rk && atomic_int_get(&self->closing)) { + CallState_begin(self, &cs); +#ifdef _WIN32 + Sleep(100); +#else + usleep(100000); +#endif + if (!CallState_end(self, &cs)) + return NULL; + } + if (!self->rk) + Py_RETURN_TRUE; + + /* The winner got interrupted by a signal */ + PyErr_SetString(PyExc_RuntimeError, + "close() was interrupted by a signal on " + "another thread"); + return NULL; } /* Signal in-flight calls to stop, and wait for them to finish @@ -586,8 +607,13 @@ Producer_close(Handle *self, PyObject *args, PyObject *kwargs) { #else usleep(100000); #endif - if (!CallState_end(self, &cs)) + if (!CallState_end(self, &cs)) { + /* Abort the attempt: rk was never touched, so + * reopen the gate for a future close() attempt. + */ + atomic_int_set(&self->closing, 0); return NULL; + } } CallState_begin(self, &cs); diff --git a/tests/concurrency/test_producer_close_race.py b/tests/concurrency/test_producer_close_race.py index ef0a5b3be..1b2d53e36 100644 --- a/tests/concurrency/test_producer_close_race.py +++ b/tests/concurrency/test_producer_close_race.py @@ -152,8 +152,6 @@ def worker(producer, stop_event): producer.init_transactions(0.05) except RuntimeError: break - except Exception: # noqa: BLE001 - librdkafka state/timeout errors are expected without a broker - pass _race_close_against(worker, conf=_TXN_PRODUCER_CONF) @@ -168,8 +166,6 @@ def worker(producer, stop_event): producer.begin_transaction() except RuntimeError: break - except Exception: # noqa: BLE001 - librdkafka state errors are expected without a broker - pass _race_close_against(worker, conf=_TXN_PRODUCER_CONF) @@ -184,8 +180,6 @@ def worker(producer, stop_event): producer.commit_transaction(0.05) except RuntimeError: break - except Exception: # noqa: BLE001 - librdkafka state/timeout errors are expected without a broker - pass _race_close_against(worker, conf=_TXN_PRODUCER_CONF) @@ -200,8 +194,6 @@ def worker(producer, stop_event): producer.abort_transaction(0.05) except RuntimeError: break - except Exception: # noqa: BLE001 - librdkafka state/timeout errors are expected without a broker - pass _race_close_against(worker, conf=_TXN_PRODUCER_CONF) @@ -222,8 +214,6 @@ def worker(producer, stop_event): producer.send_offsets_to_transaction(offsets, metadata, 0.05) except RuntimeError: break - except Exception: # noqa: BLE001 - librdkafka state/timeout errors are expected without a broker - pass _race_close_against(worker, conf=_TXN_PRODUCER_CONF) @@ -245,10 +235,8 @@ def worker(producer, stop_event): @subprocess_isolated def test_close_races_close(): """Multiple threads calling close() on the same Producer at once. The - CAS winner (whichever thread's close() actually tears down self->rk) - returns True; a losing thread returns False *unless* its own check - happens to run after self->rk is already NULL (the winner fully - completed first), in which case it also gets True.""" + CAS winner tears down self->rk itself; every losing thread blocks until + the winner finishes and then returns True too, same as a normal close().""" num_workers = 7 for i in range(ITERATIONS): @@ -274,13 +262,74 @@ def worker(): assert ( len(all_results) == num_workers + 1 ), f"iteration {i}: expected {num_workers + 1} results, got {all_results}" + assert all(all_results), f"iteration {i}: every concurrent close() call must return True, got: {all_results}" + + +@subprocess_isolated +def test_close_races_close_losers_wait_for_slow_winner(): + """Losing close() calls must actually block until the winner finishes, + not just happen to observe self->rk already NULL. A held poll(-1) call + keeps active_calls > 0, forcing the CAS winner (and therefore every + loser waiting behind it) to spin through its drain-wait loop. Proven by + checking every close() thread is still alive shortly after starting + them, while the held poll(-1) is confirmed in-flight -- not by a fixed + wall-clock lower bound on total duration, which would be racing poll0's + own ~200ms closing-flag chunk boundary at a variable, flaky-to-bound + offset depending on scheduling.""" + num_workers = 5 + + for i in range(ITERATIONS): + producer = Producer(_PRODUCER_CONF) + poll_started = threading.Event() + + def hold_active_call(): + poll_started.set() + try: + producer.poll(-1) + except RuntimeError: + pass + + holder = threading.Thread(target=hold_active_call) + holder.start() + poll_started.wait() + # Give poll() a brief moment to increment the active calls counter + # before starting any close() call. + time.sleep(0.05) + + all_results = [] + start_barrier = threading.Barrier(num_workers) + + def call_close(): + start_barrier.wait() + all_results.append(producer.close()) + + threads = [threading.Thread(target=call_close) for _ in range(num_workers)] + for t in threads: + t.start() + + # Every close() caller must still be blocked shortly after starting: + # active_calls was > 0 (from the held poll(-1)) at the moment they + # started, so none of the threads should have finished yet. + time.sleep(0.15) assert all( - isinstance(r, bool) for r in all_results - ), f"iteration {i}: every close() call must return True or False, got: {all_results}" - assert any(all_results), f"iteration {i}: expected at least one close() call to return True, got: {all_results}" + t.is_alive() for t in threads + ), f"iteration {i}: a close() call returned before the held poll(-1) released active_calls" + + for t in threads: + t.join(timeout=10) + + holder.join(timeout=10) + assert not holder.is_alive(), f"iteration {i}: poll(-1) thread did not finish" + assert all(not t.is_alive() for t in threads), f"iteration {i}: a close() thread did not finish" + + assert len(all_results) == num_workers, f"iteration {i}: expected {num_workers} results, got {all_results}" + assert all(all_results), f"iteration {i}: every concurrent close() call must return True, got: {all_results}" ############################################################################### +# End of tests for races between Producer.close() and concurrent calls +# on the same Producer instance. +############################################################################### def test_close_completes_quickly_with_indefinite_poll_in_progress(): @@ -352,7 +401,10 @@ def run_flush(): @subprocess_isolated def test_close_propagates_signal_while_waiting_for_active_calls(): """close() must raise KeyboardInterrupt if a signal arrives while its - active_calls drain-wait loop is spinning""" + active_calls drain-wait loop is spinning. Afterwards, `closing` must + have been reset so the Handle isn't left permanently bricked -- a + retried close(), once the held active_calls slot is released, should + succeed instead of failing with ERR_MSG_PRODUCER_CLOSED or hanging.""" producer = Producer(_PRODUCER_CONF) poll_started = threading.Event() @@ -383,3 +435,17 @@ def send_sigint_soon(): finally: t.join(timeout=10) assert not t.is_alive(), "poll() thread did not finish after close() was interrupted" + + # The interrupted close() never destroyed rk (poll() was still holding + # an active_calls slot when the signal hit), so a retry now that + # poll() has returned must succeed instead of failing or hanging. + assert producer.close() is True, "close() retried after a signal-interrupted attempt must return True" + + +def test_close_is_idempotent(): + """Calling close() twice, with no concurrency at all, must return True + both times.""" + producer = Producer(_PRODUCER_CONF) + + assert producer.close() is True, "first close() must return True" + assert producer.close() is True, "second close() on an already-closed producer must also return True" From d6a9c08b44bfe32c68d0a405738ae06329d6fbb7 Mon Sep 17 00:00:00 2001 From: Ojasva Jain Date: Tue, 4 Aug 2026 18:35:59 +0530 Subject: [PATCH 08/10] Fix CI flakiness in test_close_races_close_losers_wait_for_slow_winner --- tests/concurrency/test_producer_close_race.py | 32 +++++++++---------- 1 file changed, 16 insertions(+), 16 deletions(-) diff --git a/tests/concurrency/test_producer_close_race.py b/tests/concurrency/test_producer_close_race.py index 1b2d53e36..596752c1e 100644 --- a/tests/concurrency/test_producer_close_race.py +++ b/tests/concurrency/test_producer_close_race.py @@ -267,15 +267,9 @@ def worker(): @subprocess_isolated def test_close_races_close_losers_wait_for_slow_winner(): - """Losing close() calls must actually block until the winner finishes, - not just happen to observe self->rk already NULL. A held poll(-1) call - keeps active_calls > 0, forcing the CAS winner (and therefore every - loser waiting behind it) to spin through its drain-wait loop. Proven by - checking every close() thread is still alive shortly after starting - them, while the held poll(-1) is confirmed in-flight -- not by a fixed - wall-clock lower bound on total duration, which would be racing poll0's - own ~200ms closing-flag chunk boundary at a variable, flaky-to-bound - offset depending on scheduling.""" + """Losing close() calls must actually block until the winner finishes. + A held poll(-1) call keeps active_calls > 0, forcing the CAS winner (and therefore every + loser waiting behind it) to spin through its drain-wait loop.""" num_workers = 5 for i in range(ITERATIONS): @@ -307,13 +301,19 @@ def call_close(): for t in threads: t.start() - # Every close() caller must still be blocked shortly after starting: - # active_calls was > 0 (from the held poll(-1)) at the moment they - # started, so none of the threads should have finished yet. - time.sleep(0.15) - assert all( - t.is_alive() for t in threads - ), f"iteration {i}: a close() call returned before the held poll(-1) released active_calls" + # Watch continuously: as long as the holder is still alive (still + # holding its active_calls slot), no close() thread should have + # finished yet. Stop watching, without failing, once the holder + # itself has finished -- at that point active_calls may have + # legitimately dropped to 0 and close() calls are free to proceed. + deadline = time.monotonic() + 10 + while time.monotonic() < deadline and holder.is_alive(): + finished_early = [t for t in threads if not t.is_alive()] + assert not finished_early, ( + f"iteration {i}: {len(finished_early)} close() call(s) returned while the held " + f"poll(-1) was still active -- active_calls should still have been > 0" + ) + time.sleep(0.01) for t in threads: t.join(timeout=10) From c513c9c7b23c3c032cc472c669529cd93a22217b Mon Sep 17 00:00:00 2001 From: Ojasva Jain Date: Tue, 4 Aug 2026 20:16:34 +0530 Subject: [PATCH 09/10] Move signal/slow-winner close() race tests to integration suite --- tests/concurrency/test_producer_close_race.py | 107 -------------- .../integration/producer/test_concurrency.py | 133 +++++++++++++++++- 2 files changed, 132 insertions(+), 108 deletions(-) diff --git a/tests/concurrency/test_producer_close_race.py b/tests/concurrency/test_producer_close_race.py index 596752c1e..dd2ec28e2 100644 --- a/tests/concurrency/test_producer_close_race.py +++ b/tests/concurrency/test_producer_close_race.py @@ -14,8 +14,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -import os -import signal import threading import time @@ -265,67 +263,6 @@ def worker(): assert all(all_results), f"iteration {i}: every concurrent close() call must return True, got: {all_results}" -@subprocess_isolated -def test_close_races_close_losers_wait_for_slow_winner(): - """Losing close() calls must actually block until the winner finishes. - A held poll(-1) call keeps active_calls > 0, forcing the CAS winner (and therefore every - loser waiting behind it) to spin through its drain-wait loop.""" - num_workers = 5 - - for i in range(ITERATIONS): - producer = Producer(_PRODUCER_CONF) - poll_started = threading.Event() - - def hold_active_call(): - poll_started.set() - try: - producer.poll(-1) - except RuntimeError: - pass - - holder = threading.Thread(target=hold_active_call) - holder.start() - poll_started.wait() - # Give poll() a brief moment to increment the active calls counter - # before starting any close() call. - time.sleep(0.05) - - all_results = [] - start_barrier = threading.Barrier(num_workers) - - def call_close(): - start_barrier.wait() - all_results.append(producer.close()) - - threads = [threading.Thread(target=call_close) for _ in range(num_workers)] - for t in threads: - t.start() - - # Watch continuously: as long as the holder is still alive (still - # holding its active_calls slot), no close() thread should have - # finished yet. Stop watching, without failing, once the holder - # itself has finished -- at that point active_calls may have - # legitimately dropped to 0 and close() calls are free to proceed. - deadline = time.monotonic() + 10 - while time.monotonic() < deadline and holder.is_alive(): - finished_early = [t for t in threads if not t.is_alive()] - assert not finished_early, ( - f"iteration {i}: {len(finished_early)} close() call(s) returned while the held " - f"poll(-1) was still active -- active_calls should still have been > 0" - ) - time.sleep(0.01) - - for t in threads: - t.join(timeout=10) - - holder.join(timeout=10) - assert not holder.is_alive(), f"iteration {i}: poll(-1) thread did not finish" - assert all(not t.is_alive() for t in threads), f"iteration {i}: a close() thread did not finish" - - assert len(all_results) == num_workers, f"iteration {i}: expected {num_workers} results, got {all_results}" - assert all(all_results), f"iteration {i}: every concurrent close() call must return True, got: {all_results}" - - ############################################################################### # End of tests for races between Producer.close() and concurrent calls # on the same Producer instance. @@ -398,50 +335,6 @@ def run_flush(): ) -@subprocess_isolated -def test_close_propagates_signal_while_waiting_for_active_calls(): - """close() must raise KeyboardInterrupt if a signal arrives while its - active_calls drain-wait loop is spinning. Afterwards, `closing` must - have been reset so the Handle isn't left permanently bricked -- a - retried close(), once the held active_calls slot is released, should - succeed instead of failing with ERR_MSG_PRODUCER_CLOSED or hanging.""" - producer = Producer(_PRODUCER_CONF) - poll_started = threading.Event() - - def hold_active_call(): - poll_started.set() - try: - producer.poll(-1) - except BaseException: # noqa: BLE001 - just draining the thread, not asserting here - pass - - t = threading.Thread(target=hold_active_call) - t.start() - poll_started.wait() - - def send_sigint_soon(): - time.sleep(0.05) - os.kill(os.getpid(), signal.SIGINT) - - interrupt_thread = threading.Thread(target=send_sigint_soon) - interrupt_thread.daemon = True - interrupt_thread.start() - - try: - producer.close() - assert False, "close() returned normally instead of raising KeyboardInterrupt" - except KeyboardInterrupt: - assert True # expected outcome: close() correctly propagated the signal - finally: - t.join(timeout=10) - assert not t.is_alive(), "poll() thread did not finish after close() was interrupted" - - # The interrupted close() never destroyed rk (poll() was still holding - # an active_calls slot when the signal hit), so a retry now that - # poll() has returned must succeed instead of failing or hanging. - assert producer.close() is True, "close() retried after a signal-interrupted attempt must return True" - - def test_close_is_idempotent(): """Calling close() twice, with no concurrency at all, must return True both times.""" diff --git a/tests/integration/producer/test_concurrency.py b/tests/integration/producer/test_concurrency.py index c0cb47062..35abd3bf1 100644 --- a/tests/integration/producer/test_concurrency.py +++ b/tests/integration/producer/test_concurrency.py @@ -15,6 +15,8 @@ # limitations under the License. import inspect +import os +import signal import threading import time from uuid import uuid1 @@ -35,7 +37,7 @@ def error_cb(err): return error_cb -class TestCloseRaceDelivery: +class TestCloseRace: def test_close_delivers_in_flight_messages(self, kafka_cluster): """ Messages produced just before/during a concurrent close() are genuinely delivered, @@ -165,6 +167,135 @@ def flush_worker(): ), f"flush() returned early with messages still outstanding: {flush_results}" assert len(producer) == 0 + def test_close_propagates_signal_while_waiting_for_active_calls(self, kafka_cluster): + """A signal during close()'s active_calls wait raises + KeyboardInterrupt and resets `closing`, but does not unblock an + already-held poll() -- it keeps running, unaware anything + happened. Only a second, uninterrupted close() call actually + finishes the job.""" + topic = kafka_cluster.create_topic_and_wait_propogation("test_close_signal_interruption") + producer = kafka_cluster.producer( + {'error_cb': prefixed_error_cb('test_close_propagates_signal_while_waiting_for_active_calls')} + ) + + callback_started = threading.Event() + + def slow_on_delivery(err, msg): + callback_started.set() + time.sleep(2) + + producer.produce(topic, value=b'msg', on_delivery=slow_on_delivery) + + poll_started = threading.Event() + + def hold_active_call(): + poll_started.set() + try: + producer.poll(10) + except BaseException: # noqa: BLE001 - just draining the thread, not asserting here + pass + + t = threading.Thread(target=hold_active_call, daemon=True) + t.start() + poll_started.wait() + callback_started.wait(timeout=10) + assert callback_started.is_set(), "delivery callback never started" + + def send_sigint_soon(): + time.sleep(0.5) # comfortably inside the callback's 2s sleep + os.kill(os.getpid(), signal.SIGINT) + + interrupt_thread = threading.Thread(target=send_sigint_soon, daemon=True) + interrupt_thread.start() + + try: + producer.close() + assert False, "close() returned normally instead of raising KeyboardInterrupt" + except KeyboardInterrupt: + pass + + # The interrupted attempt must NOT have unblocked the held poll(): + # `closing` was reset to 0 before poll() ever got a chance to see + # it, so poll() has no reason to exit early and must still be + # running its own call. + assert t.is_alive(), "poll() thread must still be running right after the interrupted close() attempt" + + # A second, uninterrupted close() call is required to actually + # finish the job -- it must succeed, and this time the held + # poll() must notice `closing` and exit, letting the holder thread + # finish. + assert producer.close() is True, "close() retried after a signal-interrupted attempt must return True" + + t.join(timeout=10) + assert not t.is_alive(), "poll() thread did not finish after the successful retry of close()" + + def test_close_races_close_losers_wait_for_slow_winner(self, kafka_cluster): + """Losing close() calls must actually block until the winner + finishes, not just happen to observe self->rk already NULL. A + slow delivery callback forces the CAS winner and every loser waiting behind it + to spin through the drain-wait loop.""" + topic = kafka_cluster.create_topic_and_wait_propogation("test_close_race_slow_winner") + producer = kafka_cluster.producer( + {'error_cb': prefixed_error_cb('test_close_races_close_losers_wait_for_slow_winner')} + ) + + callback_started = threading.Event() + + def slow_on_delivery(err, msg): + callback_started.set() + time.sleep(2) + + producer.produce(topic, value=b'msg', on_delivery=slow_on_delivery) + + poll_started = threading.Event() + + def hold_active_call(): + poll_started.set() + try: + producer.poll(10) + except RuntimeError: + pass + + holder = threading.Thread(target=hold_active_call, daemon=True) + holder.start() + poll_started.wait() + callback_started.wait(timeout=10) + assert callback_started.is_set(), "delivery callback never started" + + num_workers = 5 + all_results = [] + start_barrier = threading.Barrier(num_workers) + + def call_close(): + start_barrier.wait() + all_results.append(producer.close()) + + threads = [threading.Thread(target=call_close) for _ in range(num_workers)] + for t in threads: + t.start() + + # Watch continuously: as long as the holder is still alive (still + # holding its active_calls slot via the slow callback), no + # close() thread should have finished yet. + deadline = time.monotonic() + 10 + while time.monotonic() < deadline and holder.is_alive(): + finished_early = [t for t in threads if not t.is_alive()] + assert not finished_early, ( + f"{len(finished_early)} close() call(s) returned while the slow delivery callback " + f"was still holding active_calls" + ) + time.sleep(0.01) + + for t in threads: + t.join(timeout=10) + holder.join(timeout=10) + + print(f"{called_by()}: all_results={all_results}") + assert not holder.is_alive(), "poll() thread did not finish" + assert all(not t.is_alive() for t in threads), "a close() thread did not finish" + assert len(all_results) == num_workers, f"expected {num_workers} results, got: {all_results}" + assert all(all_results), f"every concurrent close() call must return True, got: {all_results}" + class TestTransactionalProducerConcurrency: def test_concurrent_produce_during_open_transaction(self, kafka_cluster): From 4c750a26a95f947716278578ce4e9c6ba1d8958f Mon Sep 17 00:00:00 2001 From: Ojasva Jain Date: Thu, 13 Aug 2026 16:13:36 +0530 Subject: [PATCH 10/10] Allow reentrant Producer calls from close()'s own delivery callback --- src/confluent_kafka/src/Producer.c | 9 ++++ src/confluent_kafka/src/confluent_kafka.c | 10 ++++- src/confluent_kafka/src/confluent_kafka.h | 24 +++++++++++ .../integration/producer/test_concurrency.py | 43 +++++++++++++++++++ 4 files changed, 84 insertions(+), 2 deletions(-) diff --git a/src/confluent_kafka/src/Producer.c b/src/confluent_kafka/src/Producer.c index 8bf201893..2ec79bcad 100644 --- a/src/confluent_kafka/src/Producer.c +++ b/src/confluent_kafka/src/Producer.c @@ -595,6 +595,12 @@ Producer_close(Handle *self, PyObject *args, PyObject *kwargs) { return NULL; } + /* Record which thread won, so Handle_enter_rk_use() can let a + * reentrant call through if (and only if) it's this same thread -- + * i.e. close()'s own delivery callback, fired synchronously from + * inside the flush() call below, calling back into the Producer. */ + atomic_ulong_set(&self->closing_thread, PyThread_get_thread_ident()); + /* Signal in-flight calls to stop, and wait for them to finish * using self->rk before destroying it -- see Handle_enter_rk_use(). * New calls will see `closing` and fail with ERR_MSG_PRODUCER_CLOSED. */ @@ -611,6 +617,7 @@ Producer_close(Handle *self, PyObject *args, PyObject *kwargs) { /* Abort the attempt: rk was never touched, so * reopen the gate for a future close() attempt. */ + atomic_ulong_set(&self->closing_thread, 0); atomic_int_set(&self->closing, 0); return NULL; } @@ -1088,6 +1095,8 @@ static PyObject *Producer_abort_transaction(Handle *self, PyObject *args) { if (!PyArg_ParseTuple(args, "|d", &tmout)) return NULL; + /* TODO NOGIL: closing rejects this mandatory abort even from an + * unrelated thread once close() has started. Revisit. */ if (!Handle_enter_rk_use(self)) return NULL; diff --git a/src/confluent_kafka/src/confluent_kafka.c b/src/confluent_kafka/src/confluent_kafka.c index 1b0c7fa68..5bdf06ca3 100644 --- a/src/confluent_kafka/src/confluent_kafka.c +++ b/src/confluent_kafka/src/confluent_kafka.c @@ -3326,14 +3326,20 @@ int CallState_end(Handle *h, CallState *cs) { * Handle is closed/closing (nothing to undo). */ int Handle_enter_rk_use(Handle *h) { - if (atomic_int_get(&h->closing) || !h->rk) { + unsigned long self_tid = PyThread_get_thread_ident(); + + if ((atomic_int_get(&h->closing) && + atomic_ulong_get(&h->closing_thread) != self_tid) || + !h->rk) { PyErr_SetString(PyExc_RuntimeError, ERR_MSG_PRODUCER_CLOSED); return 0; } atomic_int_inc(&h->active_calls); /* close() may have started between our check above and the * increment; re-check now that we're counted. */ - if (atomic_int_get(&h->closing) || !h->rk) { + if ((atomic_int_get(&h->closing) && + atomic_ulong_get(&h->closing_thread) != self_tid) || + !h->rk) { atomic_int_dec(&h->active_calls); PyErr_SetString(PyExc_RuntimeError, ERR_MSG_PRODUCER_CLOSED); return 0; diff --git a/src/confluent_kafka/src/confluent_kafka.h b/src/confluent_kafka/src/confluent_kafka.h index 78c794713..d1f257971 100644 --- a/src/confluent_kafka/src/confluent_kafka.h +++ b/src/confluent_kafka/src/confluent_kafka.h @@ -83,6 +83,29 @@ static inline int atomic_int_cas(atomic_int_t *p, int expected, int desired) { } #endif +/** + * @brief Same idea as atomic_int_t above, but sized to hold a + * PyThread_get_thread_ident() value (unsigned long, + * pointer-sized on most platforms) without truncation. + */ +#if defined(_MSC_VER) +typedef volatile LONG_PTR atomic_ulong_t; + +#define atomic_ulong_init(p, v) (*(p) = (v)) +#define atomic_ulong_get(p) \ + ((unsigned long)InterlockedCompareExchangePointer( \ + (PVOID volatile *)(p), 0, 0)) +#define atomic_ulong_set(p, v) \ + InterlockedExchangePointer((PVOID volatile *)(p), (PVOID)(v)) + +#else /* gcc / clang */ +typedef unsigned long atomic_ulong_t; + +#define atomic_ulong_init(p, v) __atomic_store_n((p), (v), __ATOMIC_SEQ_CST) +#define atomic_ulong_get(p) __atomic_load_n((p), __ATOMIC_SEQ_CST) +#define atomic_ulong_set(p, v) __atomic_store_n((p), (v), __ATOMIC_SEQ_CST) +#endif + /** * @brief confluent-kafka-python version, must match that of pyproject.toml. @@ -298,6 +321,7 @@ typedef struct { */ atomic_int_t active_calls; atomic_int_t closing; + atomic_ulong_t closing_thread; union { /** diff --git a/tests/integration/producer/test_concurrency.py b/tests/integration/producer/test_concurrency.py index 35abd3bf1..0bda8a39d 100644 --- a/tests/integration/producer/test_concurrency.py +++ b/tests/integration/producer/test_concurrency.py @@ -678,3 +678,46 @@ def worker(thread_id): ), f"expected all {expected_total} reentrantly-produced messages delivered, got {len(second_delivered)}" producer.close() + + def test_delivery_callback_reentrant_produce_during_close_does_not_truncate_flush(self, kafka_cluster): + """A delivery callback fired from close()'s own internal flush + calls produce() again. Handle_enter_rk_use() must let + this reentrant call through (it's the same thread that's running + close()'s flush) instead of rejecting it.""" + topic = kafka_cluster.create_topic_and_wait_propogation("test_reentrant_produce_during_close") + producer = kafka_cluster.producer( + {'error_cb': prefixed_error_cb('test_delivery_callback_reentrant_produce_during_close')} + ) + + first_delivered = [] + second_delivered = [] + errors = [] + + def on_second_delivery(err, msg): + if err: + errors.append(err) + else: + second_delivered.append(msg) + + def on_first_delivery(err, msg): + if err: + errors.append(err) + return + first_delivered.append(msg) + producer.produce(topic, value=b'reentrant-during-close', on_delivery=on_second_delivery) + + producer.produce(topic, value=b'original', on_delivery=on_first_delivery) + # Deliberately no poll()/flush() here -- close()'s own internal + # flush must be what dispatches the delivery callback. + assert len(producer) == 1, "the original message must still be queued when close() is called" + + result = producer.close() + + print( + f"{called_by()}: close()={result}, first_delivered={len(first_delivered)}, " + f"second_delivered={len(second_delivered)}, errors={errors}" + ) + assert result is True, "close() must complete cleanly despite the reentrant produce() from its callback" + assert not errors, f"unexpected delivery errors: {errors}" + assert len(first_delivered) == 1, "the original message must be delivered" + assert len(second_delivered) == 1, "the reentrantly-produced message must itself be delivered"