Gateway recovery fixes: auth pool and replica tunnel reconnection

pyproject.toml

@@ -220,6 +220,8 @@ gateway = [
     "aiorwlock",
     "aiocache",
     "httpx>=0.28.0",
+    # Gateway pool workarounds use private HTTPcore APIs; review before upgrading.
+    "httpcore>=1.0.5,<1.1",
     "jinja2",
 ]
 server = [

src/dstack/_internal/proxy/gateway/app.py

@@ -1,6 +1,7 @@
 """FastAPI app running on a gateway."""
 
-from contextlib import asynccontextmanager
+import asyncio
+from contextlib import asynccontextmanager, suppress
 from pathlib import Path
 from typing import Optional
 
@@ -29,6 +30,7 @@
 from dstack._internal.proxy.gateway.services.server_client import HTTPMultiClient
 from dstack._internal.proxy.gateway.services.stats import StatsCollector
 from dstack._internal.proxy.lib.routers.model_proxy import router as model_proxy_router
+from dstack._internal.proxy.lib.services.service_connection import maintain_service_connections
 from dstack._internal.utils.common import run_async
 from dstack.version import __version__
 
@@ -45,9 +47,13 @@ async def lifespan(app: FastAPI):
     service_conn_pool = await injector.get_service_connection_pool()
     await run_async(nginx.write_global_conf)
     await apply_all(repo, nginx, service_conn_pool)
+    maintenance = asyncio.create_task(maintain_service_connections(service_conn_pool))
 
     yield
 
+    maintenance.cancel()
+    with suppress(asyncio.CancelledError):
+        await maintenance
     await service_conn_pool.remove_all()
 
 

src/dstack/_internal/proxy/gateway/auth.py

@@ -4,12 +4,15 @@
 from dstack._internal.proxy.lib.auth import BaseProxyAuthProvider
 from dstack._internal.proxy.lib.errors import UnexpectedProxyError
 
+_AUTH_CACHE_TTL_SECONDS = 60
+"""Seconds to cache allowed or denied access for each project/token pair."""
+
 
 class GatewayProxyAuthProvider(BaseProxyAuthProvider):
     def __init__(self, server_client: httpx.AsyncClient) -> None:
         self._server_client = server_client
 
-    @cached(ttl=60, noself=True, skip_cache_func=lambda r: r is None)
+    @cached(ttl=_AUTH_CACHE_TTL_SECONDS, noself=True, skip_cache_func=lambda r: r is None)
     async def is_project_member(self, project_name: str, token: str) -> bool:
         try:
             resp = await self._server_client.post(

src/dstack/_internal/proxy/gateway/resources/licenses/httpcore.txt

@@ -0,0 +1,27 @@
+Copyright © 2020, [Encode OSS Ltd](https://www.encode.io/).
+All rights reserved.
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+* Redistributions of source code must retain the above copyright notice, this
+  list of conditions and the following disclaimer.
+
+* Redistributions in binary form must reproduce the above copyright notice,
+  this list of conditions and the following disclaimer in the documentation
+  and/or other materials provided with the distribution.
+
+* Neither the name of the copyright holder nor the names of its
+  contributors may be used to endorse or promote products derived from
+  this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

src/dstack/_internal/proxy/gateway/resources/systemd/dstack.gateway.service

@@ -9,6 +9,7 @@ User=ubuntu
 Group=ubuntu
 Restart=always
 RestartSec=5
+LimitNOFILE=65535
 
 [Install]
 WantedBy=default.target

src/dstack/_internal/proxy/gateway/services/server_client.py

@@ -2,9 +2,11 @@
 import logging
 import random
 from dataclasses import dataclass, field
+from itertools import chain
 from pathlib import Path
 from typing import Container, Dict, Generator, List
 
+import httpcore
 import httpx
 
 logger = logging.getLogger(__name__)
@@ -13,22 +15,25 @@
 
 @dataclass
 class CachedClientInfo:
+    """HTTP client and connection-failure history for one server, owned by HTTPMultiClient."""
+
     client: httpx.AsyncClient
     socket: Path
     connect_errors: List[datetime.datetime] = field(default_factory=lambda: [])
 
     def seems_disconnected(self) -> bool:
+        # Only ConnectError is recorded; a returned HTTP response clears this history.
         if len(self.connect_errors) < 2:
             return False
         return self.connect_errors[-1] - self.connect_errors[0] >= datetime.timedelta(minutes=2)
 
 
 class HTTPMultiClient(httpx.AsyncClient):
-    """
-    An HTTP client that sends requests to randomly chosen Unix sockets from a specified
-    directory. This allows to balance the load between multiple HTTP server replicas.
-    Automatically deletes sockets that stop responding.
-    Used for requesting random dstack-server replicas from the gateway.
+    """Gateway HTTP client for uncached project-access checks against dstack-server.
+
+    GatewayProxyAuthProvider uses this client and owns authorization decisions and
+    their cache. Return the server's HTTP response, including error statuses. Raise
+    httpx.RequestError when no connected server can complete the request.
     """
 
     def __init__(self, sockets_dir: Path):
@@ -40,6 +45,7 @@ async def send(self, request: httpx.Request, *args, **kwargs) -> httpx.Response:
         errors: List[httpx.RequestError] = []
         clients_count = 0
 
+        # Try another replica after request errors; HTTP error responses are returned.
         for clients_count, client in enumerate(self._iter_clients_rand(), start=1):
             try:
                 resp = await client.client.send(request, *args, **kwargs)
@@ -69,6 +75,7 @@ async def send(self, request: httpx.Request, *args, **kwargs) -> httpx.Response:
     def _iter_clients_rand(self) -> Generator[CachedClientInfo, None, None]:
         sockets = list(self._sockets_dir.glob("*.sock"))
         self._evict_clients(stems_to_keep={s.stem for s in sockets})
+        # Each socket forwards to a server replica; shuffle to distribute auth checks.
         random.shuffle(sockets)
 
         for socket in sockets:
@@ -81,7 +88,7 @@ def _iter_clients_rand(self) -> Generator[CachedClientInfo, None, None]:
     @staticmethod
     def _make_client(socket: Path) -> CachedClientInfo:
         client = httpx.AsyncClient(
-            transport=httpx.AsyncHTTPTransport(uds=str(socket.absolute())),
+            transport=_ServerTransport(uds=str(socket.absolute())),
             base_url=BASE_URL,
         )
         return CachedClientInfo(
@@ -93,3 +100,140 @@ def _evict_clients(self, stems_to_keep: Container[str]) -> None:
         self._clients_cache = {
             stem: client for stem, client in self._clients_cache.items() if stem in stems_to_keep
         }
+
+
+class _ServerTransport(httpx.AsyncHTTPTransport):
+    """HTTPMultiClient's transport for one dstack-server Unix socket.
+
+    Preserve HTTPX's connection limits and use the request's HTTPX timeouts.
+    _ServerConnectionPool provides recovery after overload.
+    """
+
+    def __init__(self, uds: str):
+        super().__init__(uds=uds)
+        # Preserve HTTPX's existing default limits.
+        self._pool = _ServerConnectionPool(
+            uds=uds,
+            max_connections=100,
+            max_keepalive_connections=20,
+            keepalive_expiry=5.0,
+        )
+
+
+class _ServerConnectionPool(httpcore.AsyncConnectionPool):
+    """Authorization connection pool used by _ServerTransport, with overload recovery.
+
+    Preserve HTTPcore's connection limits, reuse, ordering and timeout behavior.
+    Failed or cancelled requests must not permanently consume capacity or prevent
+    later requests from proceeding once connections are available.
+    See https://github.com/dstackai/dstack/issues/4338.
+    """
+
+    async def handle_async_request(self, request: httpcore.Request) -> httpcore.Response:
+        try:
+            return await super().handle_async_request(request)
+        except BaseException:
+            # HTTPcore removes the failed/cancelled request, but can leave its newly
+            # assigned connection behind before it starts. That slot is never reused.
+            # Remove this orphan cleanup once https://github.com/encode/httpcore/pull/1099 ships.
+            with self._optional_thread_lock:
+                owners = {id(request.connection) for request in self._requests}
+                orphaned = [
+                    connection
+                    for connection in self._connections
+                    if id(connection) not in owners
+                    and not connection.is_closed()
+                    and not connection.is_available()
+                    and not connection.is_idle()
+                ]
+                if not orphaned:
+                    raise
+                for connection in orphaned:
+                    self._connections.remove(connection)
+                # Wake queued requests now that slots are free. Response streams still
+                # retain their owners, and HTTPcore shields connection closing below.
+                closing = orphaned + self._assign_requests_to_connections()
+            await self._close_connections(closing)
+            raise
+
+    def _assign_requests_to_connections(self) -> List[httpcore.AsyncConnectionInterface]:
+        # Adapted from HTTPcore; see resources/licenses/httpcore.txt.
+        if (
+            len(self._requests) > len(self._connections)
+            and len(self._connections) >= self._max_connections
+            and all(
+                not connection.is_closed()
+                and not connection.has_expired()
+                and not connection.is_idle()
+                and not connection.is_available()
+                for connection in self._connections
+            )
+        ):
+            # A full, busy pool cannot assign or close anything. Scanning every queued
+            # request here makes a backlog of timeouts quadratic to drain.
+            return []
+        closing = []
+        for connection in list(self._connections):
+            if connection.is_closed():
+                self._connections.remove(connection)
+            elif connection.has_expired():
+                self._connections.remove(connection)
+                closing.append(connection)
+            elif connection.is_idle() and len(self._connections) > self._max_keepalive_connections:
+                # Preserve HTTPcore's existing total-connection keepalive check.
+                self._connections.remove(connection)
+                closing.append(connection)
+
+        if not self._requests:
+            return closing
+        if len(self._requests) == 1:
+            request = self._requests[0]
+            if not request.is_queued():
+                return closing
+            # Reuse a connection without building lists for a single request.
+            origin = request.request.url.origin
+            for connection in self._connections:
+                if connection.can_handle_request(origin) and connection.is_available():
+                    request.assign_to_connection(connection)
+                    return closing
+
+        # Skip assignment scans when no request is waiting for a connection.
+        queued_requests = (request for request in self._requests if request.is_queued())
+        first_request = next(queued_requests, None)
+        if first_request is None:
+            return closing
+
+        # No connection state changes while this synchronous pass runs. Inspect
+        # availability once, rather than scanning every connection for every waiter.
+        available = [connection for connection in self._connections if connection.is_available()]
+        idle = [connection for connection in self._connections if connection.is_idle()]
+
+        def add_connection(origin: httpcore.Origin) -> httpcore.AsyncConnectionInterface:
+            connection = self.create_connection(origin)
+            self._connections.append(connection)
+            if connection.is_available():
+                available.append(connection)
+            if connection.is_idle():
+                idle.append(connection)
+            return connection
+
+        for request in chain((first_request,), queued_requests):
+            if len(self._connections) >= self._max_connections and not available and not idle:
+                break
+            origin = request.request.url.origin
+            connection = next(
+                (connection for connection in available if connection.can_handle_request(origin)),
+                None,
+            )
+            if connection is not None:
+                request.assign_to_connection(connection)
+            elif len(self._connections) < self._max_connections:
+                request.assign_to_connection(add_connection(origin))
+            elif idle:
+                connection = idle.pop(0)
+                self._connections.remove(connection)
+                if connection in available:
+                    available.remove(connection)
+                closing.append(connection)
+                request.assign_to_connection(add_connection(origin))
+        return closing