current_thread_executor.py 4.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118
  1. import threading
  2. from collections import deque
  3. from collections.abc import Callable
  4. from concurrent.futures import Executor, Future
  5. from typing import Any, ParamSpec, TypeVar
  6. _T = TypeVar("_T")
  7. _P = ParamSpec("_P")
  8. _R = TypeVar("_R")
  9. class _WorkItem:
  10. """
  11. Represents an item needing to be run in the executor.
  12. Copied from ThreadPoolExecutor (but it's private, so we're not going to rely on importing it)
  13. """
  14. def __init__(
  15. self,
  16. future: "Future[_R]",
  17. fn: Callable[_P, _R],
  18. *args: _P.args,
  19. **kwargs: _P.kwargs,
  20. ):
  21. self.future = future
  22. self.fn = fn
  23. self.args = args
  24. self.kwargs = kwargs
  25. def run(self) -> None:
  26. __traceback_hide__ = True # noqa: F841
  27. if not self.future.set_running_or_notify_cancel():
  28. return
  29. try:
  30. result = self.fn(*self.args, **self.kwargs)
  31. except BaseException as exc:
  32. self.future.set_exception(exc)
  33. # Break a reference cycle with the exception 'exc'
  34. self = None # type: ignore[assignment]
  35. else:
  36. self.future.set_result(result)
  37. class CurrentThreadExecutor(Executor):
  38. """
  39. An Executor that actually runs code in the thread it is instantiated in.
  40. Passed to other threads running async code, so they can run sync code in
  41. the thread they came from.
  42. """
  43. def __init__(self, old_executor: "CurrentThreadExecutor | None") -> None:
  44. self._work_thread = threading.current_thread()
  45. self._work_ready = threading.Condition(threading.Lock())
  46. self._work_items = deque[_WorkItem]() # synchronized by _work_ready
  47. self._broken = False # synchronized by _work_ready
  48. self._old_executor = old_executor
  49. def run_until_future(self, future: "Future[Any]") -> None:
  50. """
  51. Runs the code in the work queue until a result is available from the future.
  52. Should be run from the thread the executor is initialised in.
  53. """
  54. # Check we're in the right thread
  55. if threading.current_thread() != self._work_thread:
  56. raise RuntimeError(
  57. "You cannot run CurrentThreadExecutor from a different thread"
  58. )
  59. def done(future: "Future[Any]") -> None:
  60. with self._work_ready:
  61. self._broken = True
  62. self._work_ready.notify()
  63. future.add_done_callback(done)
  64. # Keep getting and running work items until the future we're waiting for
  65. # is done and the queue is empty.
  66. while True:
  67. with self._work_ready:
  68. while not self._work_items and not self._broken:
  69. self._work_ready.wait()
  70. if not self._work_items:
  71. break
  72. # Get a work item and run it
  73. work_item = self._work_items.popleft()
  74. work_item.run()
  75. del work_item
  76. def submit(
  77. self,
  78. fn: Callable[_P, _R],
  79. /,
  80. *args: _P.args,
  81. **kwargs: _P.kwargs,
  82. ) -> "Future[_R]":
  83. # Check they're not submitting from the same thread
  84. if threading.current_thread() == self._work_thread:
  85. raise RuntimeError(
  86. "You cannot submit onto CurrentThreadExecutor from its own thread"
  87. )
  88. f: "Future[_R]" = Future()
  89. work_item = _WorkItem(f, fn, *args, **kwargs)
  90. # Walk up the CurrentThreadExecutor stack to find the closest one still
  91. # running
  92. executor = self
  93. while True:
  94. with executor._work_ready:
  95. if not executor._broken:
  96. # Add to work queue
  97. executor._work_items.append(work_item)
  98. executor._work_ready.notify()
  99. break
  100. if executor._old_executor is None:
  101. raise RuntimeError("CurrentThreadExecutor already quit or is broken")
  102. executor = executor._old_executor
  103. # Return the future
  104. return f