Replace asgi "extension" with RequestResponseTiming class

Author: immerrr

uvicorn/logging.py

@@ -130,8 +130,9 @@ class GunicornSafeAtoms(abc.Mapping):
     - escapes double quotes found in atom strings
     """
 
-    def __init__(self, scope):
+    def __init__(self, scope, timing):
         self.scope = scope
+        self.timing = timing
         self.status_code = None
         self.response_headers = {}
         self.response_length = 0
@@ -142,20 +143,21 @@ def __init__(self, scope):
     def request_headers(self):
         if self._request_headers is None:
             self._request_headers = {
-                k.decode("ascii"): v.decode("ascii") for k, v in self.scope["headers"]
+                k.decode("ascii"): v.decode("ascii")
+                for k, v in self.scope["headers"]
             }
         return self._request_headers
 
     @property
     def duration(self):
-        d = self.scope["response_end_time"] - self.scope["request_start_time"]
-        return d
+        return self.timing.total_duration_seconds()
 
     def on_asgi_message(self, message):
         if message["type"] == "http.response.start":
             self.status_code = message["status"]
             self.response_headers = {
-                k.decode("ascii"): v.decode("ascii") for k, v in message["headers"]
+                k.decode("ascii"): v.decode("ascii")
+                for k, v in message["headers"]
             }
         elif message["type"] == "http.response.body":
             self.response_length += len(message.get("body", ""))

uvicorn/protocols/http/h11_impl.py

@@ -22,6 +22,7 @@
     get_path_with_query_string,
     get_remote_addr,
     is_ssl,
+    RequestResponseTiming,
 )
 
 
@@ -174,11 +175,6 @@ def handle_events(self):
                     "raw_path": raw_path,
                     "query_string": query_string,
                     "headers": self.headers,
-                    "extensions": {
-                        "uvicorn_request_duration": {
-                            "request_start_time": time.monotonic(),
-                        }
-                    },
                 }
 
                 for name, value in self.headers:
@@ -212,6 +208,7 @@ def handle_events(self):
                     message_event=asyncio.Event(),
                     on_response=self.on_response_complete,
                 )
+                self.cycle.timing.request_started()
                 task = self.loop.create_task(self.cycle.run_asgi(app))
                 task.add_done_callback(self.tasks.discard)
                 self.tasks.add(task)
@@ -229,6 +226,7 @@ def handle_events(self):
                     self.transport.resume_reading()
                     self.conn.start_next_cycle()
                     continue
+                self.cycle.timing.request_ended()
                 self.cycle.more_body = False
                 self.cycle.message_event.set()
 
@@ -373,10 +371,11 @@ def __init__(
         # Response state
         self.response_started = False
         self.response_complete = False
+        self.timing = RequestResponseTiming()
 
         # For logging
         if self.gunicorn_log:
-            self.gunicorn_atoms = GunicornSafeAtoms(self.scope)
+            self.gunicorn_atoms = GunicornSafeAtoms(self.scope, self.timing)
         else:
             self.gunicorn_atoms = None
 
@@ -443,6 +442,7 @@ async def send(self, message):
 
             self.response_started = True
             self.waiting_for_100_continue = False
+            self.timing.response_started()
 
             status_code = message["status"]
             headers = self.default_headers + message.get("headers", [])
@@ -492,8 +492,7 @@ async def send(self, message):
                 event = h11.EndOfMessage()
                 output = self.conn.send(event)
 
-                duration_scope = self.scope["extensions"]["uvicorn_request_duration"]
-                duration_scope["response_end_time"] = time.monotonic()
+                self.timing.response_ended()
                 if self.gunicorn_log is not None:
                     try:
                         self.gunicorn_log.access_log.info(

uvicorn/protocols/http/httptools_impl.py

@@ -23,6 +23,7 @@
     get_path_with_query_string,
     get_remote_addr,
     is_ssl,
+    RequestResponseTiming,
 )
 
 HEADER_RE = re.compile(b'[\x00-\x1F\x7F()<>@,;:[]={} \t\\"]')
@@ -86,6 +87,7 @@ def __init__(
         self.headers = None
         self.expect_100_continue = False
         self.cycle = None
+        self.request_start_time = None
 
     # Protocol interface
     def connection_made(self, transport):
@@ -211,13 +213,11 @@ def on_url(self, url):
             "raw_path": raw_path,
             "query_string": parsed_url.query if parsed_url.query else b"",
             "headers": self.headers,
-            "extensions": {
-                "uvicorn_request_duration": {
-                    "request_start_time": time.monotonic(),
-                }
-            },
         }
 
+    def on_message_begin(self):
+        self.request_start_time = time.monotonic()
+
     def on_header(self, name: bytes, value: bytes):
         name = name.lower()
         if name == b"expect" and value.lower() == b"100-continue":
@@ -257,6 +257,7 @@ def on_headers_complete(self):
             keep_alive=http_version != "1.0",
             on_response=self.on_response_complete,
         )
+        self.cycle.timing.request_start_time = self.request_start_time
         if existing_cycle is None or existing_cycle.response_complete:
             # Standard case - start processing the request.
             task = self.loop.create_task(self.cycle.run_asgi(app))
@@ -280,6 +281,7 @@ def on_message_complete(self):
             return
         self.cycle.more_body = False
         self.cycle.message_event.set()
+        self.cycle.timing.request_ended()
 
     def on_response_complete(self):
         # Callback for pipelined HTTP requests to be started.
@@ -376,10 +378,11 @@ def __init__(
         self.response_complete = False
         self.chunked_encoding = None
         self.expected_content_length = 0
+        self.timing = RequestResponseTiming()
 
         # For logging.
         if self.gunicorn_log:
-            self.gunicorn_atoms = GunicornSafeAtoms(self.scope)
+            self.gunicorn_atoms = GunicornSafeAtoms(self.scope, self.timing)
         else:
             self.gunicorn_atoms = None
 
@@ -445,6 +448,7 @@ async def send(self, message):
                 raise RuntimeError(msg % message_type)
 
             self.response_started = True
+            self.timing.response_started()
             self.waiting_for_100_continue = False
 
             status_code = message["status"]
@@ -528,8 +532,7 @@ async def send(self, message):
                 if self.expected_content_length != 0:
                     raise RuntimeError("Response content shorter than Content-Length")
                 self.response_complete = True
-                duration_scope = self.scope["extensions"]["uvicorn_request_duration"]
-                duration_scope["response_end_time"] = time.monotonic()
+                self.timing.response_ended()
 
                 if self.gunicorn_log is not None:
                     try:

uvicorn/protocols/utils.py

@@ -1,5 +1,6 @@
 import asyncio
 import urllib.parse
+import time
 from typing import Optional, Tuple
 
 from asgiref.typing import WWWScope
@@ -54,3 +55,32 @@ def get_path_with_query_string(scope: WWWScope) -> str:
             path_with_query_string, scope["query_string"].decode("ascii")
         )
     return path_with_query_string
+
+
+class RequestResponseTiming:
+    def __init__(self):
+        self.request_start_time: Optional[int] = None
+        self.request_end_time: Optional[int] = None
+        self.response_start_time: Optional[int] = None
+        self.response_end_time: Optional[int] = None
+
+    def request_started(self):
+        self.request_start_time = time.monotonic()
+
+    def request_ended(self):
+        self.request_end_time = time.monotonic()
+
+    def response_started(self):
+        self.response_start_time = time.monotonic()
+
+    def response_ended(self):
+        self.response_end_time = time.monotonic()
+
+    def request_duration_seconds(self):
+        return self.request_end_time - self.request_start_time
+
+    def response_duration_seconds(self):
+        return self.response_end_time - self.response_start_time
+
+    def total_duration_seconds(self):
+        return self.response_end_time - self.request_start_time