[v3] Sync with futures (#1804)

* fix: return Array from resize and update_attributes instead of AsyncArray

* test(sync): add tests for sync module

* clear up wait usage

* _get_loop must return a loop

* chore: clean up type hints

* feat: add timeout to sync and sync config class, and add a test

* chore: reword docstring

* chore: adjust line length for the linter

* update after v3 reorg merge

* improve tests using asyncmock

---------

Co-authored-by: Joseph Hamman <[email protected]>

Author: d-v-b

src/zarr/array.py

@@ -527,15 +527,19 @@ def __setitem__(self, selection: Selection, value: np.ndarray) -> None:
         )
 
     def resize(self, new_shape: ChunkCoords) -> Array:
-        return sync(
-            self._async_array.resize(new_shape),
-            self._async_array.runtime_configuration.asyncio_loop,
+        return type(self)(
+            sync(
+                self._async_array.resize(new_shape),
+                self._async_array.runtime_configuration.asyncio_loop,
+            )
         )
 
     def update_attributes(self, new_attributes: Dict[str, Any]) -> Array:
-        return sync(
-            self._async_array.update_attributes(new_attributes),
-            self._async_array.runtime_configuration.asyncio_loop,
+        return type(self)(
+            sync(
+                self._async_array.update_attributes(new_attributes),
+                self._async_array.runtime_configuration.asyncio_loop,
+            )
         )
 
     def __repr__(self):

src/zarr/config.py

@@ -9,6 +9,7 @@
 class SyncConfiguration:
     concurrency: Optional[int] = None
     asyncio_loop: Optional[AbstractEventLoop] = None
+    timeout: float | None = None
 
 
 def parse_indexing_order(data: Any) -> Literal["C", "F"]:

src/zarr/group.py

@@ -415,24 +415,36 @@ def nchildren(self) -> int:
 
     @property
     def children(self) -> List[Union[Array, Group]]:
-        _children = self._sync_iter(self._async_group.children())
-        return [Array(obj) if isinstance(obj, AsyncArray) else Group(obj) for obj in _children]
+        raise NotImplementedError
+        # Uncomment with AsyncGroup implements this method
+        # _children: List[Union[AsyncArray, AsyncGroup]] = self._sync_iter(
+        #     self._async_group.children()
+        # )
+        # return [Array(obj) if isinstance(obj, AsyncArray) else Group(obj) for obj in _children]
 
     def __contains__(self, child) -> bool:
         return self._sync(self._async_group.contains(child))
 
     def group_keys(self) -> List[str]:
-        return self._sync_iter(self._async_group.group_keys())
+        raise NotImplementedError
+        # uncomment with AsyncGroup implements this method
+        # return self._sync_iter(self._async_group.group_keys())
 
     def groups(self) -> List[Group]:
         # TODO: in v2 this was a generator that return key: Group
-        return [Group(obj) for obj in self._sync_iter(self._async_group.groups())]
+        raise NotImplementedError
+        # uncomment with AsyncGroup implements this method
+        # return [Group(obj) for obj in self._sync_iter(self._async_group.groups())]
 
     def array_keys(self) -> List[str]:
-        return self._sync_iter(self._async_group.array_keys())
+        # uncomment with AsyncGroup implements this method
+        # return self._sync_iter(self._async_group.array_keys())
+        raise NotImplementedError
 
     def arrays(self) -> List[Array]:
-        return [Array(obj) for obj in self._sync_iter(self._async_group.arrays())]
+        raise NotImplementedError
+        # uncomment with AsyncGroup implements this method
+        # return [Array(obj) for obj in self._sync_iter(self._async_group.arrays())]
 
     def tree(self, expand=False, level=None) -> Any:
         return self._sync(self._async_group.tree(expand=expand, level=level))

src/zarr/sync.py

@@ -1,30 +1,34 @@
 from __future__ import annotations
+from typing import TYPE_CHECKING, TypeVar
+
+if TYPE_CHECKING:
+    from typing import Any, AsyncIterator, Coroutine
 
 import asyncio
+from concurrent.futures import wait
 import threading
-from typing import (
-    Any,
-    AsyncIterator,
-    Coroutine,
-    List,
-    Optional,
-    TypeVar,
-)
+
 from typing_extensions import ParamSpec
 
 from zarr.config import SyncConfiguration
 
+P = ParamSpec("P")
+T = TypeVar("T")
 
 # From https://github.com/fsspec/filesystem_spec/blob/master/fsspec/asyn.py
 
-iothread: List[Optional[threading.Thread]] = [None]  # dedicated IO thread
-loop: List[Optional[asyncio.AbstractEventLoop]] = [
+iothread: list[threading.Thread | None] = [None]  # dedicated IO thread
+loop: list[asyncio.AbstractEventLoop | None] = [
     None
 ]  # global event loop for any non-async instance
-_lock: Optional[threading.Lock] = None  # global lock placeholder
+_lock: threading.Lock | None = None  # global lock placeholder
 get_running_loop = asyncio.get_running_loop
 
 
+class SyncError(Exception):
+    pass
+
+
 def _get_lock() -> threading.Lock:
     """Allocate or return a threading lock.
 
@@ -36,16 +40,22 @@ def _get_lock() -> threading.Lock:
     return _lock
 
 
-async def _runner(event: threading.Event, coro: Coroutine, result_box: List[Optional[Any]]):
+async def _runner(coro: Coroutine[Any, Any, T]) -> T | BaseException:
+    """
+    Await a coroutine and return the result of running it. If awaiting the coroutine raises an
+    exception, the exception will be returned.
+    """
     try:
-        result_box[0] = await coro
+        return await coro
     except Exception as ex:
-        result_box[0] = ex
-    finally:
-        event.set()
+        return ex
 
 
-def sync(coro: Coroutine, loop: Optional[asyncio.AbstractEventLoop] = None):
+def sync(
+    coro: Coroutine[Any, Any, T],
+    loop: asyncio.AbstractEventLoop | None = None,
+    timeout: float | None = None,
+) -> T:
     """
     Make loop run coroutine until it returns. Runs in other thread
 
@@ -57,30 +67,32 @@ def sync(coro: Coroutine, loop: Optional[asyncio.AbstractEventLoop] = None):
         # NB: if the loop is not running *yet*, it is OK to submit work
         # and we will wait for it
         loop = _get_loop()
-    if loop is None or loop.is_closed():
+    if not isinstance(loop, asyncio.AbstractEventLoop):
+        raise TypeError(f"loop cannot be of type {type(loop)}")
+    if loop.is_closed():
         raise RuntimeError("Loop is not running")
     try:
         loop0 = asyncio.events.get_running_loop()
         if loop0 is loop:
-            raise NotImplementedError("Calling sync() from within a running loop")
+            raise SyncError("Calling sync() from within a running loop")
     except RuntimeError:
         pass
-    result_box: List[Optional[Any]] = [None]
-    event = threading.Event()
-    asyncio.run_coroutine_threadsafe(_runner(event, coro, result_box), loop)
-    while True:
-        # this loops allows thread to get interrupted
-        if event.wait(1):
-            break
-
-    return_result = result_box[0]
+
+    future = asyncio.run_coroutine_threadsafe(_runner(coro), loop)
+
+    finished, unfinished = wait([future], return_when=asyncio.ALL_COMPLETED, timeout=timeout)
+    if len(unfinished) > 0:
+        raise asyncio.TimeoutError(f"Coroutine {coro} failed to finish in within {timeout}s")
+    assert len(finished) == 1
+    return_result = list(finished)[0].result()
+
     if isinstance(return_result, BaseException):
         raise return_result
     else:
         return return_result
 
 
-def _get_loop():
+def _get_loop() -> asyncio.AbstractEventLoop:
     """Create or return the default fsspec IO loop
 
     The loop will be running on a separate thread.
@@ -96,25 +108,24 @@ def _get_loop():
                 th.daemon = True
                 th.start()
                 iothread[0] = th
+    assert loop[0] is not None
     return loop[0]
 
 
-P = ParamSpec("P")
-T = TypeVar("T")
-
-
 class SyncMixin:
     _sync_configuration: SyncConfiguration
 
     def _sync(self, coroutine: Coroutine[Any, Any, T]) -> T:
         # TODO: refactor this to to take *args and **kwargs and pass those to the method
         # this should allow us to better type the sync wrapper
-        return sync(coroutine, loop=self._sync_configuration.asyncio_loop)
-
-    def _sync_iter(self, coroutine: Coroutine[Any, Any, AsyncIterator[T]]) -> List[T]:
-        async def iter_to_list() -> List[T]:
-            # TODO: replace with generators so we don't materialize the entire iterator at once
-            async_iterator = await coroutine
+        return sync(
+            coroutine,
+            loop=self._sync_configuration.asyncio_loop,
+            timeout=self._sync_configuration.timeout,
+        )
+
+    def _sync_iter(self, async_iterator: AsyncIterator[T]) -> list[T]:
+        async def iter_to_list() -> list[T]:
             return [item async for item in async_iterator]
 
         return self._sync(iter_to_list())

src/zarr/v3/sync.py

@@ -0,0 +1,131 @@
+from __future__ import annotations
+from typing import TYPE_CHECKING, TypeVar
+
+if TYPE_CHECKING:
+    from typing import Any, AsyncIterator, Coroutine
+
+import asyncio
+from concurrent.futures import wait
+import threading
+
+from typing_extensions import ParamSpec
+
+from zarr.v3.config import SyncConfiguration
+
+P = ParamSpec("P")
+T = TypeVar("T")
+
+# From https://github.com/fsspec/filesystem_spec/blob/master/fsspec/asyn.py
+
+iothread: list[threading.Thread | None] = [None]  # dedicated IO thread
+loop: list[asyncio.AbstractEventLoop | None] = [
+    None
+]  # global event loop for any non-async instance
+_lock: threading.Lock | None = None  # global lock placeholder
+get_running_loop = asyncio.get_running_loop
+
+
+class SyncError(Exception):
+    pass
+
+
+def _get_lock() -> threading.Lock:
+    """Allocate or return a threading lock.
+
+    The lock is allocated on first use to allow setting one lock per forked process.
+    """
+    global _lock
+    if not _lock:
+        _lock = threading.Lock()
+    return _lock
+
+
+async def _runner(coro: Coroutine[Any, Any, T]) -> T | BaseException:
+    """
+    Await a coroutine and return the result of running it. If awaiting the coroutine raises an
+    exception, the exception will be returned.
+    """
+    try:
+        return await coro
+    except Exception as ex:
+        return ex
+
+
+def sync(
+    coro: Coroutine[Any, Any, T],
+    loop: asyncio.AbstractEventLoop | None = None,
+    timeout: float | None = None,
+) -> T:
+    """
+    Make loop run coroutine until it returns. Runs in other thread
+
+    Examples
+    --------
+    >>> sync(async_function(), existing_loop)
+    """
+    if loop is None:
+        # NB: if the loop is not running *yet*, it is OK to submit work
+        # and we will wait for it
+        loop = _get_loop()
+    if not isinstance(loop, asyncio.AbstractEventLoop):
+        raise TypeError(f"loop cannot be of type {type(loop)}")
+    if loop.is_closed():
+        raise RuntimeError("Loop is not running")
+    try:
+        loop0 = asyncio.events.get_running_loop()
+        if loop0 is loop:
+            raise SyncError("Calling sync() from within a running loop")
+    except RuntimeError:
+        pass
+
+    future = asyncio.run_coroutine_threadsafe(_runner(coro), loop)
+
+    finished, unfinished = wait([future], return_when=asyncio.ALL_COMPLETED, timeout=timeout)
+    if len(unfinished) > 0:
+        raise asyncio.TimeoutError(f"Coroutine {coro} failed to finish in within {timeout}s")
+    assert len(finished) == 1
+    return_result = list(finished)[0].result()
+
+    if isinstance(return_result, BaseException):
+        raise return_result
+    else:
+        return return_result
+
+
+def _get_loop() -> asyncio.AbstractEventLoop:
+    """Create or return the default fsspec IO loop
+
+    The loop will be running on a separate thread.
+    """
+    if loop[0] is None:
+        with _get_lock():
+            # repeat the check just in case the loop got filled between the
+            # previous two calls from another thread
+            if loop[0] is None:
+                new_loop = asyncio.new_event_loop()
+                loop[0] = new_loop
+                th = threading.Thread(target=new_loop.run_forever, name="zarrIO")
+                th.daemon = True
+                th.start()
+                iothread[0] = th
+    assert loop[0] is not None
+    return loop[0]
+
+
+class SyncMixin:
+    _sync_configuration: SyncConfiguration
+
+    def _sync(self, coroutine: Coroutine[Any, Any, T]) -> T:
+        # TODO: refactor this to to take *args and **kwargs and pass those to the method
+        # this should allow us to better type the sync wrapper
+        return sync(
+            coroutine,
+            loop=self._sync_configuration.asyncio_loop,
+            timeout=self._sync_configuration.timeout,
+        )
+
+    def _sync_iter(self, async_iterator: AsyncIterator[T]) -> list[T]:
+        async def iter_to_list() -> list[T]:
+            return [item async for item in async_iterator]
+
+        return self._sync(iter_to_list())