server.py 12 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336
  1. from __future__ import annotations
  2. import asyncio
  3. import logging
  4. import os
  5. import platform
  6. import signal
  7. import socket
  8. import sys
  9. import threading
  10. import time
  11. from email.utils import formatdate
  12. from types import FrameType
  13. from typing import TYPE_CHECKING, Sequence, Union
  14. import click
  15. from uvicorn.config import Config
  16. if TYPE_CHECKING:
  17. from uvicorn.protocols.http.h11_impl import H11Protocol
  18. from uvicorn.protocols.http.httptools_impl import HttpToolsProtocol
  19. from uvicorn.protocols.websockets.websockets_impl import WebSocketProtocol
  20. from uvicorn.protocols.websockets.wsproto_impl import WSProtocol
  21. Protocols = Union[H11Protocol, HttpToolsProtocol, WSProtocol, WebSocketProtocol]
  22. HANDLED_SIGNALS = (
  23. signal.SIGINT, # Unix signal 2. Sent by Ctrl+C.
  24. signal.SIGTERM, # Unix signal 15. Sent by `kill <pid>`.
  25. )
  26. if sys.platform == "win32": # pragma: py-not-win32
  27. HANDLED_SIGNALS += (signal.SIGBREAK,) # Windows signal 21. Sent by Ctrl+Break.
  28. logger = logging.getLogger("uvicorn.error")
  29. class ServerState:
  30. """
  31. Shared servers state that is available between all protocol instances.
  32. """
  33. def __init__(self) -> None:
  34. self.total_requests = 0
  35. self.connections: set[Protocols] = set()
  36. self.tasks: set[asyncio.Task[None]] = set()
  37. self.default_headers: list[tuple[bytes, bytes]] = []
  38. class Server:
  39. def __init__(self, config: Config) -> None:
  40. self.config = config
  41. self.server_state = ServerState()
  42. self.started = False
  43. self.should_exit = False
  44. self.force_exit = False
  45. self.last_notified = 0.0
  46. def run(self, sockets: list[socket.socket] | None = None) -> None:
  47. self.config.setup_event_loop()
  48. return asyncio.run(self.serve(sockets=sockets))
  49. async def serve(self, sockets: list[socket.socket] | None = None) -> None:
  50. process_id = os.getpid()
  51. config = self.config
  52. if not config.loaded:
  53. config.load()
  54. self.lifespan = config.lifespan_class(config)
  55. self.install_signal_handlers()
  56. message = "Started server process [%d]"
  57. color_message = "Started server process [" + click.style("%d", fg="cyan") + "]"
  58. logger.info(message, process_id, extra={"color_message": color_message})
  59. await self.startup(sockets=sockets)
  60. if self.should_exit:
  61. return
  62. await self.main_loop()
  63. await self.shutdown(sockets=sockets)
  64. message = "Finished server process [%d]"
  65. color_message = "Finished server process [" + click.style("%d", fg="cyan") + "]"
  66. logger.info(message, process_id, extra={"color_message": color_message})
  67. async def startup(self, sockets: list[socket.socket] | None = None) -> None:
  68. await self.lifespan.startup()
  69. if self.lifespan.should_exit:
  70. self.should_exit = True
  71. return
  72. config = self.config
  73. def create_protocol(
  74. _loop: asyncio.AbstractEventLoop | None = None,
  75. ) -> asyncio.Protocol:
  76. return config.http_protocol_class( # type: ignore[call-arg]
  77. config=config,
  78. server_state=self.server_state,
  79. app_state=self.lifespan.state,
  80. _loop=_loop,
  81. )
  82. loop = asyncio.get_running_loop()
  83. listeners: Sequence[socket.SocketType]
  84. if sockets is not None:
  85. # Explicitly passed a list of open sockets.
  86. # We use this when the server is run from a Gunicorn worker.
  87. def _share_socket(
  88. sock: socket.SocketType,
  89. ) -> socket.SocketType: # pragma py-linux pragma: py-darwin
  90. # Windows requires the socket be explicitly shared across
  91. # multiple workers (processes).
  92. from socket import fromshare # type: ignore[attr-defined]
  93. sock_data = sock.share(os.getpid()) # type: ignore[attr-defined]
  94. return fromshare(sock_data)
  95. self.servers: list[asyncio.base_events.Server] = []
  96. for sock in sockets:
  97. is_windows = platform.system() == "Windows"
  98. if config.workers > 1 and is_windows: # pragma: py-not-win32
  99. sock = _share_socket(sock) # type: ignore[assignment]
  100. server = await loop.create_server(
  101. create_protocol, sock=sock, ssl=config.ssl, backlog=config.backlog
  102. )
  103. self.servers.append(server)
  104. listeners = sockets
  105. elif config.fd is not None: # pragma: py-win32
  106. # Use an existing socket, from a file descriptor.
  107. sock = socket.fromfd(config.fd, socket.AF_UNIX, socket.SOCK_STREAM)
  108. server = await loop.create_server(
  109. create_protocol, sock=sock, ssl=config.ssl, backlog=config.backlog
  110. )
  111. assert server.sockets is not None # mypy
  112. listeners = server.sockets
  113. self.servers = [server]
  114. elif config.uds is not None: # pragma: py-win32
  115. # Create a socket using UNIX domain socket.
  116. uds_perms = 0o666
  117. if os.path.exists(config.uds):
  118. uds_perms = os.stat(config.uds).st_mode
  119. server = await loop.create_unix_server(
  120. create_protocol, path=config.uds, ssl=config.ssl, backlog=config.backlog
  121. )
  122. os.chmod(config.uds, uds_perms)
  123. assert server.sockets is not None # mypy
  124. listeners = server.sockets
  125. self.servers = [server]
  126. else:
  127. # Standard case. Create a socket from a host/port pair.
  128. try:
  129. server = await loop.create_server(
  130. create_protocol,
  131. host=config.host,
  132. port=config.port,
  133. ssl=config.ssl,
  134. backlog=config.backlog,
  135. )
  136. except OSError as exc:
  137. logger.error(exc)
  138. await self.lifespan.shutdown()
  139. sys.exit(1)
  140. assert server.sockets is not None
  141. listeners = server.sockets
  142. self.servers = [server]
  143. if sockets is None:
  144. self._log_started_message(listeners)
  145. else:
  146. # We're most likely running multiple workers, so a message has already been
  147. # logged by `config.bind_socket()`.
  148. pass
  149. self.started = True
  150. def _log_started_message(self, listeners: Sequence[socket.SocketType]) -> None:
  151. config = self.config
  152. if config.fd is not None: # pragma: py-win32
  153. sock = listeners[0]
  154. logger.info(
  155. "Uvicorn running on socket %s (Press CTRL+C to quit)",
  156. sock.getsockname(),
  157. )
  158. elif config.uds is not None: # pragma: py-win32
  159. logger.info(
  160. "Uvicorn running on unix socket %s (Press CTRL+C to quit)", config.uds
  161. )
  162. else:
  163. addr_format = "%s://%s:%d"
  164. host = "0.0.0.0" if config.host is None else config.host
  165. if ":" in host:
  166. # It's an IPv6 address.
  167. addr_format = "%s://[%s]:%d"
  168. port = config.port
  169. if port == 0:
  170. port = listeners[0].getsockname()[1]
  171. protocol_name = "https" if config.ssl else "http"
  172. message = f"Uvicorn running on {addr_format} (Press CTRL+C to quit)"
  173. color_message = (
  174. "Uvicorn running on "
  175. + click.style(addr_format, bold=True)
  176. + " (Press CTRL+C to quit)"
  177. )
  178. logger.info(
  179. message,
  180. protocol_name,
  181. host,
  182. port,
  183. extra={"color_message": color_message},
  184. )
  185. async def main_loop(self) -> None:
  186. counter = 0
  187. should_exit = await self.on_tick(counter)
  188. while not should_exit:
  189. counter += 1
  190. counter = counter % 864000
  191. await asyncio.sleep(0.1)
  192. should_exit = await self.on_tick(counter)
  193. async def on_tick(self, counter: int) -> bool:
  194. # Update the default headers, once per second.
  195. if counter % 10 == 0:
  196. current_time = time.time()
  197. current_date = formatdate(current_time, usegmt=True).encode()
  198. if self.config.date_header:
  199. date_header = [(b"date", current_date)]
  200. else:
  201. date_header = []
  202. self.server_state.default_headers = (
  203. date_header + self.config.encoded_headers
  204. )
  205. # Callback to `callback_notify` once every `timeout_notify` seconds.
  206. if self.config.callback_notify is not None:
  207. if current_time - self.last_notified > self.config.timeout_notify:
  208. self.last_notified = current_time
  209. await self.config.callback_notify()
  210. # Determine if we should exit.
  211. if self.should_exit:
  212. return True
  213. if self.config.limit_max_requests is not None:
  214. return self.server_state.total_requests >= self.config.limit_max_requests
  215. return False
  216. async def shutdown(self, sockets: list[socket.socket] | None = None) -> None:
  217. logger.info("Shutting down")
  218. # Stop accepting new connections.
  219. for server in self.servers:
  220. server.close()
  221. for sock in sockets or []:
  222. sock.close()
  223. # Request shutdown on all existing connections.
  224. for connection in list(self.server_state.connections):
  225. connection.shutdown()
  226. await asyncio.sleep(0.1)
  227. # When 3.10 is not supported anymore, use `async with asyncio.timeout(...):`.
  228. try:
  229. await asyncio.wait_for(
  230. self._wait_tasks_to_complete(),
  231. timeout=self.config.timeout_graceful_shutdown,
  232. )
  233. except asyncio.TimeoutError:
  234. logger.error(
  235. "Cancel %s running task(s), timeout graceful shutdown exceeded",
  236. len(self.server_state.tasks),
  237. )
  238. for t in self.server_state.tasks:
  239. if sys.version_info < (3, 9): # pragma: py-gte-39
  240. t.cancel()
  241. else: # pragma: py-lt-39
  242. t.cancel(msg="Task cancelled, timeout graceful shutdown exceeded")
  243. # Send the lifespan shutdown event, and wait for application shutdown.
  244. if not self.force_exit:
  245. await self.lifespan.shutdown()
  246. async def _wait_tasks_to_complete(self) -> None:
  247. # Wait for existing connections to finish sending responses.
  248. if self.server_state.connections and not self.force_exit:
  249. msg = "Waiting for connections to close. (CTRL+C to force quit)"
  250. logger.info(msg)
  251. while self.server_state.connections and not self.force_exit:
  252. await asyncio.sleep(0.1)
  253. # Wait for existing tasks to complete.
  254. if self.server_state.tasks and not self.force_exit:
  255. msg = "Waiting for background tasks to complete. (CTRL+C to force quit)"
  256. logger.info(msg)
  257. while self.server_state.tasks and not self.force_exit:
  258. await asyncio.sleep(0.1)
  259. for server in self.servers:
  260. await server.wait_closed()
  261. def install_signal_handlers(self) -> None:
  262. if threading.current_thread() is not threading.main_thread():
  263. # Signals can only be listened to from the main thread.
  264. return
  265. loop = asyncio.get_event_loop()
  266. try:
  267. for sig in HANDLED_SIGNALS:
  268. loop.add_signal_handler(sig, self.handle_exit, sig, None)
  269. except NotImplementedError: # pragma: no cover
  270. # Windows
  271. for sig in HANDLED_SIGNALS:
  272. signal.signal(sig, self.handle_exit)
  273. def handle_exit(self, sig: int, frame: FrameType | None) -> None:
  274. if self.should_exit and sig == signal.SIGINT:
  275. self.force_exit = True
  276. else:
  277. self.should_exit = True