diff --git a/backend/fluksio/flow/controller.py b/backend/fluksio/flow/controller.py index 7ca2e48..963141a 100644 --- a/backend/fluksio/flow/controller.py +++ b/backend/fluksio/flow/controller.py @@ -65,6 +65,7 @@ from fluksio.flow.nodes import ( WrenNode, ) from fluksio.flow.pipeline import ( + EngineHeld, NodeOutcome, Pipeline, RunCacheLookup, @@ -1984,6 +1985,9 @@ class FlowController: Returns what the node raised, or an empty string. A node failing by hand is reported the way a queued run reports it — on the canvas and in the metrics — rather than as a stack trace in the server log. + + :raises EngineHeld: While a rename holds the engine, which is no + failure of the node's. """ node = self.get_node(node_id) if node is None: @@ -1993,6 +1997,8 @@ class FlowController: node.trigger(values, durable=False) else: node.inject(values or {}, durable=False) + except EngineHeld: + raise except Exception as exc: if self.pipeline is None: raise diff --git a/backend/fluksio/flow/rename.py b/backend/fluksio/flow/rename.py new file mode 100644 index 0000000..536f55d --- /dev/null +++ b/backend/fluksio/flow/rename.py @@ -0,0 +1,122 @@ +"""Renaming a flow's key: the door that keeps the API out while it runs. + +A rename moves everything named under a flow's key, in documents, history, +live state and queued work, so nothing may write through the API meanwhile. +`RenameGate` is that door. ``docs/private/id_refac_rename.md`` has the whole +protocol. +""" + +from __future__ import annotations + +import re +import threading +from collections import Counter +from collections.abc import Callable, Iterable + +from starlette.types import Scope + +from fluksio.core.config import settings +from fluksio.flow.controller import HOOK_PREFIX +from fluksio.flow.executor import wait_out + +#: What reading looks like. Everything else under the API changes something. +_READS = frozenset({"GET", "HEAD", "OPTIONS"}) + +#: Mutations that cannot reach a flow's namespace, under the API prefix: signing +#: in and accounts, the assistant (whose tool calls come back through the API +#: and are held there), a browser's push subscription, the portal, workers +#: attaching, artifact bytes (which are named by digest, and which a node on a +#: remote worker writes while the rename waits for it), the socket ticket and +#: the rename routes, which keep their own order. +_EXEMPT = re.compile( + r"/(?:login|password-recovery|password-recovery-html-content|reset-password" + r"|oauth|users|ai|cloud|workers|artifacts)(?:/|$)" + r"|/alerts/webpush(?:/|$)" + r"|/flows/ws-ticket$" + r"|/(?:flows|dashboards)/[^/]+/rename(?:/plan)?$" +) + + +class RenameGate: + """The door every request passes, which a rename closes. + + Open, it counts the mutating API requests it lets in, for as long as each + takes. Closed, it turns new ones away and the rename waits for those + already in, so it starts from a namespace nothing is writing to through + the API. Hooks of the flows being renamed are turned away too, since their + nodes are torn down meanwhile. Reads, websockets and the exempt routes + pass either way, uncounted. + + Thread-safe: requests pass on the event loop and `close` waits on a + thread, as the engine holds do. + """ + + def __init__(self) -> None: + self._lock = threading.Condition() + self._admitted: Counter[str] = Counter() + self._closed = False + self._affected: frozenset[str] = frozenset() + #: The rename that closed it, for whoever is turned away. + self.op = "" + + @property + def closed(self) -> bool: + with self._lock: + return self._closed + + def admit(self, scope: Scope) -> tuple[bool, str]: + """Whether *scope* may pass, and the name it is counted under until + `leave`, empty when it is not counted.""" + path: str = scope["path"] + with self._lock: + if path.startswith(f"{HOOK_PREFIX}/"): + flow_id = path[len(HOOK_PREFIX) + 1 :].split("/", 1)[0] + return not (self._closed and flow_id in self._affected), "" + api = settings.API_V1_STR + if ( + scope["type"] != "http" + or scope["method"] in _READS + or not path.startswith(f"{api}/") + or _EXEMPT.match(path, len(api)) + ): + return True, "" + if self._closed: + return False, "" + name = f"{scope['method']} {path}" + self._admitted[name] += 1 + return True, name + + def leave(self, name: str) -> None: + """A request counted under *name* is done.""" + with self._lock: + self._admitted -= Counter([name]) + self._lock.notify_all() + + def close( + self, + op: str, + affected: Iterable[str], + abort: threading.Event, + report: Callable[[list[str]], None] | None = None, + ) -> bool: + """Turn mutations and the hooks of the *affected* flow ids away, then + wait for the requests already in to finish. Blocking. + + Without a limit: ``report`` hears the requests still running each time + they change, and setting ``abort`` gives up, opens the gate again and + answers False. True means closed until `open`. + """ + with self._lock: + self._closed = True + self._affected = frozenset(affected) + self.op = op + if wait_out(self._lock, lambda: sorted(self._admitted), abort, report): + return True + self.open() + return False + + def open(self) -> None: + with self._lock: + self._closed = False + self._affected = frozenset() + self.op = "" diff --git a/backend/fluksio/flow/runs.py b/backend/fluksio/flow/runs.py index e7d1728..555d4a8 100644 --- a/backend/fluksio/flow/runs.py +++ b/backend/fluksio/flow/runs.py @@ -813,6 +813,7 @@ class RunCache: self, artifacts: ArtifactStore | None = None, key_of: Callable[[str], str] | None = None, + renaming: Callable[[], bool] | None = None, ) -> None: self.artifacts = artifacts #: The key a source run's outputs are named with, from its flow's id: @@ -820,9 +821,13 @@ class RunCache: #: one the id is taken for the key, which it is for a flow from before #: ids. self.key_of = key_of or (lambda flow_id: flow_id) + #: Whether a rename is under way, which may have moved some of a + #: source run's output names and not yet the others. Every lookup + #: misses meanwhile, which only costs running the node. + self.renaming = renaming or (lambda: False) def lookup(self, key: str) -> CacheHit | None: - if not key: + if not key or self.renaming(): return None try: with Session(db_engine) as session: @@ -874,12 +879,15 @@ class RunService: state_factory: Callable[[str], StateBackend] | None = None, parallel: int = MAX_PARALLEL, artifacts: ArtifactStore | None = None, + renaming: Callable[[], bool] | None = None, ) -> None: self.controller = controller self.queue = queue self._artifacts = artifacts self._cache = RunCache( - artifacts, lambda flow_id: tombstones.name_of(flow_id, controller.store) + artifacts, + lambda flow_id: tombstones.name_of(flow_id, controller.store), + renaming, ) # Without one, a run gets a private in-memory state — which is exactly # the isolation it wants, minus surviving the process. diff --git a/backend/fluksio/main.py b/backend/fluksio/main.py index 078e8b4..d0d0395 100644 --- a/backend/fluksio/main.py +++ b/backend/fluksio/main.py @@ -49,6 +49,7 @@ from fluksio.flow.plugins import load_plugins from fluksio.flow.provision import load_provisioners from fluksio.flow.queue import MemoryWorkQueue, RedisWorkQueue, WorkQueue from fluksio.flow.remote import RemoteWorkerHub +from fluksio.flow.rename import RenameGate from fluksio.flow.resources import ( ResourceAccountant, fair_share_env, @@ -419,6 +420,7 @@ async def lifespan(app: FastAPI) -> AsyncIterator[None]: state_factory=_run_state, artifacts=artifacts, parallel=settings.FLOW_MAX_RUNS, + renaming=lambda: rename_gate.closed, ) app.state.run_service = run_service # Said out loud because it is the only way to tell that a settings file was @@ -536,6 +538,52 @@ class DocsForThisMachine: await self.app(scope, receive, send) +class AdmissionGate: + """Turn away what a rename must not see half done (`RenameGate`): + mutating API requests, and the hooks of the flows it renames.""" + + def __init__(self, app: ASGIApp, gate: RenameGate) -> None: + self.app = app + self.gate = gate + + async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: + if scope["type"] not in ("http", "websocket"): + await self.app(scope, receive, send) + return + allowed, counted = self.gate.admit(scope) + if not allowed: + await _try_later(self.gate.op, scope, receive, send) + return + try: + await self.app(scope, receive, send) + finally: + if counted: + self.gate.leave(counted) + + +async def _try_later(op: str, scope: Scope, receive: Receive, send: Send) -> None: + """A 503 to retry, or a socket closed with 1013, try again later. + + Accepted first: a socket closed before that is a refused handshake, and + the client never sees the code. + """ + if scope["type"] == "websocket": + await receive() + await send({"type": "websocket.accept"}) + await send({"type": "websocket.close", "code": 1013}) + return + response = JSONResponse( + status_code=503, + content={ + "detail": "A rename is in progress; try again shortly", + "code": "renaming", + "op": op, + }, + headers={"Retry-After": "5"}, + ) + await response(scope, receive, send) + + def _is_loopback(host: str) -> bool: try: return ipaddress.ip_address(host).is_loopback @@ -553,6 +601,12 @@ app = FastAPI( lifespan=lifespan, ) +# Closed by a rename. Added first, so it sits inside the CORS layer and a +# browser can read the 503 it answers. +rename_gate = RenameGate() +app.state.rename_gate = rename_gate +app.add_middleware(AdmissionGate, gate=rename_gate) + # Set all CORS enabled origins if settings.all_cors_origins: app.add_middleware( diff --git a/backend/tests/api/routes/test_runs.py b/backend/tests/api/routes/test_runs.py index 55b7186..3dd4261 100644 --- a/backend/tests/api/routes/test_runs.py +++ b/backend/tests/api/routes/test_runs.py @@ -95,6 +95,8 @@ def test_a_run_cache_finds_what_an_earlier_run_recorded(tmp_path): assert cache.lookup(collected) is None assert cache.lookup("never-seen") is None assert cache.lookup("") is None + # A rename may have moved some of a row's names and not yet the others. + assert RunCache(store, renaming=lambda: True).lookup(plain) is None def test_what_may_be_stored_as_a_cache_entry(): diff --git a/backend/tests/api/test_rename_gate.py b/backend/tests/api/test_rename_gate.py new file mode 100644 index 0000000..ce740b7 --- /dev/null +++ b/backend/tests/api/test_rename_gate.py @@ -0,0 +1,193 @@ +"""The admission gate a rename closes (section 2.1 of +docs/private/id_refac_rename.md).""" + +import asyncio +import threading +import time + +import pytest +from fastapi.testclient import TestClient +from starlette.applications import Starlette +from starlette.requests import Request +from starlette.responses import PlainTextResponse +from starlette.routing import Route, WebSocketRoute +from starlette.websockets import WebSocket, WebSocketDisconnect + +from fluksio.core.config import settings +from fluksio.flow.rename import RenameGate +from fluksio.main import AdmissionGate, app +from fluksio.mcp import server + +API = settings.API_V1_STR + + +def _gated( + handler_gate: threading.Event | None = None, +) -> tuple[TestClient, RenameGate]: + """A bare app behind a gate of its own. A mutation waits for + *handler_gate*, when given, so it can be caught while admitted.""" + + async def answer(request: Request) -> PlainTextResponse: + if handler_gate is not None and request.method == "POST": + await asyncio.to_thread(handler_gate.wait) + return PlainTextResponse("ok") + + async def socket(websocket: WebSocket) -> None: + await websocket.accept() + await websocket.send_text("ok") + await websocket.close() + + methods = ["GET", "POST", "PUT", "DELETE"] + inner = Starlette( + routes=[ + Route("/{path:path}", answer, methods=methods), + WebSocketRoute("/hooks/{flow}/ws", socket), + ] + ) + gate = RenameGate() + return TestClient(AdmissionGate(inner, gate)), gate + + +def _close(gate: RenameGate, affected: list[str] | None = None) -> None: + assert gate.close("op-1", affected or [], threading.Event()) + + +def test_a_mutation_is_turned_away_while_closed_and_reads_pass(): + client, gate = _gated() + _close(gate) + + refused = client.post(f"{API}/flows/f_x/publish") + assert refused.status_code == 503 + assert refused.headers["retry-after"] == "5" + assert refused.json()["code"] == "renaming" + assert refused.json()["op"] == "op-1" + for method in ("put", "delete"): + assert getattr(client, method)(f"{API}/dashboards/d_x").status_code == 503 + assert client.post(f"{API}/messages/heat.temp").status_code == 503 + assert client.post(f"{API}/runs/flows/f_x").status_code == 503 + + assert client.get(f"{API}/flows/f_x").status_code == 200 + for path in ( + "/login/access-token", + "/users/me", + "/flows/ws-ticket", + "/ai/chat", + "/alerts/webpush/subscriptions", + "/artifacts", + "/flows/f_x/rename/plan", + "/dashboards/d_x/rename", + ): + assert client.post(f"{API}{path}").status_code == 200, path + + gate.open() + assert client.post(f"{API}/flows/f_x/publish").status_code == 200 + + +def test_an_admitted_request_is_drained_before_close_returns(): + release = threading.Event() + client, gate = _gated(release) + answered: list[int] = [] + request = threading.Thread( + target=lambda: answered.append(client.post(f"{API}/flows/").status_code) + ) + request.start() + while not gate._admitted: + time.sleep(0.01) + + heard: list[list[str]] = [] + closed = threading.Thread( + target=lambda: gate.close("op-1", [], threading.Event(), heard.append) + ) + closed.start() + closed.join(0.3) + assert closed.is_alive() + assert heard == [[f"POST {API}/flows/"]] + # Turned away meanwhile, while the one inside finishes. + assert client.post(f"{API}/flows/").status_code == 503 + + release.set() + closed.join(5) + request.join(5) + assert not closed.is_alive() + assert answered == [200] + assert gate.closed + + +def test_an_aborted_drain_reopens_the_gate(): + release = threading.Event() + client, gate = _gated(release) + request = threading.Thread(target=lambda: client.post(f"{API}/flows/")) + request.start() + while not gate._admitted: + time.sleep(0.01) + + abort = threading.Event() + result: list[bool] = [] + closed = threading.Thread( + target=lambda: result.append(gate.close("op-1", [], abort)) + ) + closed.start() + abort.set() + closed.join(5) + assert result == [False] + assert not gate.closed + release.set() + request.join(5) + assert client.post(f"{API}/flows/").status_code == 200 + + +def test_only_the_hooks_of_affected_flows_are_turned_away(): + client, gate = _gated() + _close(gate, ["f_a"]) + + assert client.get("/hooks/f_a/in").status_code == 503 + assert client.post("/hooks/f_a/in").status_code == 503 + assert client.post("/hooks/f_b/in").status_code == 200 + with pytest.raises(WebSocketDisconnect) as closed: + with client.websocket_connect("/hooks/f_a/ws") as socket: + socket.receive_text() + assert closed.value.code == 1013 + with client.websocket_connect("/hooks/f_b/ws") as socket: + assert socket.receive_text() == "ok" + + gate.open() + assert client.post("/hooks/f_a/in").status_code == 200 + + +def test_the_app_is_gated(client: TestClient, superuser_token_headers): + gate: RenameGate = app.state.rename_gate + _close(gate, ["f_x"]) + try: + assert client.get(f"{API}/flows/", headers=superuser_token_headers).is_success + made = client.post( + f"{API}/flows/", headers=superuser_token_headers, json={"name": "gated"} + ) + assert made.status_code == 503 + assert client.post( + f"{API}/login/access-token", + data={ + "username": settings.FIRST_SUPERUSER, + "password": settings.FIRST_SUPERUSER_PASSWORD, + }, + ).is_success + + async def tool_call() -> object: + previous = server._client + server.set_client(None) + server.connect(app) + token = server.caller_authorization.set( + superuser_token_headers["Authorization"] + ) + try: + return await server.create_flow({"name": "gated"}) + finally: + server.caller_authorization.reset(token) + await server.aclose() + server.set_client(previous) + + assert asyncio.run(tool_call()) == { + "error": "A rename is in progress; try again shortly", + "status": 503, + } + finally: + gate.open() diff --git a/backend/tests/flow/test_engine_hold.py b/backend/tests/flow/test_engine_hold.py index 2f6761f..c796d4c 100644 --- a/backend/tests/flow/test_engine_hold.py +++ b/backend/tests/flow/test_engine_hold.py @@ -11,6 +11,7 @@ import time import pytest from fluksio.flow import executor +from fluksio.flow.controller import FlowController, LoadedNode from fluksio.flow.executor import ExecutionService from fluksio.flow.messages import DType, MessageSpec from fluksio.flow.namespace import Rename @@ -18,6 +19,7 @@ from fluksio.flow.nodes import Node from fluksio.flow.pipeline import EngineHeld, Pipeline from fluksio.flow.queue import MemoryWorkQueue, WorkItem from fluksio.flow.state import MemoryState +from fluksio.flow.store import FlowStore def _engine(consume, queue=None) -> tuple[Pipeline, ExecutionService, Node]: @@ -214,6 +216,21 @@ def test_publish_is_refused_while_held(): assert seen == [2.0] +def test_a_trigger_by_hand_while_held_is_refused_not_failed(tmp_path, monkeypatch): + seen: list = [] + pipeline, _, source = _engine(seen.append) + controller = FlowController(FlowStore(tmp_path / "flows")) + controller.pipeline = pipeline + controller.loaded[source.id] = LoadedNode(id=source.id, flow="f", node=source) + failures: list = [] + monkeypatch.setattr(Pipeline, "publish_error", lambda *args: failures.append(args)) + assert pipeline.hold(threading.Event()) + + with pytest.raises(EngineHeld): + controller.trigger_node(source.id, {"reading": 1.0}) + assert failures == [] and seen == [] + + def test_work_arriving_during_a_hold_runs_after_it_in_order(): seen: list = [] pipeline, service, source = _engine(seen.append)