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 0bfc810b8..2ec79bcad 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); @@ -398,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 */ @@ -421,12 +433,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; @@ -465,14 +478,13 @@ 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; - 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); @@ -503,17 +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)) { - 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; @@ -525,12 +544,16 @@ Producer_flush(Handle *self, PyObject *args, PyObject *kwargs) { } if (!CallState_end(self, &cs)) - return NULL; + goto exit; /* result stays NULL */ if (err) /* Get the queue length on error (timeout) */ qlen = rd_kafka_outq_len(self->rk); - return cfl_PyInt_FromInt(qlen); + result = cfl_PyInt_FromInt(qlen); + +exit: + Handle_exit_rk_use(self); + return result; } @@ -542,6 +565,64 @@ 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 rest wait here for the winner to finish + * flushing and destroying it. + */ + if (!atomic_int_cas(&self->closing, 0, 1)) { + 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; + } + + /* 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. */ + /* 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 + Sleep(100); +#else + usleep(100000); +#endif + if (!CallState_end(self, &cs)) { + /* 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; + } + } + CallState_begin(self, &cs); /* Flush any pending messages (wait indefinitely to ensure delivery) */ @@ -817,10 +898,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,9 +928,12 @@ Producer_produce_batch(Handle *self, PyObject *args, PyObject *kwargs) { messages_list, rkt, partition, rkmessages, msgstates, message_cnt); cleanup: - /* 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) @@ -866,44 +948,49 @@ 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; - 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)) { - 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; } 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) { 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; @@ -917,108 +1004,124 @@ 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; - 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))) - return NULL; + goto exit; - if (!(cgmd = py_to_c_cgmd(metadata))) { - rd_kafka_topic_partition_list_destroy(c_offsets); - 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)) { - 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; } 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; - 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)) { - 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; } 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; - if (!self->rk) { - PyErr_SetString(PyExc_RuntimeError, ERR_MSG_PRODUCER_CLOSED); + /* 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; - } CallState_begin(self, &cs); error = rd_kafka_abort_transaction(self->rk, cfl_timeout_ms(tmout)); if (!CallState_end(self, &cs)) { - 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; } 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) { @@ -1034,10 +1137,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 +1149,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 +1506,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..5bdf06ca3 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,46 @@ 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). + */ +int Handle_enter_rk_use(Handle *h) { + 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) && + 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; + } + 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..d1f257971 100644 --- a/src/confluent_kafka/src/confluent_kafka.h +++ b/src/confluent_kafka/src/confluent_kafka.h @@ -35,6 +35,78 @@ #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 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. */ @@ -241,7 +313,15 @@ typedef struct { PyObject *logger; PyObject *oauth_cb; - int oauth_token_set; + atomic_int_t oauth_token_set; + + /* 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; + atomic_ulong_t closing_thread; union { /** @@ -350,6 +430,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/concurrency/__init__.py b/tests/concurrency/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/concurrency/_subprocess_isolation.py b/tests/concurrency/_subprocess_isolation.py new file mode 100644 index 000000000..5f1cad234 --- /dev/null +++ b/tests/concurrency/_subprocess_isolation.py @@ -0,0 +1,67 @@ +#!/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/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. +""" + +import functools +import os +import subprocess +import sys + +_SUBPROCESS_MARKER_ENV = "_CONCURRENCY_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/concurrency/test_producer_close_race.py b/tests/concurrency/test_producer_close_race.py new file mode 100644 index 000000000..dd2ec28e2 --- /dev/null +++ b/tests/concurrency/test_producer_close_race.py @@ -0,0 +1,344 @@ +#!/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 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. +############################################################################### + + +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 + + +@subprocess_isolated +def test_close_races_produce(): + """close() concurrent with produce() on another thread.""" + _race_close_against(_worker_produce) + + +@subprocess_isolated +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) + + +@subprocess_isolated +def test_close_races_poll(): + """close() concurrent with poll() on another thread.""" + _race_close_against(_worker_poll) + + +@subprocess_isolated +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) + + +@subprocess_isolated +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) + + +@subprocess_isolated +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 + + _race_close_against(worker, conf=_TXN_PRODUCER_CONF) + + +@subprocess_isolated +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 + + _race_close_against(worker, conf=_TXN_PRODUCER_CONF) + + +@subprocess_isolated +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 + + _race_close_against(worker, conf=_TXN_PRODUCER_CONF) + + +@subprocess_isolated +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 + + _race_close_against(worker, conf=_TXN_PRODUCER_CONF) + + +@subprocess_isolated +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 + + _race_close_against(worker, conf=_TXN_PRODUCER_CONF) + + +@subprocess_isolated +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) + + +@subprocess_isolated +def test_close_races_close(): + """Multiple threads calling close() on the same Producer at once. The + 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): + producer = Producer(_PRODUCER_CONF) + all_results = [] + start_barrier = threading.Barrier(num_workers + 1) + + def worker(): + start_barrier.wait() + all_results.append(producer.close()) + + 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(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(): + """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_finished_at = None + + def run_poll(): + nonlocal poll_finished_at + poll_started.set() + producer.poll(-1) + poll_finished_at = time.monotonic() + + t = threading.Thread(target=run_poll) + t.start() + poll_started.wait() + time.sleep(0.5) # Make sure poll() 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(), "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" + ) + + +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" + ) + + +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" diff --git a/tests/integration/producer/test_concurrency.py b/tests/integration/producer/test_concurrency.py new file mode 100644 index 000000000..0bda8a39d --- /dev/null +++ b/tests/integration/producer/test_concurrency.py @@ -0,0 +1,723 @@ +#!/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 os +import signal +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 TestCloseRace: + 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_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 -- + 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 + + 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): + """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}" + + +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() + + 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"