Default to no timeout, add timeout= option, closes #10

Also improves returned errors, refs #6

Author: simonw

README.md

@@ -27,6 +27,8 @@ Now `app` is an ASGI application that will proxy all incoming HTTP requests to t
 
 The function takes an optional second argument, `log=` - set this to a Python logger, or any object that has `.info(msg)` and `.error(msg)` methods, and the proxy will log information about each request it proxies.
 
+It also takes a `timeout=` option, which defaults to `None` for no timeout. This can be set to a floating point value in seconds to enforce a timeout on requests that are being proxied.
+
 ## CLI tool
 
 You can try this module out like so:
@@ -38,7 +40,7 @@ You may need to `pip install uvicorn` first for this to work.
 
 This will start a server on port 8000 that proxies to `https://datasette.io`.
 
-Add `-p PORT` to specify a different port, `--verbose` to see debug logging, and `--host 127.0.0.1` to listen on a different host (the default is `0.0.0.0`).
+Add `-p PORT` to specify a different port, `--timeout 3` to set a timeout, `--verbose` to see debug logging, and `--host 127.0.0.1` to listen on a different host (the default is `0.0.0.0`).
 
 ## Development
 

asgi_proxy/__init__.py

@@ -1,8 +1,8 @@
-from httpx import AsyncClient
+import httpx
 from urllib.parse import urlparse
 
 
-def asgi_proxy(backend, log=None):
+def asgi_proxy(backend, log=None, timeout=None):
     backend_host = urlparse(backend).netloc
 
     async def asgi_proxy(scope, receive, send):
@@ -32,7 +32,7 @@ async def asgi_proxy(scope, receive, send):
             body += message.get("body", b"")
             more_body = message.get("more_body", False)
 
-        async with AsyncClient() as client:
+        async with httpx.AsyncClient(timeout=timeout) as client:
             try:
                 # Stream it, in case of long streaming responses
                 async with client.stream(
@@ -74,10 +74,39 @@ async def asgi_proxy(scope, receive, send):
                         return
 
                     await send({"type": "http.response.body", "more_body": False})
-            except Exception as e:
+            except httpx.TimeoutException as ex:
+                if log:
+                    log.error(f"Timeout error occurred: {ex.__class__.__name__}: {ex}")
+                await send(
+                    {
+                        "type": "http.response.start",
+                        "status": 504,
+                    }
+                )
+                await send(
+                    {
+                        "type": "http.response.body",
+                        "body": b"Gateway timeout",
+                        "more_body": False,
+                    }
+                )
+            except Exception as ex:
                 # Handle any errors during the request
                 if log:
-                    log.error(f"An error occurred: {e.__class__.__name__}: {e}")
-                return
+                    log.error(f"An error occurred: {ex.__class__.__name__}: {ex}")
+                await send(
+                    {
+                        "type": "http.response.start",
+                        # Generic gateway error
+                        "status": 502,
+                    }
+                )
+                await send(
+                    {
+                        "type": "http.response.body",
+                        "body": b"Bad gateway",
+                        "more_body": False,
+                    }
+                )
 
     return asgi_proxy

asgi_proxy/__main__.py

@@ -18,6 +18,9 @@
     parser.add_argument(
         "--host", default="0.0.0.0", help="Host to bind to (default: 0.0.0.0)"
     )
+    parser.add_argument(
+        "--timeout", type=float, default=None, help="Timeout in seconds"
+    )
     parser.add_argument("-v", "--verbose", action="store_true", help="Verbose logging")
 
     args = parser.parse_args()
@@ -26,8 +29,8 @@
         import logging
 
         logging.basicConfig(level=logging.INFO)
-        app = asgi_proxy(args.url, log=logging)
+        app = asgi_proxy(args.url, log=logging, timeout=args.timeout)
     else:
-        app = asgi_proxy(args.url)
+        app = asgi_proxy(args.url, timeout=args.timeout)
 
     uvicorn.run(app, host=args.host, port=args.port)