From 5b33de3f06cfd8b2ff7b4fa238603dcff3156ee0 Mon Sep 17 00:00:00 2001 From: Dana Powers Date: Wed, 29 Jul 2026 08:14:23 -0700 Subject: [PATCH 1/3] selector task.add_done_callback() --- kafka/net/backend/selector.py | 22 ++++++++++++++++------ 1 file changed, 16 insertions(+), 6 deletions(-) diff --git a/kafka/net/backend/selector.py b/kafka/net/backend/selector.py index 2ab20c37c..aa4364992 100644 --- a/kafka/net/backend/selector.py +++ b/kafka/net/backend/selector.py @@ -90,6 +90,7 @@ def __init__(self, coro): self._stack = (_initialize_coro(coro), None) self._res = None self._exc = None + self._callbacks = [] self.scheduled_at = None self.state = TaskState.CREATED @@ -129,10 +130,9 @@ def __call__(self, arg=None): except StopIteration as final: self._stack = self._stack[1] - if not self._stack: + if self.is_done: # we're done, back to event loop - self.state = TaskState.DONE - self._res = final.value + self._complete(res=final.value) raise else: ret = final.value @@ -140,9 +140,8 @@ def __call__(self, arg=None): except BaseException as e: self._stack = self._stack[1] - if not self._stack: - self.state = TaskState.DONE - self._exc = e + if self.is_done: + self._complete(exc=e) raise else: ret = None @@ -150,6 +149,14 @@ def __call__(self, arg=None): else: exc = None + def _complete(self, res=None, exc=None): + self._res = res + self._exc = exc + self.state = TaskState.DONE + for cb in self._callbacks: + cb(self) + self._callbacks = [] + def push_stack(self, coro): self._stack = (_initialize_coro(coro), self._stack) @@ -178,6 +185,9 @@ def close(self): self.state = TaskState.CANCELLED self._exc = Errors.Cancelled() + def add_done_callback(self, fn): + self._callbacks.append(fn) + @property def is_done(self): return self._stack is None From be4de42a0033ff8ed73cfb53d709571381999a54 Mon Sep 17 00:00:00 2001 From: Dana Powers Date: Wed, 29 Jul 2026 09:20:08 -0700 Subject: [PATCH 2/3] Consolidate selector create/close w/ _create_task + add_done_callback --- kafka/net/backend/selector.py | 32 +++++++++++++------------------- 1 file changed, 13 insertions(+), 19 deletions(-) diff --git a/kafka/net/backend/selector.py b/kafka/net/backend/selector.py index aa4364992..ce200e6cd 100644 --- a/kafka/net/backend/selector.py +++ b/kafka/net/backend/selector.py @@ -149,10 +149,11 @@ def __call__(self, arg=None): else: exc = None - def _complete(self, res=None, exc=None): + def _complete(self, res=None, exc=None, state=TaskState.DONE): + assert self._stack is None, "Cannot complete Task with non-empty stack!" self._res = res self._exc = exc - self.state = TaskState.DONE + self.state = state for cb in self._callbacks: cb(self) self._callbacks = [] @@ -182,8 +183,7 @@ def close(self): except Exception: log.exception('Error closing coroutine for cancelled task') self._stack = None - self.state = TaskState.CANCELLED - self._exc = Errors.Cancelled() + self._complete(exc=Errors.Cancelled(), state=TaskState.CANCELLED) def add_done_callback(self, fn): self._callbacks.append(fn) @@ -456,16 +456,15 @@ def call_at(self, when, task): if self._closed: raise RuntimeError('NetworkSelector closed!') if not isinstance(task, Task): - task = Task(task) + task = self._create_task(task) task.scheduled_at = when task.state = TaskState.SCHEDULED heapq.heappush(self._scheduled, (when, task)) - self._pending_tasks.add(task) return task def call_later(self, delay, task): if not isinstance(task, Task): - task = Task(task) + task = self._create_task(task) self.call_at(time.monotonic() + delay, task) return task @@ -473,11 +472,11 @@ def _add_ready_task(self, task): self._ready.append(task) task.state = TaskState.READY - def _task_done(self, task): - if not task.is_done: - raise RuntimeError('Task is not done yet!') - self._pending_tasks.discard(task) - task.state = TaskState.DONE + def _create_task(self, task): + task = Task(task) + task.add_done_callback(lambda t: self._pending_tasks.discard(t)) + self._pending_tasks.add(task) + return task def call_soon(self, task): """Schedule a coroutine/callable on the loop; return its Task handle. @@ -504,9 +503,8 @@ def call_soon(self, task): elif self._closed: raise RuntimeError('NetworkSelector closed!') if not isinstance(task, Task): - task = Task(task) + task = self._create_task(task) self._add_ready_task(task) - self._pending_tasks.add(task) if threadsafe: self.wakeup() return task @@ -537,7 +535,6 @@ def cancel(self, task): # close() below drives the io_guard finalizer, which unregisters # the fileobj and cancels any paired timeout timer. pass - self._pending_tasks.discard(task) task.close() def reschedule(self, when, task): @@ -821,17 +818,14 @@ def _poll_once(self, timeout=None): event = self._current() except StopIteration: - self._task_done(self._current) + pass except BaseException: log.exception('Unhandled exception in task %s:', self._current) - # Same as StopIteration -- task is done either way. - self._task_done(self._current) else: if self._current.state is TaskState.CANCELLED: # ignores any returned KernelEvent/Future - self._pending_tasks.discard(self._current) self._current.close() elif isinstance(event, KernelEvent): log_trace('kernel event %s', event.method) From 0936501ce69ffda48fa7b194c03cce865b84d738 Mon Sep 17 00:00:00 2001 From: Dana Powers Date: Wed, 29 Jul 2026 11:17:14 -0700 Subject: [PATCH 3/3] Drop backend.cancel() in favor of task.cancel() --- kafka/net/backend/abstract.py | 6 +- kafka/net/backend/asyncio_backend.py | 4 -- kafka/net/backend/selector.py | 77 +++++++++++------------- kafka/net/backend/transport.py | 2 +- kafka/net/connection.py | 8 ++- kafka/net/wakeup_notifier.py | 2 +- test/net/backend/test_abstract.py | 2 +- test/net/backend/test_asyncio_backend.py | 2 +- test/net/backend/test_selector.py | 25 +++----- test/net/backend/test_transport.py | 2 +- 10 files changed, 55 insertions(+), 75 deletions(-) diff --git a/kafka/net/backend/abstract.py b/kafka/net/backend/abstract.py index 16432d08d..7a621d7aa 100644 --- a/kafka/net/backend/abstract.py +++ b/kafka/net/backend/abstract.py @@ -216,10 +216,6 @@ def call_at(self, when: float, task: Any) -> Any: def call_later(self, delay: float, task: Any) -> Any: """Schedule ``task`` to run after ``delay`` seconds.""" - @abc.abstractmethod - def cancel(self, task: Any) -> None: - """Cancel a scheduled task/timer previously returned by call_*.""" - # --- timing (core coroutines await this) ------------------------------ @abc.abstractmethod def sleep(self, delay: float) -> Any: @@ -313,7 +309,7 @@ def _on_timeout(): raise finally: if timer is not None: - self.cancel(timer) + timer.cancel() def wait_for(self, future: Any, timeout_ms: Optional[float], raise_error: bool=True) -> Any: """Block the calling thread until ``future`` resolves, with a timeout in ms. diff --git a/kafka/net/backend/asyncio_backend.py b/kafka/net/backend/asyncio_backend.py index 369674474..b7e3096e4 100644 --- a/kafka/net/backend/asyncio_backend.py +++ b/kafka/net/backend/asyncio_backend.py @@ -267,10 +267,6 @@ def call_later(self, delay, task): lambda: box._arm(self._loop.call_later(delay, cb))) return box - def cancel(self, task): - if task is not None: - task.cancel() - def sleep(self, delay): return asyncio.sleep(delay) diff --git a/kafka/net/backend/selector.py b/kafka/net/backend/selector.py index ce200e6cd..b8f5c50e7 100644 --- a/kafka/net/backend/selector.py +++ b/kafka/net/backend/selector.py @@ -76,9 +76,9 @@ def __await__(self): class TaskState(enum.Enum): CREATED = 'created' SCHEDULED = 'scheduled' # in _scheduled heap - UNSCHEDULED = 'unscheduled' # maybe lost READY = 'ready' # in _ready deque - RUNNING = 'running' # is _current + RUNNING = 'running' # in Task.__call__ + STOPPED = 'stopped' # suspended in Task.__call__ WAIT_IO = 'wait_io' # parked on I/O WAIT_FUTURE = 'wait_future' # waiting on Future to resolve DONE = 'done' # completed (exception is None or not) @@ -109,6 +109,7 @@ def __call__(self, arg=None): else: ret = None exc = None + self.state = TaskState.RUNNING while True: coro = self._stack[0] if callable(coro) and not inspect.isgenerator(coro) and not inspect.iscoroutine(coro): @@ -120,14 +121,6 @@ def __call__(self, arg=None): else: ret = coro.send(ret) - if isinstance(ret, (KernelEvent, Future)): - # handle in event loop - return ret - - elif inspect.isgenerator(ret) or inspect.iscoroutine(ret) or inspect.isfunction(ret): - self.push_stack(ret) - ret = None - except StopIteration as final: self._stack = self._stack[1] if self.is_done: @@ -149,6 +142,23 @@ def __call__(self, arg=None): else: exc = None + # Complete any self-cancel + if self.state is TaskState.CANCELLED: + self.cancel() + raise self._exc + + elif self.is_done: + raise RuntimeError('Unexpected inline task completion') + + if isinstance(ret, (KernelEvent, Future)): + # handle in event loop + self.state = TaskState.STOPPED + return ret + + elif inspect.isgenerator(ret) or inspect.iscoroutine(ret) or inspect.isfunction(ret): + self.push_stack(ret) + ret = None + def _complete(self, res=None, exc=None, state=TaskState.DONE): assert self._stack is None, "Cannot complete Task with non-empty stack!" self._res = res @@ -170,10 +180,13 @@ def inject_exc(self, exc): raise RuntimeError('Task exception is already set!') self._exc = exc - def close(self): + def cancel(self): if self.is_done: return - assert self.state is not TaskState.RUNNING + if self.state is TaskState.RUNNING: + # Defer cancel until task leaves RUNNING state + self.state = TaskState.CANCELLED + return stack = self._stack while stack: coro, stack = stack @@ -460,6 +473,7 @@ def call_at(self, when, task): task.scheduled_at = when task.state = TaskState.SCHEDULED heapq.heappush(self._scheduled, (when, task)) + task.add_done_callback(self._unschedule) return task def call_later(self, delay, task): @@ -510,8 +524,8 @@ def call_soon(self, task): return task def _unschedule(self, task): - assert task.state is TaskState.SCHEDULED - assert task.scheduled_at is not None + if task.scheduled_at is None: + return try: self._scheduled.remove((task.scheduled_at, task)) except ValueError: @@ -520,22 +534,6 @@ def _unschedule(self, task): # re-heapify to ensure heap structure is valid heapq.heapify(self._scheduled) task.scheduled_at = None - task.state = TaskState.UNSCHEDULED - - def cancel(self, task): - if task.state in (TaskState.DONE, TaskState.CANCELLED): - return - elif task.state is TaskState.RUNNING: - assert task is self._current - self._current.state = TaskState.CANCELLED - return - elif task.state is TaskState.SCHEDULED: - self._unschedule(task) - elif task.state is TaskState.WAIT_IO: - # close() below drives the io_guard finalizer, which unregisters - # the fileobj and cancels any paired timeout timer. - pass - task.close() def reschedule(self, when, task): if task.state is TaskState.SCHEDULED: @@ -677,7 +675,7 @@ def io_guard(): yield finally: if timer is not None and not timer.is_done: - self.cancel(timer) + timer.cancel() self.unregister_event(fileobj, event) guard = io_guard() @@ -702,7 +700,8 @@ def _schedule_tasks(self): while self._scheduled and self._scheduled[0][0] <= time.monotonic(): _, task = heapq.heappop(self._scheduled) task.scheduled_at = None - self._add_ready_task(task) + if not task.is_done: + self._add_ready_task(task) def _next_scheduled_timeout(self, now): try: @@ -807,9 +806,8 @@ def _poll_once(self, timeout=None): for i in range(n): self._current = self._ready.popleft() # Silently skip tasks that are done or cancelled - if self._current.state in (TaskState.DONE, TaskState.CANCELLED): + if self._current.is_done: continue - self._current.state = TaskState.RUNNING step_start = time.monotonic() if threshold else None try: log_trace('Calling task %s', self._current) @@ -824,10 +822,7 @@ def _poll_once(self, timeout=None): log.exception('Unhandled exception in task %s:', self._current) else: - if self._current.state is TaskState.CANCELLED: - # ignores any returned KernelEvent/Future - self._current.close() - elif isinstance(event, KernelEvent): + if isinstance(event, KernelEvent): log_trace('kernel event %s', event.method) try: getattr(self, event.method)(*event.args) @@ -846,8 +841,8 @@ def _poll_once(self, timeout=None): # No Task should leave io_loop in RUNNING state. if self._current is not None and self._current.state is TaskState.RUNNING: log.warning('Task %s left RUNNING after step; demoting to ' - 'UNSCHEDULED', self._current) - self._current.state = TaskState.UNSCHEDULED + 'STOPPED', self._current) + self._current.state = TaskState.STOPPED if threshold: elapsed = time.monotonic() - step_start @@ -898,7 +893,7 @@ def close(self): self.stop() self.drain() for task in list(self._pending_tasks): - self.cancel(task) + task.cancel() for s in (self._wakeup_r, self._wakeup_w): try: self._selector.unregister(s) diff --git a/kafka/net/backend/transport.py b/kafka/net/backend/transport.py index b876ca8a5..0dec62d3f 100644 --- a/kafka/net/backend/transport.py +++ b/kafka/net/backend/transport.py @@ -214,7 +214,7 @@ def _close(self, error=None): sock.close() for task in (self._read_task, self._write_task): if task is not None: - self._net.cancel(task) + task.cancel() self._read_task = self._write_task = None proto = self._protocol self._protocol = None diff --git a/kafka/net/connection.py b/kafka/net/connection.py index 414145b34..20c6cd6a6 100644 --- a/kafka/net/connection.py +++ b/kafka/net/connection.py @@ -209,7 +209,8 @@ def data_received(self, data): if req_correlation_id != resp_correlation_id: return self.close(Errors.KafkaConnectionError('Received unrecognized correlation id')) - self.net.cancel(timeout_task) + if timeout_task is not None: + timeout_task.cancel() latency_ms = (time.monotonic() - sent_time) * 1000 if self._sensors: self._sensors.request_time.record(latency_ms) @@ -251,7 +252,8 @@ def fail_in_flight_requests(self, error): future.failure(error) while self.in_flight_requests: _, future, _, _, timeout_task = self.in_flight_requests.popleft() - self.net.cancel(timeout_task) + if timeout_task is not None: + timeout_task.cancel() future.failure(error) def connection_made(self, transport): @@ -548,7 +550,7 @@ def cancel(self): """Cancel any pending re-auth and fail the drain awaiter if present. Called from KafkaConnection.connection_lost.""" if self._task is not None: - self._conn.net.cancel(self._task) + self._task.cancel() self._task = None if self._drain_future is not None and not self._drain_future.is_done: self._drain_future.failure(Errors.KafkaConnectionError()) diff --git a/kafka/net/wakeup_notifier.py b/kafka/net/wakeup_notifier.py index bf70534b7..228161f04 100644 --- a/kafka/net/wakeup_notifier.py +++ b/kafka/net/wakeup_notifier.py @@ -63,7 +63,7 @@ async def __call__(self, timeout_secs=None): finally: self._fut = None if timer is not None: - self._net.cancel(timer) + timer.cancel() def notify(self): # Coalesce: if a _wakeup is already scheduled and not yet consumed, diff --git a/test/net/backend/test_abstract.py b/test/net/backend/test_abstract.py index 990eaee71..b86c90ae5 100644 --- a/test/net/backend/test_abstract.py +++ b/test/net/backend/test_abstract.py @@ -23,7 +23,7 @@ CONTRACT_METHODS = ( 'start', 'stop', 'close', 'on_io_thread', 'call_soon', 'call_soon_with_future', - 'call_at', 'call_later', 'cancel', + 'call_at', 'call_later', 'sleep', 'create_connection', 'run', 'create_future', 'wakeup', ) diff --git a/test/net/backend/test_asyncio_backend.py b/test/net/backend/test_asyncio_backend.py index d3cb24872..bfc511ea5 100644 --- a/test/net/backend/test_asyncio_backend.py +++ b/test/net/backend/test_asyncio_backend.py @@ -133,7 +133,7 @@ def test_call_later_cancel_prevents_fire(self, started_backend): async def schedule_and_cancel(): handle = started_backend.call_later(0.5, fired.set) - started_backend.cancel(handle) + handle.cancel() started_backend.run(schedule_and_cancel) assert not fired.wait(timeout=0.3) diff --git a/test/net/backend/test_selector.py b/test/net/backend/test_selector.py index 3bcf91020..ba5e551d0 100644 --- a/test/net/backend/test_selector.py +++ b/test/net/backend/test_selector.py @@ -392,7 +392,7 @@ async def reader(): net.drain() # park on I/O assert task.state is TaskState.WAIT_IO assert net._selector.get_key(rsock) is not None - net.cancel(task) + task.cancel() assert task.state is TaskState.CANCELLED with pytest.raises(KeyError): net._selector.get_key(rsock) @@ -414,7 +414,7 @@ def test_cancel_running_task(self): holder = {} async def self_cancel(): - net.cancel(holder['task']) # cancel self while RUNNING + holder['task'].cancel() # cancel self while RUNNING await net.wait_read(rsock) # short-circuited; must not register raise AssertionError('cancelled task must not run past the await') @@ -452,7 +452,7 @@ async def reader(): # simulate connection teardown: drop the registration, then close net.unregister_event(rsock, selectors.EVENT_READ) rsock.close() # fileno() -> -1 - net.cancel(task) # io_guard unregisters a dead fd + task.cancel() # io_guard unregisters a dead fd assert task.state is TaskState.CANCELLED finally: wsock.close() @@ -488,7 +488,7 @@ async def reader(): def test_running_task_demoted_when_left_running(self): # Defensive backstop: a kernel-event handler that returns without # parking the task leaves it stranded in RUNNING. _poll_once's finally - # must demote it to UNSCHEDULED (with a warning) -- not crash the loop + # must demote it to STOPPED (with a warning) -- not crash the loop # -- so a later cancel()/close() can reclaim it instead of tripping # cancel()'s `task is self._current` assert. Simulate a buggy handler # with a no-op method that parks nothing. @@ -500,10 +500,10 @@ def coro(): task = net.call_soon(coro) net.drain() # demotes + warns, no crash - assert task.state is TaskState.UNSCHEDULED, task.state + assert task.state is TaskState.STOPPED, task.state assert net._current is None # cancel() must reclaim the demoted task without asserting. - net.cancel(task) + task.cancel() assert task.state is TaskState.CANCELLED assert task not in net._pending_tasks @@ -524,7 +524,7 @@ async def reader(): net.drain() assert len(net._scheduled) == 1 timer = net._scheduled[0][1] - net.cancel(task) + task.cancel() with pytest.raises(KeyError): net._selector.get_key(rsock) assert timer.is_done # paired timer cancelled/closed @@ -543,15 +543,6 @@ def task(): assert len(net._scheduled) == 0 assert t.scheduled_at is None - def test_unschedule_unscheduled_raises(self): - net = NetworkSelector() - def task(): - yield - assert len(net._scheduled) == 0 - with pytest.raises(AssertionError): - net._unschedule(Task(task)) - assert len(net._scheduled) == 0 - def test_reschedule(self): net = NetworkSelector() def task(): @@ -735,7 +726,7 @@ def test_cancel_closes_ready_task(self): assert timer in net._ready assert timer.scheduled_at is None - net.cancel(timer) + timer.cancel() assert timer in net._ready # timer still in ready queue assert timer.is_done, \ diff --git a/test/net/backend/test_transport.py b/test/net/backend/test_transport.py index 26ebc36fc..f581a917c 100644 --- a/test/net/backend/test_transport.py +++ b/test/net/backend/test_transport.py @@ -316,7 +316,7 @@ class TestTransportWaiterCleanup: read/write coroutine tasks parked in the event loop. These tests fail until ``KafkaTCPTransport._close`` cancels its read/write - waiter tasks (``net.cancel(task)``); the selector's existing WAIT_IO branch + waiter tasks (``task.cancel()``); the selector's existing WAIT_IO branch in ``cancel()`` then drives the io_guard finalizer and discards the task. """