worker.py 8.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262
  1. """Async gunicorn worker for aiohttp.web"""
  2. import asyncio
  3. import inspect
  4. import os
  5. import re
  6. import signal
  7. import sys
  8. from types import FrameType
  9. from typing import TYPE_CHECKING, Any, Optional
  10. from gunicorn.config import AccessLogFormat as GunicornAccessLogFormat
  11. from gunicorn.workers import base
  12. from aiohttp import web
  13. from .helpers import set_result
  14. from .web_app import Application
  15. from .web_log import AccessLogger
  16. if TYPE_CHECKING:
  17. import ssl
  18. SSLContext = ssl.SSLContext
  19. else:
  20. try:
  21. import ssl
  22. SSLContext = ssl.SSLContext
  23. except ImportError: # pragma: no cover
  24. ssl = None # type: ignore[assignment]
  25. SSLContext = object # type: ignore[misc,assignment]
  26. __all__ = ("GunicornWebWorker", "GunicornUVLoopWebWorker")
  27. class GunicornWebWorker(base.Worker): # type: ignore[misc,no-any-unimported]
  28. DEFAULT_AIOHTTP_LOG_FORMAT = AccessLogger.LOG_FORMAT
  29. DEFAULT_GUNICORN_LOG_FORMAT = GunicornAccessLogFormat.default
  30. def __init__(self, *args: Any, **kw: Any) -> None: # pragma: no cover
  31. super().__init__(*args, **kw)
  32. self._task: Optional[asyncio.Task[None]] = None
  33. self.exit_code = 0
  34. self._notify_waiter: Optional[asyncio.Future[bool]] = None
  35. def init_process(self) -> None:
  36. # create new event_loop after fork
  37. try:
  38. asyncio.get_event_loop().close()
  39. except RuntimeError:
  40. # No loop was running
  41. pass
  42. self.loop = asyncio.new_event_loop()
  43. asyncio.set_event_loop(self.loop)
  44. super().init_process()
  45. def run(self) -> None:
  46. self._task = self.loop.create_task(self._run())
  47. try: # ignore all finalization problems
  48. self.loop.run_until_complete(self._task)
  49. except Exception:
  50. self.log.exception("Exception in gunicorn worker")
  51. self.loop.run_until_complete(self.loop.shutdown_asyncgens())
  52. self.loop.close()
  53. sys.exit(self.exit_code)
  54. async def _run(self) -> None:
  55. runner = None
  56. if isinstance(self.wsgi, Application):
  57. app = self.wsgi
  58. elif inspect.iscoroutinefunction(self.wsgi) or (
  59. sys.version_info < (3, 14) and asyncio.iscoroutinefunction(self.wsgi)
  60. ):
  61. wsgi = await self.wsgi()
  62. if isinstance(wsgi, web.AppRunner):
  63. runner = wsgi
  64. app = runner.app
  65. else:
  66. app = wsgi
  67. else:
  68. raise RuntimeError(
  69. "wsgi app should be either Application or "
  70. "async function returning Application, got {}".format(self.wsgi)
  71. )
  72. if runner is None:
  73. access_log = self.log.access_log if self.cfg.accesslog else None
  74. runner = web.AppRunner(
  75. app,
  76. logger=self.log,
  77. keepalive_timeout=self.cfg.keepalive,
  78. access_log=access_log,
  79. access_log_format=self._get_valid_log_format(
  80. self.cfg.access_log_format
  81. ),
  82. shutdown_timeout=self.cfg.graceful_timeout / 100 * 95,
  83. )
  84. await runner.setup()
  85. ctx = self._create_ssl_context(self.cfg) if self.cfg.is_ssl else None
  86. runner = runner
  87. assert runner is not None
  88. server = runner.server
  89. assert server is not None
  90. for sock in self.sockets:
  91. site = web.SockSite(
  92. runner,
  93. sock,
  94. ssl_context=ctx,
  95. )
  96. await site.start()
  97. # If our parent changed then we shut down.
  98. pid = os.getpid()
  99. try:
  100. while self.alive: # type: ignore[has-type]
  101. self.notify()
  102. cnt = server.requests_count
  103. if self.max_requests and cnt > self.max_requests:
  104. self.alive = False
  105. self.log.info("Max requests, shutting down: %s", self)
  106. elif pid == os.getpid() and self.ppid != os.getppid():
  107. self.alive = False
  108. self.log.info("Parent changed, shutting down: %s", self)
  109. else:
  110. await self._wait_next_notify()
  111. except BaseException:
  112. pass
  113. await runner.cleanup()
  114. def _wait_next_notify(self) -> "asyncio.Future[bool]":
  115. self._notify_waiter_done()
  116. loop = self.loop
  117. assert loop is not None
  118. self._notify_waiter = waiter = loop.create_future()
  119. self.loop.call_later(1.0, self._notify_waiter_done, waiter)
  120. return waiter
  121. def _notify_waiter_done(
  122. self, waiter: Optional["asyncio.Future[bool]"] = None
  123. ) -> None:
  124. if waiter is None:
  125. waiter = self._notify_waiter
  126. if waiter is not None:
  127. set_result(waiter, True)
  128. if waiter is self._notify_waiter:
  129. self._notify_waiter = None
  130. def init_signals(self) -> None:
  131. # Set up signals through the event loop API.
  132. self.loop.add_signal_handler(
  133. signal.SIGQUIT, self.handle_quit, signal.SIGQUIT, None
  134. )
  135. self.loop.add_signal_handler(
  136. signal.SIGTERM, self.handle_exit, signal.SIGTERM, None
  137. )
  138. self.loop.add_signal_handler(
  139. signal.SIGINT, self.handle_quit, signal.SIGINT, None
  140. )
  141. self.loop.add_signal_handler(
  142. signal.SIGWINCH, self.handle_winch, signal.SIGWINCH, None
  143. )
  144. self.loop.add_signal_handler(
  145. signal.SIGUSR1, self.handle_usr1, signal.SIGUSR1, None
  146. )
  147. self.loop.add_signal_handler(
  148. signal.SIGABRT, self.handle_abort, signal.SIGABRT, None
  149. )
  150. # Don't let SIGTERM and SIGUSR1 disturb active requests
  151. # by interrupting system calls
  152. signal.siginterrupt(signal.SIGTERM, False)
  153. signal.siginterrupt(signal.SIGUSR1, False)
  154. # Reset signals so Gunicorn doesn't swallow subprocess return codes
  155. # See: https://github.com/aio-libs/aiohttp/issues/6130
  156. def handle_quit(self, sig: int, frame: Optional[FrameType]) -> None:
  157. self.alive = False
  158. # worker_int callback
  159. self.cfg.worker_int(self)
  160. # wakeup closing process
  161. self._notify_waiter_done()
  162. def handle_abort(self, sig: int, frame: Optional[FrameType]) -> None:
  163. self.alive = False
  164. self.exit_code = 1
  165. self.cfg.worker_abort(self)
  166. sys.exit(1)
  167. @staticmethod
  168. def _create_ssl_context(cfg: Any) -> "SSLContext":
  169. """Creates SSLContext instance for usage in asyncio.create_server.
  170. See ssl.SSLSocket.__init__ for more details.
  171. """
  172. if ssl is None: # pragma: no cover
  173. raise RuntimeError("SSL is not supported.")
  174. ctx = ssl.SSLContext(cfg.ssl_version)
  175. ctx.load_cert_chain(cfg.certfile, cfg.keyfile)
  176. ctx.verify_mode = cfg.cert_reqs
  177. if cfg.ca_certs:
  178. ctx.load_verify_locations(cfg.ca_certs)
  179. if cfg.ciphers:
  180. ctx.set_ciphers(cfg.ciphers)
  181. return ctx
  182. def _get_valid_log_format(self, source_format: str) -> str:
  183. if source_format == self.DEFAULT_GUNICORN_LOG_FORMAT:
  184. return self.DEFAULT_AIOHTTP_LOG_FORMAT
  185. elif re.search(r"%\([^\)]+\)", source_format):
  186. raise ValueError(
  187. "Gunicorn's style options in form of `%(name)s` are not "
  188. "supported for the log formatting. Please use aiohttp's "
  189. "format specification to configure access log formatting: "
  190. "http://docs.aiohttp.org/en/stable/logging.html"
  191. "#format-specification"
  192. )
  193. else:
  194. return source_format
  195. class GunicornUVLoopWebWorker(GunicornWebWorker):
  196. def init_process(self) -> None:
  197. import uvloop
  198. # Close any existing event loop before setting a
  199. # new policy.
  200. try:
  201. asyncio.get_event_loop().close()
  202. except RuntimeError:
  203. # No loop was running
  204. pass
  205. # Setup uvloop policy, so that every
  206. # asyncio.get_event_loop() will create an instance
  207. # of uvloop event loop.
  208. asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
  209. super().init_process()