diff --git a/python/openshell/sandbox.py b/python/openshell/sandbox.py index 2eaf6a64fa..c02a33cdbf 100644 --- a/python/openshell/sandbox.py +++ b/python/openshell/sandbox.py @@ -29,7 +29,7 @@ openshell_pb2, openshell_pb2_grpc, ) -from .errors import _error_mapping_channel +from .errors import _error_mapping_channel, from_grpc_error from .mutations import DeletionOutcome, DeletionResult _ClientCallDetailsBase = namedtuple( @@ -40,6 +40,10 @@ _OAUTH_MAX_RESPONSE_BYTES = 1 << 20 _PAGER_MAX_CONSUMED_TOKENS = 10_000 _PAGER_MAX_CONSUMED_TOKEN_BYTES = 1 << 20 +# Replica routing headers and retry rule: docs/kubernetes/ingress.mdx. +_OWNER_REPLICA_HEADER = "x-openshell-owner" +_OWNER_REQUEST_HEADER = "x-openshell-want-owner" +_ROUTE_REPLICA_HEADER = "x-openshell-replica" T = TypeVar("T") @@ -99,12 +103,62 @@ def all(self) -> builtins.list[T]: return [item for page in self for item in page.items] +class _OwnerCache: + """Owner replicas the gateway last reported, keyed by workspace and name. + + An empty owner means the gateway reported none. A stale owner is safe: a + live non-owner replica relays the call. + """ + + _TTL_SECONDS = 30.0 + _MAX_ENTRIES = 1024 + + def __init__(self) -> None: + self._lock = threading.Lock() + self._entries: dict[tuple[str, str], tuple[float, str]] = {} + + def get(self, workspace: str, name: str) -> str | None: + with self._lock: + entry = self._entries.get((workspace, name)) + if entry is None or time.monotonic() - entry[0] >= self._TTL_SECONDS: + return None + return entry[1] + + def put(self, workspace: str, name: str, owner: str) -> None: + with self._lock: + if len(self._entries) >= self._MAX_ENTRIES: + self._entries.clear() + self._entries[(workspace, name)] = (time.monotonic(), owner) + + def forget(self, workspace: str, name: str) -> None: + with self._lock: + self._entries.pop((workspace, name), None) + + +def _retry_unrouted(error: grpc.RpcError) -> bool: + """Envoy Gateway answers a replica route with no ready endpoints with + UNAVAILABLE and no message; gateway errors always carry one.""" + status = from_grpc_error(error) + details = status.details() or "" + return status.code() == grpc.StatusCode.UNAVAILABLE and ( + details == "" + or "no healthy upstream" in details + or "connection failure" in details + ) + + def _workspace_scope(workspace: str) -> datamodel_pb2.WorkspaceSelector: if not workspace: raise ValueError("workspace must be non-empty") return datamodel_pb2.WorkspaceSelector(workspace=workspace) +def _get_sandbox_request(name: str, workspace: str) -> openshell_pb2.GetSandboxRequest: + return openshell_pb2.GetSandboxRequest( + workspace_scope=_workspace_scope(workspace), name=name + ) + + def _all_workspaces_scope() -> datamodel_pb2.WorkspaceSelector: return datamodel_pb2.WorkspaceSelector(all_workspaces=datamodel_pb2.AllWorkspaces()) @@ -730,6 +784,7 @@ def __init__( self._stub = openshell_pb2_grpc.OpenShellStub( _error_mapping_channel(self._channel) ) + self._owners = _OwnerCache() @classmethod def from_active_cluster( @@ -967,10 +1022,7 @@ def providers(self) -> ProviderClient: def get(self, name: str, *, workspace: str) -> SandboxRef: response = self._stub.GetSandbox( - openshell_pb2.GetSandboxRequest( - workspace_scope=_workspace_scope(workspace), - name=name, - ), + _get_sandbox_request(name, workspace), timeout=self._timeout, ) return _sandbox_ref(response.sandbox) @@ -1280,7 +1332,7 @@ def exec_stream( grpc_deadline = self._timeout if timeout_seconds and timeout_seconds + 10 > grpc_deadline: grpc_deadline = timeout_seconds + 10 - stream = self._stub.ExecSandbox(request, timeout=grpc_deadline) + stream = self._open_exec(request, workspace, grpc_deadline) stdout_parts: list[bytes] = [] stderr_parts: list[bytes] = [] @@ -1308,6 +1360,51 @@ def exec_stream( stderr=b"".join(stderr_parts).decode("utf-8", errors="replace"), ) + def _owner_of(self, workspace: str, sandbox: str) -> str: + """Best effort: a failed lookup leaves the call unrouted.""" + owner = self._owners.get(workspace, sandbox) + if owner is not None: + return owner + try: + _, call = self._stub.GetSandbox.with_call( + _get_sandbox_request(sandbox, workspace), + timeout=self._timeout, + metadata=((_OWNER_REQUEST_HEADER, "1"),), + ) + except grpc.RpcError: + return "" + value = dict(call.initial_metadata() or ()).get(_OWNER_REPLICA_HEADER) + owner = value if isinstance(value, str) else "" + self._owners.put(workspace, sandbox, owner) + return owner + + def _open_exec( + self, + request: openshell_pb2.ExecSandboxRequest, + workspace: str, + timeout: float, + ) -> Iterator[openshell_pb2.ExecSandboxEvent]: + """Open ExecSandbox on the owner replica, retrying once unrouted.""" + owner = self._owner_of(workspace, request.sandbox) + if owner: + events = self._stub.ExecSandbox( + request, + timeout=timeout, + metadata=((_ROUTE_REPLICA_HEADER, owner),), + ) + try: + first = next(events, None) + except grpc.RpcError as error: + if not _retry_unrouted(error): + raise + self._owners.forget(workspace, request.sandbox) + else: + if first is not None: + yield first + yield from events + return + yield from self._stub.ExecSandbox(request, timeout=timeout) + def exec( self, sandbox: str, diff --git a/python/openshell/sandbox_test.py b/python/openshell/sandbox_test.py index c1df0d46e0..3fed31734e 100644 --- a/python/openshell/sandbox_test.py +++ b/python/openshell/sandbox_test.py @@ -3,22 +3,25 @@ from __future__ import annotations +import contextlib import json import os import pickle import threading import time +from concurrent import futures from copy import deepcopy from dataclasses import asdict from pathlib import Path from types import SimpleNamespace -from typing import Any, cast +from typing import TYPE_CHECKING, Any, Never, cast +import grpc import pytest from google.protobuf.struct_pb2 import Struct import openshell.sandbox as sandbox_module -from openshell._proto import openshell_pb2 +from openshell._proto import openshell_pb2, openshell_pb2_grpc from openshell.mutations import DeletionOutcome from openshell.sandbox import ( _OIDC_TOKEN_EXPIRY_GRACE_SECONDS, @@ -43,11 +46,15 @@ _make_cluster_bearer_provider, _normalize_bearer, _OidcRefresher, + _OwnerCache, _read_oidc_token_bundle, _sandbox_ref, _validate_oauth_url, ) +if TYPE_CHECKING: + from collections.abc import Generator + def _request_workspace(request: Any) -> str | None: scope = request.workspace_scope @@ -412,6 +419,11 @@ def handler(request: Any) -> Any: class _FakeStub: + class GetSandbox: + @staticmethod + def with_call(*_args: object, **_kwargs: object) -> Never: + raise grpc.RpcError() + def __init__(self) -> None: self.request: openshell_pb2.ExecSandboxRequest | None = None @@ -431,9 +443,102 @@ def _client_with_fake_stub(stub: object) -> SandboxClient: client = object.__new__(SandboxClient) client._timeout = 30.0 client._stub = cast("Any", stub) + client._owners = _OwnerCache() return client +class _RoutingGateway(openshell_pb2_grpc.OpenShellServicer): + """Gateway that reports an owner and, behind a replica route, can answer + like Envoy with a trailers-only UNAVAILABLE.""" + + def __init__( + self, owner: str, *, reject_routed: str | None = None, fail: str = "" + ) -> None: + self.owner = owner + self.reject_routed = reject_routed + self.fail = fail + self.wanted_owner: list[str] = [] + self.routes: list[str] = [] + + def GetSandbox(self, request, context): + _ = request + headers = dict(context.invocation_metadata()) + self.wanted_owner.append(headers.get("x-openshell-want-owner", "")) + if self.owner: + context.send_initial_metadata((("x-openshell-owner", self.owner),)) + return openshell_pb2.SandboxResponse() + + def ExecSandbox(self, request, context): + _ = request + route = dict(context.invocation_metadata()).get("x-openshell-replica", "") + self.routes.append(route) + if route and self.reject_routed is not None: + context.abort(grpc.StatusCode.UNAVAILABLE, self.reject_routed) + if self.fail: + context.abort(grpc.StatusCode.UNAVAILABLE, self.fail) + yield openshell_pb2.ExecSandboxEvent( + stdout=openshell_pb2.ExecSandboxStdout(data=b"ok") + ) + yield openshell_pb2.ExecSandboxEvent( + exit=openshell_pb2.ExecSandboxExit(exit_code=0) + ) + + +@contextlib.contextmanager +def _routing_client(gateway: _RoutingGateway) -> Generator[SandboxClient]: + server = grpc.server(futures.ThreadPoolExecutor(max_workers=2)) + openshell_pb2_grpc.add_OpenShellServicer_to_server(gateway, server) + port = server.add_insecure_port("127.0.0.1:0") + server.start() + try: + with SandboxClient(f"127.0.0.1:{port}", timeout=5) as client: + yield client + finally: + server.stop(None) + + +def test_exec_routes_to_reported_owner_and_caches_it() -> None: + gateway = _RoutingGateway("openshell-1") + with _routing_client(gateway) as client: + for _ in range(2): + result = client.exec("box", ["true"], workspace="default") + assert result.stdout == "ok" + + assert gateway.wanted_owner == ["1"] + assert gateway.routes == ["openshell-1", "openshell-1"] + + +def test_exec_without_owner_sends_no_hint() -> None: + gateway = _RoutingGateway("") + with _routing_client(gateway) as client: + assert client.exec("box", ["true"], workspace="default").exit_code == 0 + + assert gateway.routes == [""] + + +@pytest.mark.parametrize("message", ["", "no healthy upstream"]) +def test_exec_retries_once_unrouted_when_owner_route_is_empty(message: str) -> None: + gateway = _RoutingGateway("openshell-1", reject_routed=message) + with _routing_client(gateway) as client: + for _ in range(2): + result = client.exec("box", ["true"], workspace="default") + assert result.stdout == "ok" + + assert gateway.routes == ["openshell-1", "", "openshell-1", ""] + assert gateway.wanted_owner == ["1", "1"] + + +def test_exec_does_not_retry_gateway_errors() -> None: + gateway = _RoutingGateway("openshell-1", fail="supervisor not connected") + with ( + _routing_client(gateway) as client, + pytest.raises(grpc.RpcError, match="supervisor not connected"), + ): + client.exec("box", ["true"], workspace="default") + + assert gateway.routes == ["openshell-1"] + + def _template_client_with_fake_stub(stub: object) -> SandboxTemplateClient: client = object.__new__(SandboxTemplateClient) client._timeout = 30.0