Use threading in AsyncMapper.produce()

src/datachain/asyn.py

@@ -1,5 +1,12 @@
 import asyncio
-from collections.abc import AsyncIterable, Awaitable, Coroutine, Iterable, Iterator
+from collections.abc import (
+    AsyncIterable,
+    Awaitable,
+    Coroutine,
+    Generator,
+    Iterable,
+    Iterator,
+)
 from concurrent.futures import ThreadPoolExecutor
 from heapq import heappop, heappush
 from typing import Any, Callable, Generic, Optional, TypeVar
@@ -54,9 +61,13 @@ def start_task(self, coro: Coroutine) -> asyncio.Task:
         task.add_done_callback(self._tasks.discard)
         return task
 
-    async def produce(self) -> None:
+    def _produce(self) -> None:
         for item in self.iterable:
-            await self.work_queue.put(item)
+            fut = asyncio.run_coroutine_threadsafe(self.work_queue.put(item), self.loop)
+            fut.result()  # wait until the item is in the queue
+
+    async def produce(self) -> None:
+        await self.to_thread(self._produce)
 
     async def worker(self) -> None:
         while (item := await self.work_queue.get()) is not None:
@@ -132,7 +143,7 @@ async def _break_iteration(self) -> None:
             self.result_queue.get_nowait()
         await self.result_queue.put(None)
 
-    def iterate(self, timeout=None) -> Iterable[ResultT]:
+    def iterate(self, timeout=None) -> Generator[ResultT, None, None]:
         init = asyncio.run_coroutine_threadsafe(self.init(), self.loop)
         init.result(timeout=1)
         async_run = asyncio.run_coroutine_threadsafe(self.run(), self.loop)

tests/unit/test_asyn.py

@@ -2,6 +2,7 @@
 import functools
 from collections import Counter
 from contextlib import contextmanager
+from queue import Queue
 
 import pytest
 from fsspec.asyn import sync
@@ -111,6 +112,37 @@ async def process(row):
         list(mapper.iterate(timeout=4))
 
 
+@pytest.mark.parametrize("create_mapper", [AsyncMapper, OrderedMapper])
+def test_mapper_deadlock(create_mapper):
+    queue = Queue()
+    inputs = range(50)
+
+    def as_iter(queue):
+        while (item := queue.get()) is not None:
+            yield item
+
+    async def process(x):
+        return x
+
+    mapper = create_mapper(process, as_iter(queue), workers=10, loop=get_loop())
+    it = mapper.iterate(timeout=4)
+    for i in inputs:
+        queue.put(i)
+
+    # Check that we can get as many objects out as we put in, without deadlock
+    result = []
+    for _ in range(len(inputs)):
+        result.append(next(it))
+    if mapper.order_preserving:
+        assert result == list(inputs)
+    else:
+        assert set(result) == set(inputs)
+
+    # Check that iteration terminates cleanly
+    queue.put(None)
+    assert list(it) == []
+
+
 @pytest.mark.parametrize("create_mapper", [AsyncMapper, OrderedMapper])
 @settings(deadline=None)
 @given(