fix: drain until sentinel to prevent deadlock on early generator close

When concurrent_files > 1 and max_buffered_batches is small, multiple
workers can be blocked on batch_queue.put() at the moment the consumer
closes early (e.g. due to a limit). The previous drain loop used
get_nowait() + empty() which had a race: empty() could return True
before a just-notified worker had a chance to put, leaving remaining
workers stuck on put() forever while executor.shutdown(wait=True) hung.

Fix: replace the racy drain loop with a blocking drain-until-sentinel
loop. Each get() naturally wakes one blocked worker via not_full.notify();
that worker checks cancel and returns, eventually allowing the last worker
to put the sentinel. Stopping only on the sentinel guarantees all workers
have finished before we exit.

Also move batch_queue.put(_QUEUE_SENTINEL) outside remaining_lock to
avoid holding a lock during a potentially blocking call.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

Author: sumedhsakdeo

pyiceberg/io/pyarrow.py

@@ -1727,18 +1727,21 @@ def worker(task: FileScanTask) -> None:
         finally:
             with remaining_lock:
                 remaining -= 1
-                if remaining == 0:
-                    batch_queue.put(_QUEUE_SENTINEL)
+                is_last = remaining == 0
+            if is_last:
+                batch_queue.put(_QUEUE_SENTINEL)
 
     with ThreadPoolExecutor(max_workers=concurrent_streams) as executor:
         for task in tasks:
             executor.submit(worker, task)
 
+        saw_sentinel = False
         try:
             while True:
                 item = batch_queue.get()
 
                 if item is _QUEUE_SENTINEL:
+                    saw_sentinel = True
                     break
 
                 if isinstance(item, BaseException):
@@ -1747,12 +1750,16 @@ def worker(task: FileScanTask) -> None:
                 yield item
         finally:
             cancel.set()
-            # Drain the queue to unblock any workers stuck on put()
-            while not batch_queue.empty():
-                try:
-                    batch_queue.get_nowait()
-                except queue.Empty:
-                    break
+            if not saw_sentinel:
+                # Drain the queue to unblock workers stuck on put().
+                # Each get() wakes one waiting producer; that producer checks
+                # cancel and returns, eventually allowing the last worker to
+                # put the sentinel.  We stop only when we see the sentinel,
+                # which guarantees all workers have finished.
+                while True:
+                    item = batch_queue.get()
+                    if item is _QUEUE_SENTINEL:
+                        break
 
 
 _DEFAULT_SCAN_ORDER: ScanOrder = TaskOrder()

tests/io/test_bounded_concurrent_batches.py

@@ -142,6 +142,53 @@ def batch_fn(t: FileScanTask) -> Iterator[pa.RecordBatch]:
         list(gen)
 
 
+def test_no_deadlock_on_early_termination_with_full_queue() -> None:
+    """Regression test for deadlock when concurrent_files > 1 and max_buffered_batches=1.
+
+    Scenario:
+      - max_buffered_batches=1, concurrent_files=3
+      - Queue is full (1 item), Workers A/B/C all blocked on put()
+      - Consumer closes the generator (GeneratorExit → finally → cancel.set())
+      - Buggy drain loop: get_nowait() removes the 1 item, empty() → True,
+        exits before A/B/C have a chance to put. Workers remain stuck on put()
+        and executor.shutdown(wait=True) hangs forever.
+    """
+    tasks = [_make_task() for _ in range(3)]
+
+    blocked_on_put = 0
+    blocked_lock = threading.Lock()
+    at_least_two_blocked = threading.Event()
+
+    def batch_fn(t: FileScanTask) -> Iterator[pa.RecordBatch]:
+        nonlocal blocked_on_put
+        for i in range(100):
+            with blocked_lock:
+                blocked_on_put += 1
+                if blocked_on_put >= 2:
+                    at_least_two_blocked.set()
+            yield pa.record_batch({"col": [i]})
+            with blocked_lock:
+                blocked_on_put -= 1
+
+    gen = _bounded_concurrent_batches(tasks, batch_fn, concurrent_streams=3, max_buffered_batches=1)
+    next(gen)  # consume 1 batch, workers immediately try to put the next
+
+    # Wait until at least 2 workers are inside batch_fn trying to put — confirming
+    # they will block on put() since the queue is full (maxsize=1)
+    assert at_least_two_blocked.wait(timeout=5.0), "Workers did not reach put() in time"
+
+    closed = threading.Event()
+
+    def close_gen() -> None:
+        gen.close()
+        closed.set()
+
+    t = threading.Thread(target=close_gen, daemon=True)
+    t.start()
+    assert closed.wait(timeout=10.0), "Deadlock: gen.close() did not complete — workers stuck on put()"
+    t.join(timeout=1.0)
+
+
 def test_early_termination() -> None:
     """Test that stopping consumption cancels workers."""
     tasks = [_make_task() for _ in range(5)]