You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

current_thread_executor.py 2.9KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586
  1. import queue
  2. import threading
  3. import time
  4. from concurrent.futures import Executor, Future
  5. class _WorkItem(object):
  6. """
  7. Represents an item needing to be run in the executor.
  8. Copied from ThreadPoolExecutor (but it's private, so we're not going to rely on importing it)
  9. """
  10. def __init__(self, future, fn, args, kwargs):
  11. self.future = future
  12. self.fn = fn
  13. self.args = args
  14. self.kwargs = kwargs
  15. def run(self):
  16. if not self.future.set_running_or_notify_cancel():
  17. return
  18. try:
  19. result = self.fn(*self.args, **self.kwargs)
  20. except BaseException as exc:
  21. self.future.set_exception(exc)
  22. # Break a reference cycle with the exception 'exc'
  23. self = None
  24. else:
  25. self.future.set_result(result)
  26. class CurrentThreadExecutor(Executor):
  27. """
  28. An Executor that actually runs code in the thread it is instantiated in.
  29. Passed to other threads running async code, so they can run sync code in
  30. the thread they came from.
  31. """
  32. def __init__(self):
  33. self._work_thread = threading.current_thread()
  34. self._work_queue = queue.Queue()
  35. self._broken = False
  36. def run_until_future(self, future):
  37. """
  38. Runs the code in the work queue until a result is available from the future.
  39. Should be run from the thread the executor is initialised in.
  40. """
  41. # Check we're in the right thread
  42. if threading.current_thread() != self._work_thread:
  43. raise RuntimeError(
  44. "You cannot run CurrentThreadExecutor from a different thread"
  45. )
  46. # Keep getting work items and checking the future
  47. try:
  48. while True:
  49. # Get a work item and run it
  50. try:
  51. work_item = self._work_queue.get(block=False)
  52. except queue.Empty:
  53. # See if the future is done (we only exit if the work queue is empty)
  54. if future.done():
  55. return
  56. # Prevent hot-looping on nothing
  57. time.sleep(0.001)
  58. else:
  59. work_item.run()
  60. del work_item
  61. finally:
  62. self._broken = True
  63. def submit(self, fn, *args, **kwargs):
  64. # Check they're not submitting from the same thread
  65. if threading.current_thread() == self._work_thread:
  66. raise RuntimeError(
  67. "You cannot submit onto CurrentThreadExecutor from its own thread"
  68. )
  69. # Check they're not too late or the executor errored
  70. if self._broken:
  71. raise RuntimeError("CurrentThreadExecutor already quit or is broken")
  72. # Add to work queue
  73. f = Future()
  74. work_item = _WorkItem(f, fn, args, kwargs)
  75. self._work_queue.put(work_item)
  76. # Return the future
  77. return f