diff --git a/wgpu/backends/wgpu_native/_helpers.py b/wgpu/backends/wgpu_native/_helpers.py index 6a77fc32..23283e44 100644 --- a/wgpu/backends/wgpu_native/_helpers.py +++ b/wgpu/backends/wgpu_native/_helpers.py @@ -3,8 +3,10 @@ import ctypes import sys import time +from collections.abc import Generator from queue import deque +import anyio import sniffio from ._ffi import ffi, lib, lib_path @@ -247,22 +249,21 @@ def __init__(self, title, callback, finalizer, poll_function=None): self.callback = callback # only used to prevent it from being gc'd self.finalizer = finalizer # function to finish the result self.poll_function = poll_function # call this to poll wgpu + self.event = None # only used in asynchronous cases self.result = None def set_result(self, result): self.result = (result, None) + if self.event: + self.event.set() def set_error(self, error): self.result = (None, error) - - def _is_done(self): - self.poll_function() - return self.result is not None + if self.event: + self.event.set() def _finish(self): try: - if not self.result: - raise RuntimeError(f"Waiting for {self.title} timed out.") result, error = self.result if error: raise RuntimeError(error) @@ -278,28 +279,46 @@ def sync_wait(self): elif not self.poll_function: raise RuntimeError("Expected callback to have already happened") else: - while not self._is_done(): - time.sleep(0) + backoff_time_generator = self._get_backoff_time_generator() + while True: + self.poll_function() + if self.result is not None: + break + time.sleep(next(backoff_time_generator)) + # We check the result after sleeping just in case another thread + # causes the callback to happen + if self.result is not None: + break + return self._finish() def __await__(self): # There is no documentation on what __await__() is supposed to return, but we # can certainly copy from a function that *does* know what to return async def wait_for_callback(): - # In all the async cases that I've tried, the result is either already set, or - # resolves after the first call to the poll function. To make sure that our - # sleep-logic actually works, we always do at least one sleep call. - await async_sleep(0) + # Set self.event before checking self.result to prevent data race. + self.event = anyio.Event() if self.result is not None: return if not self.poll_function: raise RuntimeError("Expected callback to have already happened") - while not self._is_done(): - await async_sleep(0) + sleep_generator = self._get_backoff_time_generator() + while not self.event.is_set(): + self.poll_function() + with anyio.move_on_after(next(sleep_generator)): + await self.event.wait() yield from wait_for_callback().__await__() return self._finish() + def _get_backoff_time_generator(self) -> Generator[float, None, None]: + for _ in range(5): + yield 0 + for i in range(1, 20): + yield i / 2000.0 # ramp up from 0ms to 10ms + while True: + yield 0.01 + class ErrorHandler: """Object that logs errors, with the option to collect incoming