Add tests and other fixes

Author: jvstme

requirements_dev.txt

@@ -3,6 +3,7 @@ pre-commit
 httpx>=0.23
 pytest~=7.2
 pytest-asyncio>=0.21
+pytest-httpbin==2.1.0
 freezegun>=1.2.0
 ruff==0.5.3  # Should match .pre-commit-config.yaml
 testcontainers # testcontainers<4 may not work with asyncpg

src/dstack/_internal/gateway/deps.py

@@ -8,6 +8,13 @@
 
 
 class BaseGatewayDependencyInjector(ABC):
+    """
+    The gateway uses different implementations of this injector in different
+    environments: in-serer and on a remote host. An instance with the injector interface
+    stored in FastAPI's app.state.gateway_dependency_injector configures the gateway to
+    use a specific set of dependencies, e.g. a specific repo implementation.
+    """
+
     @abstractmethod
     async def get_repo(self) -> AsyncGenerator[BaseGatewayRepo, None]:
         if False:

src/dstack/_internal/gateway/repos/base.py

@@ -28,9 +28,17 @@ class Project(BaseModel):
 
 class BaseGatewayRepo(ABC):
     @abstractmethod
-    async def get_service(self, project_name: str, name: str) -> Optional[Service]:
+    async def get_service(self, project_name: str, run_name: str) -> Optional[Service]:
+        pass
+
+    @abstractmethod
+    async def add_service(self, project_name: str, service: Service) -> None:
         pass
 
     @abstractmethod
     async def get_project(self, name: str) -> Optional[Project]:
         pass
+
+    @abstractmethod
+    async def add_project(self, project: Project) -> None:
+        pass

src/dstack/_internal/gateway/repos/memory.py

@@ -0,0 +1,21 @@
+from typing import Dict, Optional
+
+from dstack._internal.gateway.repos.base import BaseGatewayRepo, Project, Service
+
+
+class InMemoryGatewayRepo(BaseGatewayRepo):
+    def __init__(self) -> None:
+        self.services: Dict[str, Dict[str, Service]] = {}
+        self.projects: Dict[str, Project] = {}
+
+    async def get_service(self, project_name: str, run_name: str) -> Optional[Service]:
+        return self.services.get(project_name, {}).get(run_name)
+
+    async def add_service(self, project_name: str, service: Service) -> None:
+        self.services.setdefault(project_name, {})[service.run_name] = service
+
+    async def get_project(self, name: str) -> Optional[Project]:
+        return self.projects.get(name)
+
+    async def add_project(self, project: Project) -> None:
+        self.projects[project.name] = project

src/dstack/_internal/gateway/services/service_proxy.py

@@ -5,7 +5,7 @@
 import httpx
 from starlette.requests import ClientDisconnect
 
-from dstack._internal.gateway.repos.base import BaseGatewayRepo
+from dstack._internal.gateway.repos.base import BaseGatewayRepo, Replica, Service
 from dstack._internal.gateway.services.service_connection import service_replica_connection_pool
 from dstack._internal.utils.logging import get_logger
 
@@ -38,15 +38,7 @@ async def proxy(
         )
 
     replica = random.choice(service.replicas)
-
-    connection = await service_replica_connection_pool.get(replica.id)
-    if connection is None:
-        project = await repo.get_project(project_name)
-        if project is None:
-            raise RuntimeError(f"Expected to find project {project_name} but could not")
-        connection = await service_replica_connection_pool.add(project, service, replica)
-
-    client = await connection.client()
+    client = await get_replica_client(project_name, service, replica, repo)
 
     try:
         upstream_request = await build_upstream_request(request, path, client, replica.id)
@@ -68,7 +60,7 @@ async def proxy(
             replica.id,
             e,
         )
-        if isinstance(e, TimeoutError):
+        if isinstance(e, httpx.TimeoutException):
             raise fastapi.HTTPException(fastapi.status.HTTP_504_GATEWAY_TIMEOUT)
         raise fastapi.HTTPException(fastapi.status.HTTP_502_BAD_GATEWAY)
 
@@ -79,6 +71,18 @@ async def proxy(
     )
 
 
+async def get_replica_client(
+    project_name: str, service: Service, replica: Replica, repo: BaseGatewayRepo
+) -> httpx.AsyncClient:
+    connection = await service_replica_connection_pool.get(replica.id)
+    if connection is None:
+        project = await repo.get_project(project_name)
+        if project is None:
+            raise RuntimeError(f"Expected to find project {project_name} but could not")
+        connection = await service_replica_connection_pool.add(project, service, replica)
+    return await connection.client()
+
+
 async def stream_response(
     response: httpx.Response, replica_id: str
 ) -> AsyncGenerator[bytes, None]:

src/dstack/_internal/server/services/gateway_in_server/repo.py

@@ -13,16 +13,22 @@
 
 
 class DBGatewayRepo(BaseGatewayRepo):
+    """
+    A gateway repo implementation used for gateway-in-server that retrieves data from
+    dstack-server's database. Since the database is populated by dstack-server, all or
+    most writer methods in this implementation are expected to be empty.
+    """
+
     def __init__(self, session: AsyncSession) -> None:
         self.session = session
 
-    async def get_service(self, project_name: str, name: str) -> Optional[Service]:
+    async def get_service(self, project_name: str, run_name: str) -> Optional[Service]:
         res = await self.session.execute(
             select(JobModel)
             .join(JobModel.project)
             .where(
                 ProjectModel.name == project_name,
-                JobModel.run_name == name,
+                JobModel.run_name == run_name,
                 JobModel.status == JobStatus.RUNNING,
                 JobModel.job_num == 0,
             )
@@ -67,6 +73,9 @@ async def get_service(self, project_name: str, name: str) -> Optional[Service]:
             replicas=replicas,
         )
 
+    async def add_service(self, project_name: str, service: Service) -> None:
+        pass
+
     async def get_project(self, name: str) -> Optional[Project]:
         res = await self.session.execute(select(ProjectModel).where(ProjectModel.name == name))
         project = res.scalar_one_or_none()
@@ -76,3 +85,6 @@ async def get_project(self, name: str) -> Optional[Project]:
             name=project.name,
             ssh_private_key=project.ssh_private_key,
         )
+
+    async def add_project(self, project: Project) -> None:
+        pass

src/tests/_internal/gateway/routers/test_service_proxy.py

@@ -0,0 +1,196 @@
+from typing import AsyncGenerator, Generator, Tuple
+from unittest.mock import patch
+
+import httpx
+import pytest
+from fastapi import FastAPI
+
+from dstack._internal.gateway.deps import BaseGatewayDependencyInjector
+from dstack._internal.gateway.repos.base import BaseGatewayRepo, Project, Replica, Service
+from dstack._internal.gateway.repos.memory import InMemoryGatewayRepo
+from dstack._internal.gateway.routers.service_proxy import router
+
+
+def make_app(repo: BaseGatewayRepo) -> FastAPI:
+    class DependencyInjector(BaseGatewayDependencyInjector):
+        async def get_repo(self) -> AsyncGenerator[BaseGatewayRepo, None]:
+            yield repo
+
+    app = FastAPI()
+    app.state.gateway_dependency_injector = DependencyInjector()
+    app.include_router(router, prefix="/gateway/services")
+    return app
+
+
+def make_client(app: FastAPI) -> httpx.AsyncClient:
+    return httpx.AsyncClient(transport=httpx.ASGITransport(app=app))
+
+
+def make_app_client(repo: BaseGatewayRepo) -> Tuple[FastAPI, httpx.AsyncClient]:
+    app = make_app(repo)
+    client = make_client(app)
+    return app, client
+
+
+def make_project(name: str) -> Project:
+    return Project(name=name, ssh_private_key="secret")
+
+
+def make_service(run_name: str) -> Service:
+    return Service(
+        id="xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx",
+        run_name=run_name,
+        auth=False,
+        app_port=80,
+        replicas=[
+            Replica(
+                id="xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx",
+                ssh_destination="ubuntu@server",
+                ssh_port=22,
+                ssh_proxy=None,
+            )
+        ],
+    )
+
+
+MOCK_REPLICA_CLIENT_TIMEOUT = 8
+
+
+@pytest.fixture
+def mock_replica_client_httpbin(httpbin) -> Generator[None, None, None]:
+    with patch(
+        "dstack._internal.gateway.services.service_proxy.get_replica_client"
+    ) as get_replica_client_mock:
+        get_replica_client_mock.return_value = httpx.AsyncClient(
+            base_url=httpbin.url, timeout=MOCK_REPLICA_CLIENT_TIMEOUT
+        )
+        yield
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("method", ["get", "post", "put", "patch", "delete"])
+async def test_proxy(mock_replica_client_httpbin, method: str) -> None:
+    methods_without_body = "get", "delete"
+    repo = InMemoryGatewayRepo()
+    await repo.add_project(make_project("test-proj"))
+    await repo.add_service(project_name="test-proj", service=make_service("httpbin"))
+    _, client = make_app_client(repo)
+    req_body = "." * 20 * 2**20 if method not in methods_without_body else None
+    resp = await client.request(
+        method,
+        f"http://test-host:8888/gateway/services/test-proj/httpbin/{method}?a=b&c=",
+        headers={"User-Agent": "test-ua", "Connection": "keep-alive"},
+        content=req_body,
+    )
+    assert resp.status_code == 200
+    assert resp.headers["server"].startswith("Pytest-HTTPBIN")
+    resp_body = resp.json()
+    assert resp_body["url"] == f"http://test-host:8888/{method}?a=b&c="
+    assert resp_body["args"] == {"a": "b", "c": ""}
+    assert resp_body["headers"]["Host"] == "test-host:8888"
+    assert resp_body["headers"]["User-Agent"] == "test-ua"
+    assert resp_body["headers"]["Connection"] == "keep-alive"
+    if method not in methods_without_body:
+        assert resp_body["data"] == req_body
+
+
+@pytest.mark.asyncio
+async def test_proxy_method_head(mock_replica_client_httpbin) -> None:
+    repo = InMemoryGatewayRepo()
+    await repo.add_project(make_project("test-proj"))
+    await repo.add_service(project_name="test-proj", service=make_service("httpbin"))
+    _, client = make_app_client(repo)
+    url = "http://test-host/gateway/services/test-proj/httpbin/"
+    get_resp = await client.get(url)
+    head_resp = await client.head(url)
+    assert get_resp.status_code == head_resp.status_code == 200
+    assert head_resp.headers["Content-Length"] == get_resp.headers["Content-Length"]
+    assert int(head_resp.headers["Content-Length"]) > 0
+    assert head_resp.content == b""
+
+
+@pytest.mark.asyncio
+async def test_proxy_method_options(mock_replica_client_httpbin) -> None:
+    repo = InMemoryGatewayRepo()
+    await repo.add_project(make_project("test-proj"))
+    await repo.add_service(project_name="test-proj", service=make_service("httpbin"))
+    _, client = make_app_client(repo)
+    resp = await client.options("http://test-host/gateway/services/test-proj/httpbin/get")
+    assert resp.status_code == 200
+    assert set(resp.headers["Allow"].split(", ")) == {"HEAD", "GET", "OPTIONS"}
+    assert resp.content == b""
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("code", [204, 304, 418, 503])
+async def test_proxy_status_codes(mock_replica_client_httpbin, code: int) -> None:
+    repo = InMemoryGatewayRepo()
+    await repo.add_project(make_project("test-proj"))
+    await repo.add_service(project_name="test-proj", service=make_service("httpbin"))
+    _, client = make_app_client(repo)
+    resp = await client.get(f"http://test-host/gateway/services/test-proj/httpbin/status/{code}")
+    assert resp.status_code == code
+
+
+@pytest.mark.asyncio
+async def test_proxy_not_leaks_cookies(mock_replica_client_httpbin) -> None:
+    repo = InMemoryGatewayRepo()
+    await repo.add_project(make_project("test-proj"))
+    await repo.add_service(project_name="test-proj", service=make_service("httpbin"))
+    app = make_app(repo)
+    client1 = make_client(app)
+    client2 = make_client(app)
+    cookies_url = "http://test-host/gateway/services/test-proj/httpbin/cookies"
+    await client1.get(cookies_url + "/set?a=1")
+    await client1.get(cookies_url + "/set?b=2")
+    await client2.get(cookies_url + "/set?a=3")
+    resp1 = await client1.get(cookies_url)
+    resp2 = await client2.get(cookies_url)
+    assert resp1.json()["cookies"] == {"a": "1", "b": "2"}
+    assert resp2.json()["cookies"] == {"a": "3"}
+
+
+@pytest.mark.asyncio
+async def test_proxy_gateway_timeout(mock_replica_client_httpbin) -> None:
+    repo = InMemoryGatewayRepo()
+    await repo.add_project(make_project("test-proj"))
+    await repo.add_service(project_name="test-proj", service=make_service("httpbin"))
+    _, client = make_app_client(repo)
+    assert MOCK_REPLICA_CLIENT_TIMEOUT < 10
+    resp = await client.get("http://test-host/gateway/services/test-proj/httpbin/delay/10")
+    assert resp.status_code == 504
+    assert resp.json()["detail"] == "Gateway Timeout"
+
+
+@pytest.mark.asyncio
+async def test_proxy_run_not_found(mock_replica_client_httpbin) -> None:
+    repo = InMemoryGatewayRepo()
+    await repo.add_project(make_project("test-proj"))
+    await repo.add_service(project_name="test-proj", service=make_service("test-run"))
+    _, client = make_app_client(repo)
+    resp = await client.get("http://test-host/gateway/services/test-proj/unknown/")
+    assert resp.status_code == 404
+    assert resp.json()["detail"] == "Service test-proj/unknown not found"
+
+
+@pytest.mark.asyncio
+async def test_proxy_project_not_found(mock_replica_client_httpbin) -> None:
+    _, client = make_app_client(InMemoryGatewayRepo())
+    resp = await client.get("http://test-host/gateway/services/unknown/test-run/")
+    assert resp.status_code == 404
+    assert resp.json()["detail"] == "Service unknown/test-run not found"
+
+
+@pytest.mark.asyncio
+async def test_redirect_to_service_root(mock_replica_client_httpbin) -> None:
+    repo = InMemoryGatewayRepo()
+    await repo.add_project(make_project("test-proj"))
+    await repo.add_service(project_name="test-proj", service=make_service("httpbin"))
+    _, client = make_app_client(repo)
+    url = "http://test-host/gateway/services/test-proj/httpbin"
+    resp = await client.get(url, follow_redirects=False)
+    assert resp.status_code == 308
+    assert resp.headers["Location"] == url + "/"
+    resp = await client.get(url, follow_redirects=True)
+    assert resp.status_code == 200
+    assert resp.request.url == url + "/"