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.

306 lines
10 KiB

import asyncio
import functools
import logging
import os
import platform
import signal
import socket
import sys
import threading
import time
from email.utils import formatdate
from types import FrameType
from typing import TYPE_CHECKING, List, Optional, Set, Tuple, Union
import click
from uvicorn.config import Config
if TYPE_CHECKING:
from uvicorn.protocols.http.h11_impl import H11Protocol
from uvicorn.protocols.http.httptools_impl import HttpToolsProtocol
from uvicorn.protocols.websockets.websockets_impl import WebSocketProtocol
from uvicorn.protocols.websockets.wsproto_impl import WSProtocol
Protocols = Union[H11Protocol, HttpToolsProtocol, WSProtocol, WebSocketProtocol]
HANDLED_SIGNALS = (
signal.SIGINT, # Unix signal 2. Sent by Ctrl+C.
signal.SIGTERM, # Unix signal 15. Sent by `kill <pid>`.
)
logger = logging.getLogger("uvicorn.error")
class ServerState:
"""
Shared servers state that is available between all protocol instances.
"""
def __init__(self) -> None:
self.total_requests = 0
self.connections: Set["Protocols"] = set()
self.tasks: Set[asyncio.Task] = set()
self.default_headers: List[Tuple[bytes, bytes]] = []
class Server:
def __init__(self, config: Config) -> None:
self.config = config
self.server_state = ServerState()
self.started = False
self.should_exit = False
self.force_exit = False
self.last_notified = 0.0
def run(self, sockets: Optional[List[socket.socket]] = None) -> None:
self.config.setup_event_loop()
return asyncio.run(self.serve(sockets=sockets))
async def serve(self, sockets: Optional[List[socket.socket]] = None) -> None:
process_id = os.getpid()
config = self.config
if not config.loaded:
config.load()
self.lifespan = config.lifespan_class(config)
self.install_signal_handlers()
message = "Started server process [%d]"
color_message = "Started server process [" + click.style("%d", fg="cyan") + "]"
logger.info(message, process_id, extra={"color_message": color_message})
await self.startup(sockets=sockets)
if self.should_exit:
return
await self.main_loop()
await self.shutdown(sockets=sockets)
message = "Finished server process [%d]"
color_message = "Finished server process [" + click.style("%d", fg="cyan") + "]"
logger.info(message, process_id, extra={"color_message": color_message})
async def startup(self, sockets: list = None) -> None:
await self.lifespan.startup()
if self.lifespan.should_exit:
self.should_exit = True
return
config = self.config
create_protocol = functools.partial(
config.http_protocol_class, config=config, server_state=self.server_state
)
loop = asyncio.get_running_loop()
if sockets is not None:
# Explicitly passed a list of open sockets.
# We use this when the server is run from a Gunicorn worker.
def _share_socket(sock: socket.SocketType) -> socket.SocketType:
# Windows requires the socket be explicitly shared across
# multiple workers (processes).
from socket import fromshare # type: ignore
sock_data = sock.share(os.getpid()) # type: ignore
return fromshare(sock_data)
self.servers = []
for sock in sockets:
if config.workers > 1 and platform.system() == "Windows":
sock = _share_socket(sock)
server = await loop.create_server(
create_protocol, sock=sock, ssl=config.ssl, backlog=config.backlog
)
self.servers.append(server)
listeners = sockets
elif config.fd is not None:
# Use an existing socket, from a file descriptor.
sock = socket.fromfd(config.fd, socket.AF_UNIX, socket.SOCK_STREAM)
server = await loop.create_server(
create_protocol, sock=sock, ssl=config.ssl, backlog=config.backlog
)
assert server.sockets is not None # mypy
listeners = server.sockets
self.servers = [server]
elif config.uds is not None:
# Create a socket using UNIX domain socket.
uds_perms = 0o666
if os.path.exists(config.uds):
uds_perms = os.stat(config.uds).st_mode
server = await loop.create_unix_server(
create_protocol, path=config.uds, ssl=config.ssl, backlog=config.backlog
)
os.chmod(config.uds, uds_perms)
assert server.sockets is not None # mypy
listeners = server.sockets
self.servers = [server]
else:
# Standard case. Create a socket from a host/port pair.
try:
server = await loop.create_server(
create_protocol,
host=config.host,
port=config.port,
ssl=config.ssl,
backlog=config.backlog,
)
except OSError as exc:
logger.error(exc)
await self.lifespan.shutdown()
sys.exit(1)
assert server.sockets is not None
listeners = server.sockets
self.servers = [server]
if sockets is None:
self._log_started_message(listeners)
else:
# We're most likely running multiple workers, so a message has already been
# logged by `config.bind_socket()`.
pass
self.started = True
def _log_started_message(self, listeners: List[socket.SocketType]) -> None:
config = self.config
if config.fd is not None:
sock = listeners[0]
logger.info(
"Uvicorn running on socket %s (Press CTRL+C to quit)",
sock.getsockname(),
)
elif config.uds is not None:
logger.info(
"Uvicorn running on unix socket %s (Press CTRL+C to quit)", config.uds
)
else:
addr_format = "%s://%s:%d"
host = "0.0.0.0" if config.host is None else config.host
if ":" in host:
# It's an IPv6 address.
addr_format = "%s://[%s]:%d"
port = config.port
if port == 0:
port = listeners[0].getsockname()[1]
protocol_name = "https" if config.ssl else "http"
message = f"Uvicorn running on {addr_format} (Press CTRL+C to quit)"
color_message = (
"Uvicorn running on "
+ click.style(addr_format, bold=True)
+ " (Press CTRL+C to quit)"
)
logger.info(
message,
protocol_name,
host,
port,
extra={"color_message": color_message},
)
async def main_loop(self) -> None:
counter = 0
should_exit = await self.on_tick(counter)
while not should_exit:
counter += 1
counter = counter % 864000
await asyncio.sleep(0.1)
should_exit = await self.on_tick(counter)
async def on_tick(self, counter: int) -> bool:
# Update the default headers, once per second.
if counter % 10 == 0:
current_time = time.time()
current_date = formatdate(current_time, usegmt=True).encode()
if self.config.date_header:
date_header = [(b"date", current_date)]
else:
date_header = []
self.server_state.default_headers = (
date_header + self.config.encoded_headers
)
# Callback to `callback_notify` once every `timeout_notify` seconds.
if self.config.callback_notify is not None:
if current_time - self.last_notified > self.config.timeout_notify:
self.last_notified = current_time
await self.config.callback_notify()
# Determine if we should exit.
if self.should_exit:
return True
if self.config.limit_max_requests is not None:
return self.server_state.total_requests >= self.config.limit_max_requests
return False
async def shutdown(self, sockets: Optional[List[socket.socket]] = None) -> None:
logger.info("Shutting down")
# Stop accepting new connections.
for server in self.servers:
server.close()
for sock in sockets or []:
sock.close()
for server in self.servers:
await server.wait_closed()
# Request shutdown on all existing connections.
for connection in list(self.server_state.connections):
connection.shutdown()
await asyncio.sleep(0.1)
# Wait for existing connections to finish sending responses.
if self.server_state.connections and not self.force_exit:
msg = "Waiting for connections to close. (CTRL+C to force quit)"
logger.info(msg)
while self.server_state.connections and not self.force_exit:
await asyncio.sleep(0.1)
# Wait for existing tasks to complete.
if self.server_state.tasks and not self.force_exit:
msg = "Waiting for background tasks to complete. (CTRL+C to force quit)"
logger.info(msg)
while self.server_state.tasks and not self.force_exit:
await asyncio.sleep(0.1)
# Send the lifespan shutdown event, and wait for application shutdown.
if not self.force_exit:
await self.lifespan.shutdown()
def install_signal_handlers(self) -> None:
if threading.current_thread() is not threading.main_thread():
# Signals can only be listened to from the main thread.
return
loop = asyncio.get_event_loop()
try:
for sig in HANDLED_SIGNALS:
loop.add_signal_handler(sig, self.handle_exit, sig, None)
except NotImplementedError: # pragma: no cover
# Windows
for sig in HANDLED_SIGNALS:
signal.signal(sig, self.handle_exit)
def handle_exit(self, sig: signal.Signals, frame: FrameType) -> None:
if self.should_exit and sig == signal.SIGINT:
self.force_exit = True
else:
self.should_exit = True