diff --git a/DESIGN_ISSUES.md b/DESIGN_ISSUES.md index 9938514d..07e1cc15 100644 --- a/DESIGN_ISSUES.md +++ b/DESIGN_ISSUES.md @@ -150,6 +150,28 @@ change its data schema on the second new-data poll and crash on the third. --- +### PS4 — Concurrent sync access during async run silently loses rows +**Status:** resolved +**Severity:** critical +**Issue:** ITL-617 + +`_get_latest_stream()`, called by `iter_data()` or `as_table()` while +`async_iter_data()` is running, could advance `_cursor` and fold new rows into +`_accumulated_stream` without the async loop emitting them — permanent silent data +loss. `ITL-615` (PR #255) partially addressed this for `output_schema()` and `keys()` +by caching the stream reference. `ITL-617` fixes the root. + +**Fix:** Replaced single `_accumulated_stream: ArrowTableStream | None` with +append-only `_batches: list[ArrowTableStream]` and `_state_lock: threading.Lock`. +All callers (sync and async) use an optimistic lock protocol: snapshot cursor (brief +lock) → perform I/O freely with no lock held → commit only if cursor is unchanged +(brief lock). The async loop tracks its yield position with a per-iterator +``local_batch_idx`` local variable and drains ``_batches`` at the top of each +iteration, ensuring rows committed by a concurrent sync caller are yielded even when +the async loop loses the commit race. The lock is never held across ``await``. + +--- + ## `src/orcapod/core/nodes/function_node.py` ### FN1 — `FunctionNodeBase.as_table()` returned empty schema when no data existed diff --git a/src/orcapod/core/sources/polling_source.py b/src/orcapod/core/sources/polling_source.py index d729e051..d022d9f0 100644 --- a/src/orcapod/core/sources/polling_source.py +++ b/src/orcapod/core/sources/polling_source.py @@ -12,6 +12,7 @@ import dataclasses import functools import logging +import threading from collections.abc import Collection from math import floor from typing import TYPE_CHECKING, Any, Generic, TypeVar @@ -255,7 +256,8 @@ def __init__( self._tag_columns: tuple[str, ...] = tuple(_normalize_column_list(tag_columns)) self._polling_config = polling_config self._cursor: Cursor[T] | None = None - self._accumulated_stream: ArrowTableStream | None = None + self._batches: list[ArrowTableStream] = [] + self._state_lock: threading.Lock = threading.Lock() # Derive source_id from impl identity if not explicitly provided if self._source_id is None: self._source_id = str(self._impl.identity()) @@ -330,7 +332,7 @@ def identity_structure(self) -> Any: return (self._impl.identity(), self._tag_columns) # ------------------------------------------------------------------------- - # Sync stream delegation — all route through _get_latest_stream() + # Sync stream delegation # ------------------------------------------------------------------------- def output_schema( @@ -352,17 +354,18 @@ def output_schema( declared schemas (``system_tags``); flags requiring actual persisted data (``meta``, ``source``, ``context``, ``content_hash``) fall through to the next level. - 2. **Cached-stream path** — if ``accumulated_stream`` is already - populated (at least one batch fetched via sync or async), delegates - to it without triggering a new poll+fetch cycle. - 3. **Fallback** — calls ``_get_latest_stream()``, which runs a - synchronous poll+fetch via ``_run_sync``. Only reached when the impl - declared no schema and no batch has been fetched yet. + 2. **Cached-stream path** — if ``_batches`` is already non-empty (at + least one batch fetched via sync or async), delegates to the first + batch without triggering a new poll+fetch cycle. + 3. **Fallback** — calls ``_sync_poll_and_commit()`` then + ``_get_combined_stream()``, running a synchronous poll+fetch via + ``_run_sync``. Only reached when the impl declared no schema and no + batch has been fetched yet. """ if self._tag_schema is not None and self._data_schema is not None: columns_config = ColumnConfig.handle_config(columns, all_info=all_info) # meta / source / context / content_hash require actual stream data — - # fall through to accumulated_stream or _get_latest_stream for those. + # fall through to _batches[0] or _sync_poll_and_commit for those. if not ( columns_config.meta or columns_config.source @@ -380,9 +383,10 @@ def output_schema( {**dict(tag_schema), src_col: str, rec_col: bytes} ) return tag_schema, self._data_schema - if self._accumulated_stream is not None: - return self._accumulated_stream.output_schema(columns=columns, all_info=all_info) - return self._get_latest_stream().output_schema(columns=columns, all_info=all_info) + if self._batches: + return self._batches[0].output_schema(columns=columns, all_info=all_info) + self._sync_poll_and_commit() + return self._get_combined_stream().output_schema(columns=columns, all_info=all_info) def keys( self, @@ -393,9 +397,9 @@ def keys( """Return tag and data column keys. Uses the same three-level resolution as ``output_schema``: impl-declared - schemas → cached stream → ``_get_latest_stream()`` fallback. When the - impl declared a schema, system-tag column names are derived from - ``_declared_schema_hash`` without triggering a fetch. + schemas → cached ``_batches[0]`` → ``_sync_poll_and_commit()`` fallback. + When the impl declared a schema, system-tag column names are derived + from ``_declared_schema_hash`` without triggering a fetch. """ if self._tag_schema is not None and self._data_schema is not None: columns_config = ColumnConfig.handle_config(columns, all_info=all_info) @@ -414,13 +418,15 @@ def keys( src_col, rec_col = system_tag_column_names(schema_hash) tag_keys = tag_keys + (src_col, rec_col) return tag_keys, tuple(self._data_schema.keys()) - if self._accumulated_stream is not None: - return self._accumulated_stream.keys(columns=columns, all_info=all_info) - return self._get_latest_stream().keys(columns=columns, all_info=all_info) + if self._batches: + return self._batches[0].keys(columns=columns, all_info=all_info) + self._sync_poll_and_commit() + return self._get_combined_stream().keys(columns=columns, all_info=all_info) def iter_data(self): """Iterate over (tag, data) pairs from the current snapshot.""" - return self._get_latest_stream().iter_data() + self._sync_poll_and_commit() + return self._get_combined_stream().iter_data() def as_table( self, @@ -429,7 +435,8 @@ def as_table( all_info: bool = False, ) -> pa.Table: """Return the accumulated rows as a PyArrow table.""" - return self._get_latest_stream().as_table(columns=columns, all_info=all_info) + self._sync_poll_and_commit() + return self._get_combined_stream().as_table(columns=columns, all_info=all_info) # ------------------------------------------------------------------------- # Serialization @@ -510,40 +517,86 @@ def from_config(cls, config: dict[str, Any], db_registry: Any = None) -> Polling # Internal sync helpers # ------------------------------------------------------------------------- - def _get_latest_stream(self) -> ArrowTableStream: - """Return the current accumulated stream, fetching/polling as needed.""" - if self._accumulated_stream is None: - # First access — no cache yet; fetch immediately + def _sync_poll_and_commit(self) -> None: + """Poll for new data and commit a new batch under the optimistic lock. + + Implements the check-snapshot → I/O → check-and-commit protocol: + reads the cursor snapshot under a brief lock, performs poll and fetch + with no lock held, then commits only if the cursor has not changed. + If another caller advanced the cursor in between (winning the commit + race), this call discards its fetched data — the batch is already in + ``_batches`` and will be visible on the next read. + + Safe to call concurrently with ``async_iter_data``. + """ + with self._state_lock: + cursor_snapshot = self._cursor + + if cursor_snapshot is None: + # First access — no poll needed, fetch unconditionally. logger.debug("PollingSource %r: first sync access — fetching", self._source_id) new_cursor, data = _run_sync(self._impl.fetch, cursor=None) new_stream = self._try_build_stream(data) if new_stream is not None: if self._tag_schema is not None or self._data_schema is not None: self._validate_against_declared_schemas(new_stream) - self._accumulated_stream = new_stream - self._update_last_modified_from_cursor(new_cursor) - self._cursor = new_cursor + committed = False + with self._state_lock: + if self._cursor is None: + if new_stream is not None: + self._batches.append(new_stream) + self._cursor = new_cursor + committed = True + if committed: + self._update_last_modified_from_cursor(new_cursor) else: - # Have cache — poll for updates - has_new = _run_sync(self._impl.poll, cursor=self._cursor) + has_new = _run_sync(self._impl.poll, cursor=cursor_snapshot) if has_new: - logger.debug("PollingSource %r: sync poll found new data — fetching", self._source_id) - new_cursor, data = _run_sync(self._impl.fetch, cursor=self._cursor) + logger.debug( + "PollingSource %r: sync poll found new data — fetching", self._source_id + ) + new_cursor, data = _run_sync(self._impl.fetch, cursor=cursor_snapshot) new_stream = self._try_build_stream(data) if new_stream is not None: if self._tag_schema is not None or self._data_schema is not None: self._validate_against_declared_schemas(new_stream) - self._accumulated_stream = self._combine(self._accumulated_stream, new_stream) - self._update_last_modified_from_cursor(new_cursor) - self._cursor = new_cursor + if self._batches: + self._validate_combining_schemas(self._batches[0], new_stream) + committed = False + with self._state_lock: + if self._cursor == cursor_snapshot: + if new_stream is not None: + self._batches.append(new_stream) + self._cursor = new_cursor + committed = True + if committed: + self._update_last_modified_from_cursor(new_cursor) else: - logger.debug("PollingSource %r: sync poll — cache still valid", self._source_id) + logger.debug( + "PollingSource %r: sync poll — cache still valid", self._source_id + ) - if self._accumulated_stream is None: + def _get_combined_stream(self) -> ArrowTableStream: + """Return all committed batches concatenated as a single stream. + + Takes a snapshot of ``_batches`` (no lock needed — list is append-only) + then combines using the existing ``_combine`` helper. + + Returns: + A single ``ArrowTableStream`` containing all rows from all batches. + + Raises: + ValueError: If no data has been fetched yet (``_batches`` is empty). + """ + batches = list(self._batches) # snapshot — safe, list is append-only + if not batches: raise ValueError( "PollingSource: no data available yet — first fetch returned empty data." ) - return self._accumulated_stream + result = batches[0] + for batch in batches[1:]: + result = self._combine(result, batch) + return result def _try_build_stream(self, data: FrameInitTypes) -> ArrowTableStream | None: """Build an ``ArrowTableStream`` from raw data, returning ``None`` for empty data. @@ -683,26 +736,22 @@ def _update_last_modified_from_cursor(self, cursor: Cursor[T]) -> None: async def async_iter_data(self): """Async generator that continuously emits (tag, data) pairs. - Pre-seeds from the cached stream (if any) before entering the polling - loop. The loop runs until: the configured duration elapses, the maximum - consecutive error or overrun threshold is exceeded, or the task is - cancelled. ``CursorInvalidatedError`` is re-raised to the caller after - logging — the caller will see it as an exception rather than a clean - end-of-stream. ``SchemaInconsistencyError`` is also propagated - immediately. + Uses a per-iterator ``local_batch_idx`` to track the next position in + the append-only ``_batches`` list. A drain step at the top of each + loop iteration yields any newly committed batches (including those + committed by concurrent sync callers) before sleeping. + + The optimistic lock protocol — snapshot cursor (brief lock) → poll and + fetch with no lock held → commit only if cursor unchanged (brief lock) + — prevents TOCTOU races without ever holding ``_state_lock`` across an + ``await``. ``impl.close()`` is always awaited before returning or raising. """ - # Pre-seed from cache - if self._accumulated_stream is not None: - pre_seed_count = 0 - for pre_seed_count, item in enumerate( - self._accumulated_stream.iter_data(), start=1 - ): - yield item - logger.debug( - "PollingSource %r: pre-seeded %d row(s)", self._source_id, pre_seed_count - ) + # local_batch_idx tracks the next _batches index to yield. + # Starting at 0 means the first drain covers any pre-existing batches + # (the pre-seed case from the old implementation). + local_batch_idx = 0 cfg = self._polling_config loop = asyncio.get_running_loop() @@ -720,44 +769,70 @@ async def async_iter_data(self): try: while True: - # 1. Sleep until next scheduled tick + # ── 1. Drain: yield any batches not yet emitted ────────────── + # Covers pre-existing rows on first iteration AND rows committed + # by concurrent sync callers while this iterator was polling. + while local_batch_idx < len(self._batches): + for item in self._batches[local_batch_idx].iter_data(): + yield item + local_batch_idx += 1 + + # ── 2. Sleep to next scheduled tick ────────────────────────── now = loop.time() if next_tick > now: await asyncio.sleep(next_tick - now) - # 2. Poll + fetch + # ── 3. Poll + fetch + commit (optimistic lock) ─────────────── try: - has_new = await self._impl.poll(cursor=self._cursor) + # Brief lock: snapshot cursor only. + with self._state_lock: + cursor_snapshot = self._cursor + + # Poll — no lock held across await. + has_new = await self._impl.poll(cursor=cursor_snapshot) if has_new: logger.debug( "PollingSource %r: new data detected, fetching", self._source_id, ) - new_cursor, data = await self._impl.fetch(cursor=self._cursor) + # Fetch — no lock held across await. + new_cursor, data = await self._impl.fetch(cursor=cursor_snapshot) new_stream = self._try_build_stream(data) - self._cursor = new_cursor - self._update_last_modified_from_cursor(new_cursor) + if new_stream is not None: if self._tag_schema is not None or self._data_schema is not None: self._validate_against_declared_schemas(new_stream) - if self._accumulated_stream is None: - self._accumulated_stream = new_stream - else: - self._accumulated_stream = self._combine( - self._accumulated_stream, new_stream + if self._batches: + self._validate_combining_schemas( + self._batches[0], new_stream ) - # Emit new rows from just this fetch - emitted = 0 - for item in new_stream.iter_data(): - yield item - emitted += 1 - logger.debug( - "PollingSource %r: emitted %d row(s)", - self._source_id, - emitted, - ) + # Brief lock: check cursor then commit atomically. + committed = False + with self._state_lock: + if self._cursor == cursor_snapshot: + if new_stream is not None: + self._batches.append(new_stream) + self._cursor = new_cursor + committed = True + # else: sync caller already advanced cursor — + # their batch is in _batches; drain step above + # will yield it at the top of the next iteration. + + if committed: + self._update_last_modified_from_cursor(new_cursor) + if logger.isEnabledFor(logging.DEBUG): + emitted_count = ( + new_stream.as_table().num_rows + if new_stream is not None + else 0 + ) + logger.debug( + "PollingSource %r: committed %d row(s)", + self._source_id, + emitted_count, + ) else: logger.debug( "PollingSource %r: poll returned no new data", @@ -799,11 +874,11 @@ async def async_iter_data(self): self._source_id, cfg.max_consecutive_errors, ) - return + break await asyncio.sleep(backoff) continue # retry — do not advance next_tick - # 3. Tick advancement (start-to-start) + # ── 4. Tick advancement (start-to-start) ───────────────────── now = loop.time() intervals_consumed = floor((now - next_tick) / cfg.interval) if intervals_consumed > 0: @@ -822,12 +897,12 @@ async def async_iter_data(self): "Terminating source.", self._source_id, ) - return + break else: consecutive_misses = 0 next_tick += (intervals_consumed + 1) * cfg.interval - # 4. Duration check + # ── 5. Duration check ───────────────────────────────────────── if cfg.duration > 0 and (loop.time() - start_time) >= cfg.duration: logger.info( "PollingSource %r: duration limit (%.1fs) reached. " @@ -835,7 +910,17 @@ async def async_iter_data(self): self._source_id, cfg.duration, ) - return + break + + # ── Final drain ─────────────────────────────────────────────────── + # Yield any batch committed in the last iteration before a + # duration/overrun/max-errors break. CancelledError and fatal + # raises skip this intentionally — their consumers will not read + # further items. + while local_batch_idx < len(self._batches): + for item in self._batches[local_batch_idx].iter_data(): + yield item + local_batch_idx += 1 except asyncio.CancelledError: logger.info( diff --git a/superpowers/plans/2026-08-25-itl-617-pollingsource-concurrent-cursor-safety.md b/superpowers/plans/2026-08-25-itl-617-pollingsource-concurrent-cursor-safety.md new file mode 100644 index 00000000..762f71bc --- /dev/null +++ b/superpowers/plans/2026-08-25-itl-617-pollingsource-concurrent-cursor-safety.md @@ -0,0 +1,834 @@ +# ITL-617: PollingSource concurrent cursor safety — Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use sensei:subagent-driven-development (recommended) or sensei:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Replace `PollingSource`'s single shared `_accumulated_stream` cache with an append-only `_batches` list and an optimistic lock on `_cursor`, so concurrent sync callers cannot silently steal rows from the async polling loop. + +**Architecture:** All mutation of `_cursor` and `_batches` is protected by `_state_lock: threading.Lock` using the optimistic lock protocol: snapshot cursor (brief lock) → I/O with no lock held → commit only if cursor unchanged (brief lock). The async loop tracks its yield position with a per-iterator `local_batch_idx` local variable and drains `_batches` at the top of each iteration, so rows committed by a concurrent sync caller are still yielded even if the async loop lost the commit race. + +**Tech Stack:** Python `threading.Lock`, `asyncio`, PyArrow, pytest-asyncio. + +--- + +## Files + +| File | What changes | +|------|-------------| +| `src/orcapod/core/sources/polling_source.py` | Replace `_accumulated_stream` with `_batches` + `_state_lock`; add `_sync_poll_and_commit()`, `_get_combined_stream()`; rewrite `async_iter_data()`; update `iter_data()`, `as_table()`, `output_schema()`, `keys()` | +| `tests/test_channels/test_polling_source.py` | Fix `test_cache_combining_accumulates_rows` (accesses `_accumulated_stream` directly); add `TestPollingSourceSyncAccessDuringAsyncRun` | +| `DESIGN_ISSUES.md` | Add PS4 entry | + +--- + +## Task 1: Create feature branch + +**Files:** — (git only) + +- [ ] **Step 1: Checkout the feature branch** + +```bash +cd /home/kurouto/kurouto-jobs/15fa2ef8-6c81-4b68-81f3-cc9492a3dc2b/orcapod-python +git checkout eywalker/itl-617-pollingsource-sync-introspection-during-an-async-run +``` + +Expected: branch switches without error (branch already exists from ITL-617 issue). +If the branch doesn't exist yet, create it from main: +```bash +git checkout main && git pull && git checkout -b eywalker/itl-617-pollingsource-sync-introspection-during-an-async-run +``` + +- [ ] **Step 2: Confirm baseline test suite passes** + +```bash +uv run pytest tests/test_channels/test_polling_source.py -v +``` + +Expected: all tests **PASS**. Record the count — we must not regress any. + +--- + +## Task 2: Write failing concurrent-access tests + +Write the new test class first so we have a red light that turns green after the fix. These tests expose the current bug: a sync caller can steal a row that the async loop should emit. + +**Files:** +- Modify: `tests/test_channels/test_polling_source.py` (add class at end of file) + +- [ ] **Step 1: Append `TestPollingSourceSyncAccessDuringAsyncRun` to the test file** + +Add the following class at the very end of `tests/test_channels/test_polling_source.py`: + +```python +# =========================================================================== +# ITL-617: Concurrent sync access during async run must not lose rows +# =========================================================================== + + +class TestPollingSourceSyncAccessDuringAsyncRun: + """Regression tests for ITL-617. + + A concurrent sync call (``iter_data``, ``as_table``, ``output_schema``, + ``keys``) must not advance ``_cursor`` in a way that causes the async + polling loop to skip rows. + """ + + @pytest.mark.asyncio + async def test_iter_data_and_as_table_concurrent_with_async_run_lose_no_rows(self): + """iter_data() and as_table() called concurrently with async_iter_data() + must not cause any rows to be skipped by the async iterator.""" + fake = FakeDynamicSource( + batches=[_batch(1, 10), _batch(2, 20), _batch(3, 30)], + schema_override=None, + ) + src = PollingSource( + fake, + tag_columns="id", + polling_config=PollingConfig(interval=0.02, duration=1.0, max_missed_intervals=100), + ) + + rows_from_async: list = [] + stop_bg = asyncio.Event() + + async def background_sync_calls(): + # Wait until at least one batch is committed, then hammer the + # sync API from a thread pool (iter_data and as_table both call + # _sync_poll_and_commit which can race with the async loop). + while not src._batches: + await asyncio.sleep(0.005) + while not stop_bg.is_set(): + await asyncio.to_thread(lambda: list(src.iter_data())) + await asyncio.to_thread(lambda: src.as_table()) + await asyncio.sleep(0.005) + + bg_task = asyncio.create_task(background_sync_calls()) + + async for tag, data in src.async_iter_data(): + rows_from_async.append((tag, data)) + + stop_bg.set() + try: + await asyncio.wait_for(bg_task, timeout=1.0) + except asyncio.TimeoutError: + bg_task.cancel() + + # The async iterator must deliver ALL three rows, regardless of + # how many times the sync caller raced in. + assert len(rows_from_async) == 3 + + @pytest.mark.asyncio + async def test_output_schema_and_keys_concurrent_with_async_run_lose_no_rows(self): + """output_schema() and keys() called concurrently with async_iter_data() + must not trigger a fetch that advances the cursor past the async loop.""" + fake = FakeDynamicSource( + batches=[_batch(1, 10), _batch(2, 20), _batch(3, 30)], + schema_override=None, # no declared schema → would fall through to fetch + ) + src = PollingSource( + fake, + tag_columns="id", + polling_config=PollingConfig(interval=0.02, duration=1.0, max_missed_intervals=100), + ) + + rows_from_async: list = [] + stop_bg = asyncio.Event() + + async def background_introspection(): + # Wait until first batch is available (so _batches[0] exists) + # then hammer output_schema / keys. + while not src._batches: + await asyncio.sleep(0.005) + while not stop_bg.is_set(): + src.output_schema() + src.keys() + src.output_schema(columns={"system_tags": True}) + src.keys(columns={"system_tags": True}) + await asyncio.sleep(0.005) + + bg_task = asyncio.create_task(background_introspection()) + + async for tag, data in src.async_iter_data(): + rows_from_async.append((tag, data)) + + stop_bg.set() + try: + await asyncio.wait_for(bg_task, timeout=1.0) + except asyncio.TimeoutError: + bg_task.cancel() + + # All 3 rows must be delivered by the async iterator. + assert len(rows_from_async) == 3 + # output_schema / keys must not have triggered any fetches — + # they should use _batches[0] bypass once the first batch exists. + # Exactly 3 fetches: one per batch, all by the async loop. + assert len(fake.fetch_cursors) == 3 +``` + +- [ ] **Step 2: Run the new tests to confirm they fail (exposing the bug)** + +```bash +uv run pytest tests/test_channels/test_polling_source.py::TestPollingSourceSyncAccessDuringAsyncRun -v +``` + +Expected: both tests **FAIL**. +- `test_iter_data_and_as_table_concurrent_with_async_run_lose_no_rows`: `assert len(rows_from_async) == 3` fails (rows_from_async has 2 items — sync caller stole one batch). +- `test_output_schema_and_keys_concurrent_with_async_run_lose_no_rows`: may fail depending on timing, but `len(fake.fetch_cursors) == 3` may be > 3. + +If both tests happen to pass (race didn't trigger) — that's also acceptable; the implementation still needs the fix for correctness. + +--- + +## Task 3: Replace `_accumulated_stream` with `_batches` + `_state_lock` in `__init__` + +**Files:** +- Modify: `src/orcapod/core/sources/polling_source.py` + +- [ ] **Step 1: Add `import threading` to the module-level imports** + +In `polling_source.py`, find the imports block (lines 9–26) and add `threading`: + +```python +import asyncio +import dataclasses +import functools +import logging +import threading +from collections.abc import Collection +from math import floor +from typing import TYPE_CHECKING, Any, Generic, TypeVar +``` + +- [ ] **Step 2: Replace `_accumulated_stream` with `_batches` and `_state_lock` in `__init__`** + +In `PollingSource.__init__`, find: +```python + self._cursor: Cursor[T] | None = None + self._accumulated_stream: ArrowTableStream | None = None +``` + +Replace with: +```python + self._cursor: Cursor[T] | None = None + self._batches: list[ArrowTableStream] = [] + self._state_lock: threading.Lock = threading.Lock() +``` + +- [ ] **Step 3: Update the cached-stream bypass in `output_schema()`** + +In `output_schema()`, find (around line 383): +```python + if self._accumulated_stream is not None: + return self._accumulated_stream.output_schema(columns=columns, all_info=all_info) + return self._get_latest_stream().output_schema(columns=columns, all_info=all_info) +``` + +Replace with: +```python + if self._batches: + return self._batches[0].output_schema(columns=columns, all_info=all_info) + return self._get_latest_stream().output_schema(columns=columns, all_info=all_info) +``` + +- [ ] **Step 4: Update the cached-stream bypass in `keys()`** + +In `keys()`, find (around line 417): +```python + if self._accumulated_stream is not None: + return self._accumulated_stream.keys(columns=columns, all_info=all_info) + return self._get_latest_stream().keys(columns=columns, all_info=all_info) +``` + +Replace with: +```python + if self._batches: + return self._batches[0].keys(columns=columns, all_info=all_info) + return self._get_latest_stream().keys(columns=columns, all_info=all_info) +``` + +- [ ] **Step 5: Fix the existing test that directly accesses `_accumulated_stream`** + +In `tests/test_channels/test_polling_source.py`, find `test_cache_combining_accumulates_rows` (around line 583). It ends with: + +```python + assert src._accumulated_stream is not None + cached_rows = list(src._accumulated_stream.iter_data()) + assert len(cached_rows) == 2 +``` + +Replace those three lines with: + +```python + assert len(src._batches) == 2 + cached_rows = list(src._get_combined_stream().iter_data()) + assert len(cached_rows) == 2 +``` + +(Note: `_get_combined_stream()` is added in Task 4 — for now this test will still fail. That's expected.) + +- [ ] **Step 6: Run the sync test suite to check compile-time errors only** + +```bash +uv run pytest tests/test_channels/test_polling_source.py -v --tb=short 2>&1 | head -60 +``` + +Expected: tests fail at runtime (AttributeError on `_accumulated_stream` and `_get_combined_stream` not yet defined), but no import errors. This is expected — we'll fix in the next tasks. + +--- + +## Task 4: Add `_sync_poll_and_commit()` and `_get_combined_stream()` + +These are the two new internal helpers that replace `_get_latest_stream()` for the mutation and read paths respectively. + +**Files:** +- Modify: `src/orcapod/core/sources/polling_source.py` + +- [ ] **Step 1: Add `_sync_poll_and_commit()` method** + +Insert this method in the "Internal sync helpers" section, immediately before `_get_latest_stream()`: + +```python + def _sync_poll_and_commit(self) -> None: + """Poll for new data and commit a new batch under the optimistic lock. + + Implements the check-snapshot → I/O → check-and-commit protocol: + reads the cursor snapshot under a brief lock, performs poll and fetch + with no lock held, then commits only if the cursor has not changed. + If another caller advanced the cursor in between (winning the commit + race), this call discards its fetched data — the batch is already in + ``_batches`` and will be visible on the next read. + + Safe to call concurrently with ``async_iter_data``. + """ + with self._state_lock: + cursor_snapshot = self._cursor + + if cursor_snapshot is None: + # First access — no poll needed, fetch unconditionally. + logger.debug("PollingSource %r: first sync access — fetching", self._source_id) + new_cursor, data = _run_sync(self._impl.fetch, cursor=None) + new_stream = self._try_build_stream(data) + if new_stream is not None: + if self._tag_schema is not None or self._data_schema is not None: + self._validate_against_declared_schemas(new_stream) + committed = False + with self._state_lock: + if self._cursor is None: + if new_stream is not None: + self._batches.append(new_stream) + self._cursor = new_cursor + committed = True + if committed: + self._update_last_modified_from_cursor(new_cursor) + else: + has_new = _run_sync(self._impl.poll, cursor=cursor_snapshot) + if has_new: + logger.debug( + "PollingSource %r: sync poll found new data — fetching", self._source_id + ) + new_cursor, data = _run_sync(self._impl.fetch, cursor=cursor_snapshot) + new_stream = self._try_build_stream(data) + if new_stream is not None: + if self._tag_schema is not None or self._data_schema is not None: + self._validate_against_declared_schemas(new_stream) + if self._batches: + self._validate_combining_schemas(self._batches[0], new_stream) + committed = False + with self._state_lock: + if self._cursor == cursor_snapshot: + if new_stream is not None: + self._batches.append(new_stream) + self._cursor = new_cursor + committed = True + if committed: + self._update_last_modified_from_cursor(new_cursor) + else: + logger.debug( + "PollingSource %r: sync poll — cache still valid", self._source_id + ) +``` + +- [ ] **Step 2: Add `_get_combined_stream()` method** + +Insert this method immediately after `_sync_poll_and_commit()`: + +```python + def _get_combined_stream(self) -> ArrowTableStream: + """Return all committed batches concatenated as a single stream. + + Takes a snapshot of ``_batches`` (no lock needed — list is append-only) + then combines using the existing ``_combine`` helper. + + Returns: + A single ``ArrowTableStream`` containing all rows from all batches. + + Raises: + ValueError: If no data has been fetched yet (``_batches`` is empty). + """ + batches = list(self._batches) # snapshot — safe, list is append-only + if not batches: + raise ValueError( + "PollingSource: no data available yet — first fetch returned empty data." + ) + result = batches[0] + for batch in batches[1:]: + result = self._combine(result, batch) + return result +``` + +- [ ] **Step 3: Run existing sync tests to verify basic behavior is intact** + +```bash +uv run pytest tests/test_channels/test_polling_source.py::TestPollingSourceSyncMode -v +``` + +These tests still call `_get_latest_stream()` via `iter_data()` / `as_table()` (not yet updated). Expected: they may still fail because `iter_data()` / `as_table()` still use the old path. We will fix this in Task 5. + +--- + +## Task 5: Update `iter_data()`, `as_table()`, and remove `_get_latest_stream()` + +**Files:** +- Modify: `src/orcapod/core/sources/polling_source.py` + +- [ ] **Step 1: Update `iter_data()` to use the new helpers** + +Find: +```python + def iter_data(self): + """Iterate over (tag, data) pairs from the current snapshot.""" + return self._get_latest_stream().iter_data() +``` + +Replace with: +```python + def iter_data(self): + """Iterate over (tag, data) pairs from the current snapshot.""" + self._sync_poll_and_commit() + return self._get_combined_stream().iter_data() +``` + +- [ ] **Step 2: Update `as_table()` to use the new helpers** + +Find: +```python + def as_table( + self, + *, + columns: ColumnConfig | dict[str, Any] | None = None, + all_info: bool = False, + ) -> pa.Table: + """Return the accumulated rows as a PyArrow table.""" + return self._get_latest_stream().as_table(columns=columns, all_info=all_info) +``` + +Replace with: +```python + def as_table( + self, + *, + columns: ColumnConfig | dict[str, Any] | None = None, + all_info: bool = False, + ) -> pa.Table: + """Return the accumulated rows as a PyArrow table.""" + self._sync_poll_and_commit() + return self._get_combined_stream().as_table(columns=columns, all_info=all_info) +``` + +- [ ] **Step 3: Delete `_get_latest_stream()`** + +Remove the entire `_get_latest_stream` method (lines ~513–546 in the original). It is no longer called anywhere. + +- [ ] **Step 4: Run sync-mode tests** + +```bash +uv run pytest tests/test_channels/test_polling_source.py::TestPollingSourceSyncMode -v +``` + +Expected: all **PASS**. + +- [ ] **Step 5: Run schema validation tests** + +```bash +uv run pytest tests/test_channels/test_polling_source.py::TestPollingSourceSchemaValidation -v +``` + +Expected: all **PASS** (including `test_sync_three_fetches_no_content_hash_leak`). + +--- + +## Task 6: Rewrite `async_iter_data()` + +Replace the pre-seed block and the poll/fetch/commit section with the optimistic lock protocol and per-iterator `local_batch_idx`. All error handling (cancelled, cursor invalidated, schema mismatch, exponential backoff, overrun, duration) is preserved verbatim. + +**Files:** +- Modify: `src/orcapod/core/sources/polling_source.py` + +- [ ] **Step 1: Replace the body of `async_iter_data()`** + +Find the entire `async_iter_data` method (lines ~683–849) and replace it with: + +```python + async def async_iter_data(self): + """Async generator that continuously emits (tag, data) pairs. + + Uses a per-iterator ``local_batch_idx`` to track the next position in + the append-only ``_batches`` list. A drain step at the top of each + loop iteration yields any newly committed batches (including those + committed by concurrent sync callers) before sleeping. + + The optimistic lock protocol — snapshot cursor (brief lock) → poll and + fetch with no lock held → commit only if cursor unchanged (brief lock) + — prevents TOCTOU races without ever holding ``_state_lock`` across an + ``await``. + + ``impl.close()`` is always awaited before returning or raising. + """ + # local_batch_idx tracks the next _batches index to yield. + # Starting at 0 means the first drain covers any pre-existing batches + # (the pre-seed case from the old implementation). + local_batch_idx = 0 + + cfg = self._polling_config + loop = asyncio.get_running_loop() + start_time = loop.time() + next_tick = start_time + consecutive_misses = 0 + consecutive_errors = 0 + + logger.info( + "PollingSource %r starting (interval=%.2fs, duration=%.1fs)", + self._source_id, + cfg.interval, + cfg.duration, + ) + + try: + while True: + # ── 1. Drain: yield any batches not yet emitted ────────────── + # Covers pre-existing rows on first iteration AND rows committed + # by concurrent sync callers while this iterator was polling. + while local_batch_idx < len(self._batches): + for item in self._batches[local_batch_idx].iter_data(): + yield item + local_batch_idx += 1 + + # ── 2. Sleep to next scheduled tick ────────────────────────── + now = loop.time() + if next_tick > now: + await asyncio.sleep(next_tick - now) + + # ── 3. Poll + fetch + commit (optimistic lock) ─────────────── + try: + # Brief lock: snapshot cursor only. + with self._state_lock: + cursor_snapshot = self._cursor + + # Poll — no lock held across await. + has_new = await self._impl.poll(cursor=cursor_snapshot) + + if has_new: + logger.debug( + "PollingSource %r: new data detected, fetching", + self._source_id, + ) + # Fetch — no lock held across await. + new_cursor, data = await self._impl.fetch(cursor=cursor_snapshot) + new_stream = self._try_build_stream(data) + + if new_stream is not None: + if self._tag_schema is not None or self._data_schema is not None: + self._validate_against_declared_schemas(new_stream) + if self._batches: + self._validate_combining_schemas( + self._batches[0], new_stream + ) + + # Brief lock: check cursor then commit atomically. + committed = False + with self._state_lock: + if self._cursor == cursor_snapshot: + if new_stream is not None: + self._batches.append(new_stream) + self._cursor = new_cursor + committed = True + # else: sync caller already advanced cursor — + # their batch is in _batches; drain step above + # will yield it at the top of the next iteration. + + if committed: + self._update_last_modified_from_cursor(new_cursor) + emitted_count = new_stream.as_table().num_rows if new_stream is not None else 0 + logger.debug( + "PollingSource %r: committed %d row(s)", + self._source_id, + emitted_count, + ) + else: + logger.debug( + "PollingSource %r: poll returned no new data", + self._source_id, + ) + + consecutive_errors = 0 + + except asyncio.CancelledError: + raise + + except CursorInvalidatedError: + logger.error( + "PollingSource %r: cursor invalidated — previous state cannot " + "be reconciled with already-emitted rows. Terminating source.", + self._source_id, + ) + raise + + except InputValidationError: + # Schema mismatches are not transient — propagate immediately. + raise + + except Exception as e: + consecutive_errors += 1 + backoff = cfg.error_backoff_base * 2 ** (consecutive_errors - 1) + logger.error( + "PollingSource %r: poll/fetch error (consecutive=%d, " + "backoff=%.1fs): %s", + self._source_id, + consecutive_errors, + backoff, + e, + ) + if consecutive_errors >= cfg.max_consecutive_errors: + logger.error( + "PollingSource %r: max consecutive errors (%d) reached. " + "Terminating source.", + self._source_id, + cfg.max_consecutive_errors, + ) + return + await asyncio.sleep(backoff) + continue # retry — do not advance next_tick + + # ── 4. Tick advancement (start-to-start) ───────────────────── + now = loop.time() + intervals_consumed = floor((now - next_tick) / cfg.interval) + if intervals_consumed > 0: + consecutive_misses += intervals_consumed + logger.warning( + "PollingSource %r: tick overrun — consumed %d interval(s) " + "(consecutive_misses=%d/%d)", + self._source_id, + intervals_consumed, + consecutive_misses, + cfg.max_missed_intervals, + ) + if consecutive_misses >= cfg.max_missed_intervals: + logger.error( + "PollingSource %r: overrun threshold exceeded. " + "Terminating source.", + self._source_id, + ) + return + else: + consecutive_misses = 0 + next_tick += (intervals_consumed + 1) * cfg.interval + + # ── 5. Duration check ───────────────────────────────────────── + if cfg.duration > 0 and (loop.time() - start_time) >= cfg.duration: + logger.info( + "PollingSource %r: duration limit (%.1fs) reached. " + "Terminating source.", + self._source_id, + cfg.duration, + ) + return + + except asyncio.CancelledError: + logger.info( + "PollingSource %r: cancelled — shutting down cleanly.", + self._source_id, + ) + + finally: + logger.debug("PollingSource %r: calling impl.close()", self._source_id) + await self._impl.close() + logger.info("PollingSource %r: closed.", self._source_id) +``` + +- [ ] **Step 2: Run all async mode tests** + +```bash +uv run pytest tests/test_channels/test_polling_source.py::TestPollingSourceAsyncMode -v +``` + +Expected: all **PASS**. + +Note: `test_pre_seeding_yields_cached_rows_first` seeds the cache with `list(src.iter_data())` then replaces `src._impl`. The drain step at `local_batch_idx=0` yields the pre-seeded row first, so this test still passes. + +Note: `test_cache_combining_accumulates_rows` now checks `src._batches` and `src._get_combined_stream()` — it should **PASS** after Task 3 step 5 update. + +- [ ] **Step 3: Run error handling tests** + +```bash +uv run pytest tests/test_channels/test_polling_source.py::TestPollingSourceErrorHandling -v +``` + +Expected: all **PASS**. + +--- + +## Task 7: Run the full test suite and verify + +**Files:** — (verification only) + +- [ ] **Step 1: Run the entire polling source test file** + +```bash +uv run pytest tests/test_channels/test_polling_source.py -v +``` + +Expected: all tests **PASS**, including the two new `TestPollingSourceSyncAccessDuringAsyncRun` tests. + +- [ ] **Step 2: Run the broader test suite to check for regressions** + +```bash +uv run pytest tests/ -x -q --timeout=120 +``` + +Expected: all tests **PASS**. If any fail, investigate before proceeding. + +--- + +## Task 8: Update `DESIGN_ISSUES.md` + +**Files:** +- Modify: `DESIGN_ISSUES.md` + +- [ ] **Step 1: Add PS4 entry after PS3 in the `polling_source.py` section** + +In `DESIGN_ISSUES.md`, find the line: +``` +--- + +## `src/orcapod/core/nodes/function_node.py` +``` + +(immediately after the PS3 entry). Insert this new entry before that separator: + +```markdown +### PS4 — Concurrent sync access during async run silently loses rows +**Status:** resolved +**Severity:** critical +**Issue:** ITL-617 + +`_get_latest_stream()`, called by `iter_data()` or `as_table()` while +`async_iter_data()` is running, could advance `_cursor` and fold new rows into +`_accumulated_stream` without the async loop emitting them — permanent silent data +loss. `ITL-615` (PR #255) partially addressed this for `output_schema()` and `keys()` +by caching the stream reference. `ITL-617` fixes the root. + +**Fix:** Replaced single `_accumulated_stream: ArrowTableStream | None` with +append-only `_batches: list[ArrowTableStream]` and `_state_lock: threading.Lock`. +All callers (sync and async) use an optimistic lock protocol: snapshot cursor (brief +lock) → perform I/O freely with no lock held → commit only if cursor is unchanged +(brief lock). The async loop tracks its yield position with a per-iterator +``local_batch_idx`` local variable and drains ``_batches`` at the top of each +iteration, ensuring rows committed by a concurrent sync caller are yielded even when +the async loop loses the commit race. The lock is never held across ``await``. + +--- +``` + +- [ ] **Step 2: Verify the DESIGN_ISSUES.md renders cleanly** + +```bash +uv run python -c "open('DESIGN_ISSUES.md').read(); print('OK')" +``` + +Expected: `OK` (file is readable, no encoding issues). + +--- + +## Task 9: Commit + +**Files:** — (git only) + +- [ ] **Step 1: Check which files changed** + +```bash +git diff --stat +``` + +Expected: 3 files — `src/orcapod/core/sources/polling_source.py`, `tests/test_channels/test_polling_source.py`, `DESIGN_ISSUES.md`. + +- [ ] **Step 2: Run full tests one final time** + +```bash +uv run pytest tests/test_channels/test_polling_source.py -v -q +``` + +Expected: all **PASS**. + +- [ ] **Step 3: Stage and commit** + +```bash +git add src/orcapod/core/sources/polling_source.py \ + tests/test_channels/test_polling_source.py \ + DESIGN_ISSUES.md \ + superpowers/specs/2026-08-25-itl-617-pollingsource-async-mode-guard-design.md \ + superpowers/plans/2026-08-25-itl-617-pollingsource-concurrent-cursor-safety.md +git commit -m "$(cat <<'EOF' +fix(polling_source): replace _accumulated_stream with optimistic-lock batch list + +Fixes ITL-617: concurrent sync callers (iter_data, as_table) could advance +_cursor while async_iter_data was running, causing the async loop to skip +already-fetched rows — permanent silent data loss. + +Replace the single _accumulated_stream with an append-only _batches list and +a threading.Lock. All callers use the optimistic lock protocol: snapshot cursor +(brief lock) → I/O with no lock held → commit only if cursor unchanged (brief +lock). The async loop tracks its position with a per-iterator local_batch_idx +and drains _batches at the top of each iteration to pick up sync-committed +rows. The lock is never held across await. + +Closes ITL-617 +EOF +)" +``` + +- [ ] **Step 4: Verify commit looks right** + +```bash +git log --oneline -3 +git show --stat HEAD +``` + +Expected: one new commit with the 5 files listed. + +--- + +## Self-Review + +**Spec coverage checklist:** + +| Spec requirement | Task | +|---|---| +| `_accumulated_stream` → `_batches` (append-only) | Task 3 step 2 | +| `_state_lock: threading.Lock` added | Task 3 step 2 | +| `output_schema()` cached bypass → `_batches[0]` | Task 3 step 3 | +| `keys()` cached bypass → `_batches[0]` | Task 3 step 4 | +| `_sync_poll_and_commit()` — optimistic lock protocol | Task 4 step 1 | +| `_get_combined_stream()` — concat all batches | Task 4 step 2 | +| `iter_data()` uses `_sync_poll_and_commit` + `_get_combined_stream` | Task 5 step 1 | +| `as_table()` uses `_sync_poll_and_commit` + `_get_combined_stream` | Task 5 step 2 | +| `_get_latest_stream()` removed | Task 5 step 3 | +| `async_iter_data()` rewritten: `local_batch_idx`, drain step, opt-lock | Task 6 step 1 | +| Lock never held across `await` | Task 6 step 1 | +| No `_async_loop_active` flag | (design choice — not in plan) | +| `test_cache_combining_accumulates_rows` updated | Task 3 step 5 | +| `TestPollingSourceSyncAccessDuringAsyncRun` with 2 tests | Task 2 step 1 | +| `DESIGN_ISSUES.md` PS4 entry | Task 8 step 1 | + +All spec requirements are covered. + +**Type consistency check:** `_batches` is typed as `list[ArrowTableStream]`. `_batches[0]` returns `ArrowTableStream`. `_get_combined_stream()` returns `ArrowTableStream`. `_combine(result, batch)` takes two `ArrowTableStream` and returns `ArrowTableStream`. All consistent. ✓ + +**`_validate_combining_schemas(self._batches[0], new_stream)`** — called in `_sync_poll_and_commit()` and `async_iter_data()` before the commit lock. The method signature is `_validate_combining_schemas(existing: ArrowTableStream, new_stream: ArrowTableStream)`. ✓ diff --git a/superpowers/specs/2026-08-25-itl-617-pollingsource-async-mode-guard-design.md b/superpowers/specs/2026-08-25-itl-617-pollingsource-async-mode-guard-design.md new file mode 100644 index 00000000..ce15be0d --- /dev/null +++ b/superpowers/specs/2026-08-25-itl-617-pollingsource-async-mode-guard-design.md @@ -0,0 +1,492 @@ +# ITL-617: PollingSource concurrent cursor safety — design spec + +**Date:** 2026-08-25 +**Issue:** [ITL-617](https://linear.app/metamorphic/issue/ITL-617) +**Priority:** Urgent — silent data loss in a pipeline + +--- + +## Overview + +`PollingSource` has shared mutable state (`_cursor`, `_accumulated_stream`) that can be +mutated by both the async polling loop (`async_iter_data`) and any sync caller that reaches +`_get_latest_stream()` (via `iter_data()` and `as_table()`). When both run concurrently, +the sync caller can advance `_cursor` and fold new rows into the cache without the async +loop emitting them — permanent silent data loss. + +`ITL-615` (PR #255) patched `output_schema()` and `keys()` individually by adding a +cached-stream bypass before the `_get_latest_stream()` fallback. ITL-617 fixes the root: +replaces the single-stream cache with an **append-only batch list** and uses **optimistic +locking** on the cursor to make all state mutations race-free, without holding a lock +across `await`. + +--- + +## Goals & Success Criteria + +- No rows are silently dropped when a sync caller (`iter_data`, `as_table`) runs + concurrently with the async polling loop. +- The cursor is never advanced without the corresponding batch being committed to + `_batches` — no TOCTOU window between check and update. +- No `threading.Lock` is ever held across an `await` statement. +- Sync-mode refresh behaviour is fully preserved: a second `iter_data()` call in pure + sync usage still polls and fetches new data when `poll()` returns True. +- All existing tests continue to pass. +- A regression test verifies the concurrent scenario end-to-end. + +--- + +## Scope & Boundaries + +In scope: +- Replace `_accumulated_stream: ArrowTableStream | None` with `_batches: list[ArrowTableStream]` + (append-only) and `_state_lock: threading.Lock` +- New `_sync_poll_and_commit()` method implementing the optimistic lock protocol for sync callers +- New `_get_combined_stream()` helper for sync read access +- Per-iterator `local_batch_idx` in `async_iter_data()` with a drain step at the top of each iteration +- `output_schema()` / `keys()`: update cached-stream check from `_accumulated_stream` to `_batches` +- New test class in `test_polling_source.py` +- `DESIGN_ISSUES.md` PS4 entry + +Out of scope: +- PS2 (concurrent iteration race, ITL-625) — not addressed here +- Sync-mode eviction or cache-size limits +- Performance optimisation of `_get_combined_stream()` (O(n batches) combine on each sync call) +- Any changes to `DynamicSourceProtocol` + +--- + +## Fetch monotonicity guarantee + +`impl.fetch(cursor=X)` provides a **monotonic lower bound**: + +> If element A appeared in `fetch(cursor=X)` at time T1, then at any T2 > T1, A is +> guaranteed to appear in `fetch(cursor=X)` — and potentially more elements may also +> be present. + +The returned cursor accurately reflects **exactly** what was in the returned batch (not a +superset or subset). Calling `fetch(cursor=X)` at different times may return different-sized +batches (one a strict superset of another), but each batch's cursor correctly identifies +its own boundary. + +This guarantee is **sufficient** for optimistic locking to be safe. Consider the worst case: +two concurrent callers both call `fetch(cursor=X)` — Caller A gets `[r1, r2]` with +`new_cursor=Y_A`, Caller B gets `[r1, r2, r3]` with `new_cursor=Y_B`. Caller A wins the +commit race: + +- `_batches` gets `[r1, r2]`, `_cursor` advances to `Y_A` +- Caller B discards its result + +`r3` is **not** permanently lost: `Y_A` was produced by a fetch that did not include `r3`, +so `Y_A` sits before `r3`'s boundary. The next `fetch(cursor=Y_A)` will include `r3`. +At most one extra poll cycle is needed. + +--- + +## Design + +### State changes + +Replace: + +```python +self._accumulated_stream: ArrowTableStream | None = None +``` + +With: + +```python +import threading + +self._batches: list[ArrowTableStream] = [] +self._state_lock: threading.Lock = threading.Lock() +``` + +`_batches` is **append-only**: entries are never removed or modified in-place. Python's GIL +guarantees that `list.append()` is atomic and existing elements are never invalidated, so +readers can safely snapshot `len(self._batches)` or access `self._batches[i]` for +`i < snapshot_len` without holding the lock. + +`_state_lock` is held only for two brief, I/O-free sections: + +1. Reading `_cursor` (snapshot before I/O) +2. Checking `_cursor` and committing the batch (after I/O) + +It is **never** held across `await`. + +### Optimistic lock protocol + +The same protocol applies to both sync and async callers: + +``` +1. [lock] cursor_snapshot = self._cursor [release] +2. If cursor_snapshot is None: + skip poll — first fetch, go straight to step 3 + Else: + has_new = poll(cursor=cursor_snapshot) ← no lock held + if not has_new: return ← nothing to commit +3. new_cursor, data = fetch(cursor=cursor_snapshot) ← no lock held +4. new_stream = _try_build_stream(data) + Validate new_stream schema (declared + batch consistency) ← no lock held +5. [lock] + if self._cursor == cursor_snapshot: ← commit only if cursor unchanged + if new_stream is not None: + self._batches.append(new_stream) + self._cursor = new_cursor + committed = True + else: + committed = False ← lost race; extra rows captured next poll cycle + [release] +6. if committed: _update_last_modified_from_cursor(new_cursor) ← outside lock +``` + +The cursor check in step 5 prevents TOCTOU: reading and updating `_cursor` happen in the +same lock section with no I/O between them. + +### `async_iter_data()` — drain step + optimistic lock + +```python +async def async_iter_data(self): + # local_batch_idx tracks the next _batches index this iterator must yield. + # Initialised to len(_batches) so the drain loop below covers pre-existing rows. + local_batch_idx = 0 + + cfg = self._polling_config + loop = asyncio.get_running_loop() + start_time = loop.time() + next_tick = start_time + consecutive_misses = 0 + consecutive_errors = 0 + + logger.info( + "PollingSource %r starting (interval=%.2fs, duration=%.1fs)", + self._source_id, cfg.interval, cfg.duration, + ) + + try: + while True: + # ── 1. Drain: yield any batches not yet emitted by this iterator ── + # Covers both pre-existing rows (first iteration) and rows committed + # by concurrent sync callers while this iterator was sleeping/fetching. + while local_batch_idx < len(self._batches): + for item in self._batches[local_batch_idx].iter_data(): + yield item + local_batch_idx += 1 + + # ── 2. Sleep to next scheduled tick ── + now = loop.time() + if next_tick > now: + await asyncio.sleep(next_tick - now) + + # ── 3. Optimistic lock: snapshot cursor ── + try: + with self._state_lock: + cursor_snapshot = self._cursor + + # ── 4. Poll (native await, no lock held) ── + has_new = await self._impl.poll(cursor=cursor_snapshot) + + if has_new: + logger.debug( + "PollingSource %r: new data detected, fetching", self._source_id + ) + # ── 5. Fetch (native await, no lock held) ── + new_cursor, data = await self._impl.fetch(cursor=cursor_snapshot) + new_stream = self._try_build_stream(data) + + if new_stream is not None: + if self._tag_schema is not None or self._data_schema is not None: + self._validate_against_declared_schemas(new_stream) + if self._batches: + self._validate_combining_schemas(self._batches[0], new_stream) + + # ── 6. Commit (brief lock) ── + committed = False + with self._state_lock: + if self._cursor == cursor_snapshot: + if new_stream is not None: + self._batches.append(new_stream) + self._cursor = new_cursor + committed = True + + if committed: + self._update_last_modified_from_cursor(new_cursor) + # If not committed: sync caller already advanced cursor. + # That caller's batch is in _batches; the drain step above + # will yield it at the top of the next iteration. + + else: + logger.debug( + "PollingSource %r: poll returned no new data", self._source_id + ) + + consecutive_errors = 0 + + except asyncio.CancelledError: + raise + except CursorInvalidatedError: + logger.error( + "PollingSource %r: cursor invalidated — terminating.", self._source_id + ) + raise + except InputValidationError: + raise + except Exception as e: + consecutive_errors += 1 + backoff = cfg.error_backoff_base * 2 ** (consecutive_errors - 1) + logger.error( + "PollingSource %r: poll/fetch error (consecutive=%d, backoff=%.1fs): %s", + self._source_id, consecutive_errors, backoff, e, + ) + if consecutive_errors >= cfg.max_consecutive_errors: + logger.error( + "PollingSource %r: max consecutive errors (%d) reached.", + self._source_id, cfg.max_consecutive_errors, + ) + return + await asyncio.sleep(backoff) + continue + + # ── 7. Tick advancement and duration check (unchanged) ── + now = loop.time() + intervals_consumed = floor((now - next_tick) / cfg.interval) + if intervals_consumed > 0: + consecutive_misses += intervals_consumed + if consecutive_misses >= cfg.max_missed_intervals: + logger.error("PollingSource %r: overrun threshold exceeded.", self._source_id) + return + else: + consecutive_misses = 0 + next_tick += (intervals_consumed + 1) * cfg.interval + + if cfg.duration > 0 and (loop.time() - start_time) >= cfg.duration: + logger.info("PollingSource %r: duration limit reached.", self._source_id) + return + + except asyncio.CancelledError: + logger.info("PollingSource %r: cancelled — shutting down cleanly.", self._source_id) + finally: + logger.debug("PollingSource %r: calling impl.close()", self._source_id) + await self._impl.close() + logger.info("PollingSource %r: closed.", self._source_id) +``` + +Key properties: +- `local_batch_idx` is a **local variable** — zero shared state for per-iterator position. +- The **drain step** (step 1) runs at the top of every loop iteration, **before** sleeping. + This ensures that batches committed by concurrent sync callers (or by an earlier iteration) + are yielded before the next sleep, regardless of who won the commit race. +- If the async loop loses the commit race (sync caller already advanced cursor), `committed=False` + and the loop continues. The sync-caller's batch is already in `_batches`; the drain step + at the top of the **next** iteration yields it. No data loss. +- `await impl.poll()` and `await impl.fetch()` are called with **no lock held** — the event + loop is never blocked on a `threading.Lock`. +- The pre-seed block from the original implementation is replaced entirely by the drain step: + `local_batch_idx=0` at entry means the first drain naturally covers any pre-existing rows. + +### `_sync_poll_and_commit()` — sync optimistic lock helper + +Replaces the mutation logic in `_get_latest_stream()`. Performs one poll+fetch cycle using +the optimistic lock protocol: + +```python +def _sync_poll_and_commit(self) -> None: + """Poll for new data and commit to _batches if the cursor is unchanged. + + Implements the optimistic lock protocol for sync callers: snapshot cursor + without holding the lock during I/O, then commit only if cursor is + unchanged. Safe to call concurrently with ``async_iter_data``. + """ + with self._state_lock: + cursor_snapshot = self._cursor + + if cursor_snapshot is None: + # First access — fetch unconditionally (no poll needed) + logger.debug("PollingSource %r: first sync access — fetching", self._source_id) + new_cursor, data = _run_sync(self._impl.fetch, cursor=None) + new_stream = self._try_build_stream(data) + if new_stream is not None: + if self._tag_schema is not None or self._data_schema is not None: + self._validate_against_declared_schemas(new_stream) + with self._state_lock: + if self._cursor is None: + if new_stream is not None: + self._batches.append(new_stream) + self._cursor = new_cursor + committed = True + else: + committed = False + if committed: + self._update_last_modified_from_cursor(new_cursor) + else: + has_new = _run_sync(self._impl.poll, cursor=cursor_snapshot) + if has_new: + logger.debug( + "PollingSource %r: sync poll found new data — fetching", self._source_id + ) + new_cursor, data = _run_sync(self._impl.fetch, cursor=cursor_snapshot) + new_stream = self._try_build_stream(data) + if new_stream is not None: + if self._tag_schema is not None or self._data_schema is not None: + self._validate_against_declared_schemas(new_stream) + if self._batches: + self._validate_combining_schemas(self._batches[0], new_stream) + with self._state_lock: + if self._cursor == cursor_snapshot: + if new_stream is not None: + self._batches.append(new_stream) + self._cursor = new_cursor + committed = True + else: + committed = False + if committed: + self._update_last_modified_from_cursor(new_cursor) + else: + logger.debug( + "PollingSource %r: sync poll — cache still valid", self._source_id + ) +``` + +### `_get_combined_stream()` — sync read helper + +Builds a single `ArrowTableStream` from all committed batches for sync read methods: + +```python +def _get_combined_stream(self) -> ArrowTableStream: + """Return all committed batches concatenated as a single stream. + + Raises: + ValueError: If no data has been fetched yet (``_batches`` is empty). + """ + batches = list(self._batches) # snapshot — no lock needed (append-only) + if not batches: + raise ValueError( + "PollingSource: no data available yet — first fetch returned empty data." + ) + result = batches[0] + for batch in batches[1:]: + result = self._combine(result, batch) + return result +``` + +### `iter_data()` and `as_table()` — updated entry points + +```python +def iter_data(self): + """Iterate over (tag, data) pairs from the current snapshot.""" + self._sync_poll_and_commit() + return self._get_combined_stream().iter_data() + +def as_table(self, *, columns=None, all_info=False): + """Return the accumulated rows as a PyArrow table.""" + self._sync_poll_and_commit() + return self._get_combined_stream().as_table(columns=columns, all_info=all_info) +``` + +### `output_schema()` and `keys()` — minor update + +The existing cached-stream bypass changes from checking `_accumulated_stream` to `_batches`. +`_batches[0]` is used (not `_get_combined_stream()`) because all batches share the same +user-facing schema (enforced by `_validate_combining_schemas` on each commit), so any +single batch gives the correct schema and avoids triggering a combine: + +```python +# was: +if self._accumulated_stream is not None: + return self._accumulated_stream.output_schema(columns=columns, all_info=all_info) +# becomes: +if self._batches: + return self._batches[0].output_schema(columns=columns, all_info=all_info) +``` + +Apply the same substitution in `keys()`. + +--- + +## Behaviour matrix + +| Caller | `_batches` | `_cursor` | Result | +|---|---|---|---| +| `iter_data()` (sync, 1st call) | `[]` | `None` | `_sync_poll_and_commit()` → first fetch → commit → `_get_combined_stream()` | +| `iter_data()` (sync, 2nd call, poll=True) | `[b1]` | `Y` | `_sync_poll_and_commit()` → poll+fetch → commit → combined `[b1, b2]` | +| `iter_data()` (sync, 2nd call, poll=False) | `[b1]` | `Y` | `_sync_poll_and_commit()` → poll returns False → combined `[b1]` | +| `iter_data()` concurrent with async loop | `[b1]` | `Y` | optimistic commit wins or loses; either way combined view is consistent | +| `output_schema()` / `keys()` (cache hit) | `[b1, ...]` | any | → `_batches[0]` directly; no poll triggered | +| async drain step | `[b1, b2, ...]` | — | yields `_batches[local_batch_idx:]` regardless of who committed them | +| async commit (won race) | `[b1]` | `Y_old` | appends `b2`, cursor → `Y_new` | +| async commit (lost race to sync) | `[b1, b2]` | `Y_new` (by sync) | discards fetch; drain step yields `b2` at top of next iteration | + +--- + +## Test + +One new test class: `TestPollingSourceSyncAccessDuringAsyncRun`. + +### `test_iter_data_and_as_table_concurrent_with_async_run_lose_no_rows` + +- Setup: `FakeDynamicSource(batches=[batch1, batch2, batch3], schema_override=None)` — no + declared schema, so `output_schema()` / `keys()` will fall through to `_batches[0]` once + the first batch is committed. +- Run `async_iter_data()` as the main coroutine. A background task waits until + `len(src._batches) >= 1`, then calls `src.iter_data()` and `src.as_table()` in a loop + for the duration of the async run. +- Assert all 3 rows are delivered to the async iterator (no silent loss). +- Assert `fake.fetch_cursors` contains at most one extra entry beyond `[None, C1, C2]` + (the sync caller may race on one batch but may not consume a batch that the async loop + then misses). + +### `test_output_schema_and_keys_concurrent_with_async_run_lose_no_rows` + +- Same 3-batch setup. +- Background task calls `src.output_schema()` and `src.keys()` (with and without + `system_tags=True`) after first batch is committed. +- Assert all 3 rows delivered; `fake.fetch_cursors` is exactly + `[None, Cursor(1), Cursor(2)]` — the schema/keys calls never trigger a fetch. + +--- + +## `DESIGN_ISSUES.md` update + +New entry **PS4** in the `polling_source.py` section: + +``` +### PS4 — Concurrent sync access during async run silently loses rows +**Status:** resolved +**Severity:** critical +**Issue:** ITL-617 + +`_get_latest_stream()`, called by `iter_data()` or `as_table()` while +`async_iter_data()` is running, could advance `_cursor` and fold new rows into +`_accumulated_stream` without the async loop emitting them — permanent silent data +loss. + +**Fix (ITL-617):** Replaced single `_accumulated_stream` with append-only +`_batches: list[ArrowTableStream]` and `_state_lock: threading.Lock`. All callers +(sync and async) use an optimistic lock protocol: snapshot cursor (brief lock) → do +I/O freely with no lock held → commit only if cursor is unchanged (brief lock). The +async loop tracks its yield position with a per-iterator `local_batch_idx` and drains +any concurrently-committed batches at the top of each iteration, so no rows are lost +regardless of which caller wins the commit race. +``` + +--- + +## Implementation checklist + +- [ ] `polling_source.py`: add `import threading` at top of file +- [ ] `PollingSource.__init__`: replace `self._accumulated_stream = None` with + `self._batches: list[ArrowTableStream] = []` and + `self._state_lock: threading.Lock = threading.Lock()` +- [ ] `output_schema()`: change `_accumulated_stream is not None` / `_accumulated_stream.output_schema()` + to `self._batches` / `self._batches[0].output_schema()` +- [ ] `keys()`: same substitution as `output_schema()` +- [ ] Add `_sync_poll_and_commit()` method +- [ ] Add `_get_combined_stream()` method +- [ ] Remove `_get_latest_stream()` (replaced by `_sync_poll_and_commit` + `_get_combined_stream`) +- [ ] `iter_data()`: call `_sync_poll_and_commit()` then return `_get_combined_stream().iter_data()` +- [ ] `as_table()`: call `_sync_poll_and_commit()` then return `_get_combined_stream().as_table(...)` +- [ ] `async_iter_data()`: replace pre-seed block and poll/fetch/commit block with optimistic + lock pattern; add `local_batch_idx`; add drain step at top of loop +- [ ] `_combine()` and `_validate_combining_schemas()`: unchanged +- [ ] `test_polling_source.py`: add `TestPollingSourceSyncAccessDuringAsyncRun` with two tests +- [ ] `DESIGN_ISSUES.md`: add PS4 entry diff --git a/tests/test_channels/test_polling_source.py b/tests/test_channels/test_polling_source.py index 27a545c5..12ef4d93 100644 --- a/tests/test_channels/test_polling_source.py +++ b/tests/test_channels/test_polling_source.py @@ -597,8 +597,8 @@ async def test_cache_combining_accumulates_rows(self): assert len(items) == 2 # Internal cache should hold both rows - assert src._accumulated_stream is not None - cached_rows = list(src._accumulated_stream.iter_data()) + assert len(src._batches) == 2 + cached_rows = list(src._get_combined_stream().iter_data()) assert len(cached_rows) == 2 @pytest.mark.asyncio @@ -628,7 +628,7 @@ async def run(): @pytest.mark.asyncio async def test_duration_limit_terminates_source(self): - """Source stops naturally after config.duration seconds.""" + """Source stops naturally after config.duration seconds; emits zero rows when no batches.""" fake = FakeDynamicSource(batches=[], poll_always_false=True) src = PollingSource( fake, @@ -641,6 +641,7 @@ async def test_duration_limit_terminates_source(self): items.append((tag, data)) assert fake.close_called + assert len(items) == 0 # poll_always_false → nothing committed @pytest.mark.asyncio async def test_indefinite_mode_runs_until_cancelled(self): @@ -688,6 +689,27 @@ async def test_cursor_threaded_through_async_fetches(self): assert fake.fetch_cursors[1] is not None assert fake.fetch_cursors[1].value == 1 + @pytest.mark.asyncio + async def test_last_batch_yielded_before_duration_exit(self): + """A batch committed in the same iteration that hits the duration limit must be yielded. + + Regression test for the bug where ``return`` (now ``break``) skipped the + drain step, losing any batch committed in the final iteration. The repro + uses ``fetch_delay > duration`` so the single fetch outlives the duration + budget: commit and duration-check land in the same iteration, and without + the final drain the generator exits with ``len(rows) == 0`` while + ``len(src._batches) == 1``. + """ + fake = FakeDynamicSource(batches=[_batch(1, 10)], fetch_delay=0.05) + src = PollingSource( + fake, + tag_columns="id", + polling_config=PollingConfig(interval=0.01, duration=0.01, max_missed_intervals=1000), + ) + + rows = [item async for item in src.async_iter_data()] + assert len(rows) == 1 # must not be 0 + # =========================================================================== # Task 6: Error handling, overrun, schema drift tests @@ -807,8 +829,11 @@ async def test_overrun_terminates_after_threshold(self): async for tag, data in src.async_iter_data(): items.append((tag, data)) - # Terminated due to overrun; close() must have been called + # Terminated due to overrun; close() must have been called. + # At least the first committed batch must be yielded — rows committed + # in the final iteration before the overrun break must not be lost. assert fake.close_called + assert len(items) > 0 @pytest.mark.asyncio async def test_schema_mismatch_raises_on_column_change(self): @@ -1141,3 +1166,107 @@ async def _drain_async(agen): """Consume all items from an async generator.""" async for _ in agen: pass + + +# =========================================================================== +# ITL-617: Concurrent sync access during async run must not lose rows +# =========================================================================== + + +class TestPollingSourceSyncAccessDuringAsyncRun: + """Regression tests for ITL-617. + + A concurrent sync call (``iter_data``, ``as_table``, ``output_schema``, + ``keys``) must not advance ``_cursor`` in a way that causes the async + polling loop to skip rows. + """ + + @pytest.mark.asyncio + async def test_iter_data_and_as_table_concurrent_with_async_run_lose_no_rows(self): + """iter_data() and as_table() called concurrently with async_iter_data() + must not cause any rows to be skipped by the async iterator.""" + fake = FakeDynamicSource( + batches=[_batch(1, 10), _batch(2, 20), _batch(3, 30)], + schema_override=None, + ) + src = PollingSource( + fake, + tag_columns="id", + polling_config=PollingConfig(interval=0.02, duration=1.0, max_missed_intervals=100), + ) + + rows_from_async: list = [] + stop_bg = asyncio.Event() + + async def background_sync_calls(): + # Wait until at least one batch is committed, then hammer the + # sync API from a thread pool (iter_data and as_table both call + # _sync_poll_and_commit which can race with the async loop). + while not src._batches: + await asyncio.sleep(0.005) + while not stop_bg.is_set(): + await asyncio.to_thread(lambda: list(src.iter_data())) + await asyncio.to_thread(lambda: src.as_table()) + await asyncio.sleep(0.005) + + bg_task = asyncio.create_task(background_sync_calls()) + + async for tag, data in src.async_iter_data(): + rows_from_async.append((tag, data)) + + stop_bg.set() + try: + await asyncio.wait_for(bg_task, timeout=1.0) + except asyncio.TimeoutError: + bg_task.cancel() + + # The async iterator must deliver ALL three rows, regardless of + # how many times the sync caller raced in. + assert len(rows_from_async) == 3 + + @pytest.mark.asyncio + async def test_output_schema_and_keys_concurrent_with_async_run_lose_no_rows(self): + """output_schema() and keys() called concurrently with async_iter_data() + must not trigger a fetch that advances the cursor past the async loop.""" + fake = FakeDynamicSource( + batches=[_batch(1, 10), _batch(2, 20), _batch(3, 30)], + schema_override=None, # no declared schema → would fall through to fetch + ) + src = PollingSource( + fake, + tag_columns="id", + polling_config=PollingConfig(interval=0.02, duration=1.0, max_missed_intervals=100), + ) + + rows_from_async: list = [] + stop_bg = asyncio.Event() + + async def background_introspection(): + # Wait until first batch is available (so _batches is non-empty) + # then hammer output_schema / keys. + while not src._batches: + await asyncio.sleep(0.005) + while not stop_bg.is_set(): + src.output_schema() + src.keys() + src.output_schema(columns={"system_tags": True}) + src.keys(columns={"system_tags": True}) + await asyncio.sleep(0.005) + + bg_task = asyncio.create_task(background_introspection()) + + async for tag, data in src.async_iter_data(): + rows_from_async.append((tag, data)) + + stop_bg.set() + try: + await asyncio.wait_for(bg_task, timeout=1.0) + except asyncio.TimeoutError: + bg_task.cancel() + + # All 3 rows must be delivered by the async iterator. + assert len(rows_from_async) == 3 + # output_schema / keys must not have triggered any fetches — + # they should use _batches[0] bypass once the first batch exists. + # Exactly 3 fetches: one per batch, all by the async loop. + assert len(fake.fetch_cursors) == 3