/lightbus/utilities/async_tools.py

https://github.com/adamcharnock/lightbus · Python · 180 lines · 108 code · 21 blank · 51 comment · 33 complexity · c540b1e81ee42460dc8f8e5ac7b4ccae MD5 · raw file

  1. import sys
  2. import asyncio
  3. import logging
  4. import threading
  5. from concurrent.futures.thread import ThreadPoolExecutor
  6. from time import time
  7. from typing import Coroutine, TYPE_CHECKING
  8. import datetime
  9. import aioredis
  10. from lightbus.exceptions import CannotBlockHere
  11. if TYPE_CHECKING:
  12. # pylint: disable=unused-import,cyclic-import,cyclic-import
  13. from schedule import Job
  14. logger = logging.getLogger(__name__)
  15. PYTHON_VERSION = tuple(sys.version_info)
  16. def block(coroutine: Coroutine, loop=None, *, timeout=None):
  17. """Call asynchronous code synchronously
  18. Note that this cannot be used inside an event loop.
  19. """
  20. loop = loop or get_event_loop()
  21. if loop.is_running():
  22. if hasattr(coroutine, "close"):
  23. coroutine.close()
  24. raise CannotBlockHere(
  25. "It appears you have tried to use a blocking API method "
  26. "from within an event loop. Unfortunately this is unsupported. "
  27. "Instead, use the async version of the method."
  28. )
  29. try:
  30. if timeout is None:
  31. val = loop.run_until_complete(coroutine)
  32. else:
  33. val = loop.run_until_complete(asyncio.wait_for(coroutine, timeout=timeout))
  34. except Exception as e:
  35. # The intention here is to get sensible stack traces from exceptions within blocking calls
  36. raise e
  37. return val
  38. def get_event_loop():
  39. try:
  40. loop = asyncio.get_event_loop()
  41. except RuntimeError:
  42. loop = asyncio.new_event_loop()
  43. asyncio.set_event_loop(loop)
  44. return loop
  45. async def cancel(*tasks):
  46. """Useful for cleaning up tasks in tests"""
  47. # pylint: disable=broad-except
  48. ex = None
  49. for task in tasks:
  50. if task is None:
  51. continue
  52. # Cancel all the tasks any pull out any exceptions
  53. if not task.cancelled():
  54. task.cancel()
  55. try:
  56. await task
  57. task.result()
  58. except asyncio.CancelledError:
  59. pass
  60. except Exception as e:
  61. # aioredis hides errors which occurred within a pipline
  62. # by wrapping them up in a PipelineError. So see if that is
  63. # the case here
  64. if (
  65. # An aioredis pipeline error
  66. isinstance(e, aioredis.PipelineError)
  67. # Where one of the errors that caused it was a CancelledError
  68. and len(e.args) > 1
  69. and asyncio.CancelledError in map(type, e.args[1])
  70. ):
  71. pass
  72. else:
  73. # If there was an exception, and this is the first
  74. # exception we've seen, then stash it away for later
  75. if ex is None:
  76. ex = e
  77. # Now raise the first exception we saw, if any
  78. if ex:
  79. raise ex
  80. async def cancel_and_log_exceptions(*tasks):
  81. """Cancel tasks and log any exceptions
  82. This is useful when shutting down, when tasks need to be cancelled any anything
  83. that goes wrong should be logged but will not otherwise be dealt with.
  84. """
  85. for task in tasks:
  86. try:
  87. await cancel(task)
  88. except Exception as e:
  89. # pylint: disable=broad-except
  90. logger.info(
  91. "Error encountered when shutting down task %s. Exception logged, will now move on.",
  92. task,
  93. )
  94. async def run_user_provided_callable(callable_, args, kwargs):
  95. """Run user provided code
  96. If the callable is blocking (i.e. a regular function) it will be
  97. moved to its own thread and awaited. The purpose of this thread is to:
  98. 1. Allow other tasks to continue as normal
  99. 2. Allow user code to start its own event loop where needed
  100. If an async function is provided it will be awaited as-is, and no
  101. child thread will be created.
  102. The callable will be called with the given args and kwargs
  103. """
  104. if asyncio.iscoroutinefunction(callable_):
  105. try:
  106. return await callable_(*args, **kwargs)
  107. except Exception as e:
  108. exception = e
  109. else:
  110. try:
  111. thread_pool_executor = ThreadPoolExecutor(
  112. thread_name_prefix="user_provided_callable_tpe"
  113. )
  114. future = asyncio.get_event_loop().run_in_executor(
  115. executor=thread_pool_executor, func=lambda: callable_(*args, **kwargs)
  116. )
  117. return await future
  118. except Exception as e:
  119. exception = e
  120. logger.debug(f"Error in user provided callable: {repr(exception)}")
  121. raise exception
  122. async def call_every(*, callback, timedelta: datetime.timedelta, also_run_immediately: bool):
  123. """Call callback every timedelta
  124. If also_run_immediately is set then the callback will be called before any waiting
  125. happens. If the callback's result is awaitable then it will be awaited
  126. Callback execution time is accounted for in the scheduling. If the execution takes
  127. 2 seconds, and timedelta is 10 seconds, then call_every() will wait 8 seconds
  128. before the subsequent execution.
  129. """
  130. first_run = True
  131. while True:
  132. start_time = time()
  133. if not first_run or also_run_immediately:
  134. await run_user_provided_callable(callback, args=[], kwargs={})
  135. total_execution_time = time() - start_time
  136. sleep_time = max(0.0, timedelta.total_seconds() - total_execution_time)
  137. await asyncio.sleep(sleep_time)
  138. first_run = False
  139. async def call_on_schedule(callback, schedule: "Job", also_run_immediately: bool):
  140. first_run = True
  141. while True:
  142. schedule._schedule_next_run()
  143. if not first_run or also_run_immediately:
  144. schedule.last_run = datetime.datetime.now()
  145. await run_user_provided_callable(callback, args=[], kwargs={})
  146. td = schedule.next_run - datetime.datetime.now()
  147. await asyncio.sleep(td.total_seconds())
  148. first_run = False