flowstash-runtime 0.9.0__tar.gz → 0.9.2__tar.gz

{flowstash_runtime-0.9.0 → flowstash_runtime-0.9.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: flowstash-runtime
-Version: 0.9.0
+Version: 0.9.2
 Summary: Actual runtime engine and bootstrap layer for the flowstash platform.
 Author: juraj.bezdek@gmail.com
 Author-email: juraj.bezdek@gmail.com
@@ -12,7 +12,8 @@ Classifier: Programming Language :: Python :: 3.13
 Requires-Dist: apscheduler (>=3.11.0)
 Requires-Dist: dramatiq (>=1.16.0)
 Requires-Dist: fastapi (>=0.110.0)
-Requires-Dist: flowstash-clients (>=0.9.0,<0.10.0)
-Requires-Dist: flowstash-lib (>=0.9.0,<0.10.0)
+Requires-Dist: flowstash-clients (>=0.9.2,<0.10.0)
+Requires-Dist: flowstash-lib (>=0.9.2,<0.10.0)
 Requires-Dist: pyyaml (>=6.0.1)
 Requires-Dist: uvicorn (>=0.29.0)
+Requires-Dist: websockets (>=14.0)

{flowstash_runtime-0.9.0 → flowstash_runtime-0.9.2}/pyproject.toml

@@ -1,17 +1,18 @@
 [project]
 name = "flowstash-runtime"
-version = "0.9.0"
+version = "0.9.2"
 description = "Actual runtime engine and bootstrap layer for the flowstash platform."
 authors = [{name = "juraj.bezdek@gmail.com", email = "juraj.bezdek@gmail.com"}]
 requires-python = ">=3.11"
 dependencies = [
-    "flowstash-clients>=0.9.0,<0.10.0",
-    "flowstash-lib>=0.9.0,<0.10.0",
+    "flowstash-clients>=0.9.2,<0.10.0",
+    "flowstash-lib>=0.9.2,<0.10.0",
     "fastapi>=0.110.0",
     "uvicorn>=0.29.0",
     "dramatiq>=1.16.0",
     "apscheduler>=3.11.0",
-    "pyyaml>=6.0.1"
+    "pyyaml>=6.0.1",
+    "websockets>=14.0"
 ]
 
 [build-system]

{flowstash_runtime-0.9.0 → flowstash_runtime-0.9.2}/src/flowstash/runtime/worker/backends/managed/http_entrypoint.py

@@ -34,6 +34,12 @@ from .task_resolver import (
     resolve_function as _registry_resolve,
     _invoke_task_callable,
 )
+from .lease_client import (
+    get_lease_client,
+    LeaseBusy,
+    SUCCEEDED as LEASE_SUCCEEDED,
+    FAILED as LEASE_FAILED,
+)
 from flowstash.pipelines.record_serialization import from_jsonable
 
 logger = logging.getLogger(__name__)
@@ -218,63 +224,91 @@ async def _execute_managed_task(payload: TaskPayload) -> dict:
     # Fall back to a fresh id for older payloads that don't carry one.
     execution_run_id = payload.run_id or str(uuid.uuid4())
 
-    with integration_context(
-        integration=integration,
-        integration_pipeline=pipeline,
-        run_id=execution_run_id,
-        parent_run_id=parent_run_id,
-        operation_id=operation_id,
-        tags=tags,
-        attrs={"args": raw_args},
-        record_lifecycle=False,
-    ) as ctx:
-        try:
-            func = _resolve_task_callable(
-                payload.task_name or payload.task_id, func_ref
+    # Idempotency guard: claim an exclusive lease on this run_id from the broker so a
+    # Cloud Tasks redelivery (or a concurrent duplicate) never re-executes a run that is
+    # already running or just completed. Disabled-safe: get_lease_client() returns None
+    # when the broker isn't configured; an unreachable broker yields UNAVAILABLE → we
+    # fail closed with a retryable 503 rather than risk a duplicate.
+    lease = get_lease_client()
+    lease_held = False
+    if lease is not None:
+        res = await lease.acquire(execution_run_id, entry_point=entry_point)
+        if res.duplicate:
+            logger.info(
+                "run %s already completed — skipping redelivery", execution_run_id
             )
-        except ValueError as e:
-            logger.warning(f"Could not resolve task '{entry_point}': {e}")
+            return {"status": "duplicate", "run_id": execution_run_id, "task": func_ref}
+        if not res.acquired:
+            # BUSY (held elsewhere), RECOVERING (broker restarting), or UNAVAILABLE
+            # (broker unreachable) — all map to a retryable 503. Only ACQUIRED runs.
+            raise LeaseBusy(execution_run_id)
+        lease_held = True
+
+    lease_status = LEASE_FAILED
+    try:
+        with integration_context(
+            integration=integration,
+            integration_pipeline=pipeline,
+            run_id=execution_run_id,
+            parent_run_id=parent_run_id,
+            operation_id=operation_id,
+            tags=tags,
+            attrs={"args": raw_args},
+            record_lifecycle=False,
+        ) as ctx:
+            try:
+                func = _resolve_task_callable(
+                    payload.task_name or payload.task_id, func_ref
+                )
+            except ValueError as e:
+                logger.warning(f"Could not resolve task '{entry_point}': {e}")
+                await record_run_started(
+                    correlation=ctx.corelation,
+                    entry_point=entry_point,
+                    attrs={"args": raw_args},
+                )
+                await record_run_ended(
+                    status="FAILED",
+                    correlation=ctx.corelation,
+                    attrs={"error": str(e), "task_resolution_failed": True},
+                )
+                raise
+
+            normalized_args = normalize_arguments(func, args, kwargs)
             await record_run_started(
                 correlation=ctx.corelation,
                 entry_point=entry_point,
-                attrs={"args": raw_args},
-            )
-            await record_run_ended(
-                status="FAILED",
-                correlation=ctx.corelation,
-                attrs={"error": str(e), "task_resolution_failed": True},
+                attrs={"args": normalized_args},
             )
-            raise
 
-        normalized_args = normalize_arguments(func, args, kwargs)
-        await record_run_started(
-            correlation=ctx.corelation,
-            entry_point=entry_point,
-            attrs={"args": normalized_args},
-        )
-
-        try:
-            await _invoke_task_callable(func, args, kwargs)
-        except Exception as e:
-            import traceback
+            try:
+                await _invoke_task_callable(func, args, kwargs)
+            except Exception as e:
+                import traceback
+
+                error = str(e)
+                tb = traceback.format_exc()
+                logger.error(f"Task execution failed: {func_ref}: {e}", exc_info=True)
+                await record_run_ended(
+                    status="FAILED",
+                    correlation=ctx.corelation,
+                    attrs={"error": error, "traceback": tb},
+                )
+                # Re-raise so the run ends as FAILED only (not ALSO SUCCEEDED) and the handler
+                # returns 500 → Cloud Tasks retries the delivery. Without this, a raising task
+                # was recorded as both FAILED and SUCCEEDED and the worker returned 200.
+                raise
 
-            error = str(e)
-            tb = traceback.format_exc()
-            logger.error(f"Task execution failed: {func_ref}: {e}", exc_info=True)
             await record_run_ended(
-                status="FAILED",
+                status="SUCCEEDED",
                 correlation=ctx.corelation,
-                attrs={"error": error, "traceback": tb},
             )
-            # Re-raise so the run ends as FAILED only (not ALSO SUCCEEDED) and the handler
-            # returns 500 → Cloud Tasks retries the delivery. Without this, a raising task
-            # was recorded as both FAILED and SUCCEEDED and the worker returned 200.
-            raise
-
-        await record_run_ended(
-            status="SUCCEEDED",
-            correlation=ctx.corelation,
-        )
+            lease_status = LEASE_SUCCEEDED
+    finally:
+        if lease_held:
+            # SUCCEEDED → completed-set tombstone (skips post-completion redeliveries);
+            # FAILED → free the lease so Cloud Tasks can genuinely retry.
+            await lease.release(execution_run_id, lease_status)
 
     return {
         "status": "ok",
@@ -317,6 +351,15 @@ async def handle_task(request: Request, payload: TaskPayload):
                 status_code=status.HTTP_404_NOT_FOUND,
                 content={"status": "TASK_NOT_FOUND", "detail": str(e)},
             )
+        except LeaseBusy as e:
+            # Lease held elsewhere (or broker unreachable). 503 is retryable by
+            # Cloud Tasks (unlike a 4xx), so the redelivery becomes the liveness
+            # backstop: by the next attempt the holder has finished (→ COMPLETED)
+            # or its lease expired (→ this attempt acquires it).
+            return JSONResponse(
+                status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
+                content={"status": "LEASE_HELD", "run_id": e.run_id},
+            )
         except Exception as e:
             # Task execution failed (already recorded as FAILED on the run). Return 500 so
             # Cloud Tasks retries the delivery; the request scope still releases its slot

flowstash_runtime-0.9.2/src/flowstash/runtime/worker/backends/managed/lease_client.py

@@ -0,0 +1,267 @@
+"""
+Worker-side client for the lease broker.
+
+One process-singleton WebSocket connection multiplexes the acquire/release of
+every task this worker process runs. The connection itself is the liveness
+signal — if it drops, the broker frees this worker's leases after a grace window
+unless we reconnect and re-assert them. We track our own held run_ids locally and
+re-assert the whole set on every (re)connect.
+
+Safe rollout: if the broker is not configured (no URL/token), or the optional
+``websockets`` dependency is missing, the client is *disabled* and callers run
+without the guard (no behaviour change). Once enabled, an unreachable broker
+makes ``acquire`` return ``UNAVAILABLE`` so the caller fails closed (HTTP 503 /
+job no-op) rather than risk a duplicate.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import logging
+import os
+from dataclasses import dataclass
+from typing import Dict, List, Optional, Set
+
+try:  # optional dependency — absence simply disables the guard
+    import websockets
+except Exception:  # pragma: no cover - import guard
+    websockets = None  # type: ignore
+
+logger = logging.getLogger(__name__)
+
+# acquire() outcomes
+ACQUIRED = "ACQUIRED"
+BUSY = "BUSY"
+COMPLETED = "COMPLETED"
+UNAVAILABLE = "UNAVAILABLE"
+
+# release() statuses
+SUCCEEDED = "SUCCEEDED"
+FAILED = "FAILED"
+
+
+class LeaseBusy(Exception):
+    """Raised when the run's lease is held elsewhere (or the broker is
+    unreachable) — the caller should respond with a retryable 503 so Cloud Tasks
+    redelivers later, by which point the holder has finished or its lease expired."""
+
+    def __init__(self, run_id: str):
+        super().__init__(f"lease busy for run_id={run_id}")
+        self.run_id = run_id
+
+
+@dataclass
+class AcquireResult:
+    outcome: str  # ACQUIRED | BUSY | COMPLETED | UNAVAILABLE
+
+    @property
+    def acquired(self) -> bool:
+        return self.outcome == ACQUIRED
+
+    @property
+    def duplicate(self) -> bool:
+        return self.outcome == COMPLETED
+
+    @property
+    def busy(self) -> bool:
+        return self.outcome in (BUSY, UNAVAILABLE)
+
+
+class LeaseBrokerClient:
+    def __init__(
+        self,
+        url: str,
+        token: str,
+        *,
+        acquire_timeout: float = 10.0,
+        reconnect_min: float = 0.5,
+        reconnect_max: float = 10.0,
+    ) -> None:
+        self._url = url
+        self._token = token
+        self._acquire_timeout = acquire_timeout
+        self._reconnect_min = reconnect_min
+        self._reconnect_max = reconnect_max
+
+        self._held: Set[str] = set()
+        self._waiters: Dict[str, List[asyncio.Future]] = {}
+        self._ws = None
+        self._connected = asyncio.Event()
+        self._send_lock = asyncio.Lock()
+        self._run_task: Optional[asyncio.Task] = None
+        self._closing = False
+
+    # ── lifecycle ────────────────────────────────────────────────────────
+
+    def ensure_started(self) -> None:
+        """Start the background connect loop once, on the current event loop."""
+        if self._run_task is None or self._run_task.done():
+            self._run_task = asyncio.ensure_future(self._run())
+
+    async def stop(self) -> None:
+        self._closing = True
+        if self._run_task is not None:
+            self._run_task.cancel()
+        ws = self._ws
+        if ws is not None:
+            try:
+                await ws.close()
+            except Exception:
+                pass
+
+    async def _run(self) -> None:
+        backoff = self._reconnect_min
+        headers = {"Authorization": f"Bearer {self._token}"}
+        while not self._closing:
+            try:
+                async with websockets.connect(
+                    self._url, additional_headers=headers, open_timeout=10
+                ) as ws:
+                    self._ws = ws
+                    backoff = self._reconnect_min
+                    if self._held:
+                        # Re-assert everything we still believe we own.
+                        await self._send({"op": "reassert", "run_ids": list(self._held)})
+                    self._connected.set()
+                    async for raw in ws:
+                        try:
+                            self._on_message(json.loads(raw))
+                        except Exception:
+                            logger.exception("lease client: bad message %r", raw)
+            except asyncio.CancelledError:
+                raise
+            except Exception as e:
+                logger.warning("lease broker connection lost: %s", e)
+            finally:
+                self._connected.clear()
+                self._ws = None
+            if self._closing:
+                break
+            await asyncio.sleep(backoff)
+            backoff = min(backoff * 2, self._reconnect_max)
+
+    # ── messaging ──────────────────────────────────────────────────────────
+
+    def _on_message(self, msg: dict) -> None:
+        op = msg.get("op")
+        if op in ("granted", "refused"):
+            run_id = msg.get("run_id")
+            outcome = ACQUIRED if op == "granted" else msg.get("reason", BUSY)
+            self._resolve(run_id, outcome)
+        elif op == "revoked":
+            # Another worker took over (split-brain after a long stall). We can't
+            # force-cancel in-flight work; drop ownership and rely on step
+            # idempotency. Logged loudly so it's visible.
+            run_id = msg.get("run_id")
+            self._held.discard(run_id)
+            logger.error("lease REVOKED run_id=%s reason=%s", run_id, msg.get("reason"))
+        elif op == "ping":
+            asyncio.ensure_future(self._send_safe({"op": "pong"}))
+
+    def _resolve(self, run_id: Optional[str], outcome: str) -> None:
+        for fut in self._waiters.pop(run_id, []):
+            if not fut.done():
+                fut.set_result(outcome)
+
+    async def _send(self, msg: dict) -> None:
+        async with self._send_lock:
+            if self._ws is None:
+                raise ConnectionError("lease broker not connected")
+            await self._ws.send(json.dumps(msg))
+
+    async def _send_safe(self, msg: dict) -> None:
+        try:
+            await self._send(msg)
+        except Exception:
+            pass
+
+    # ── public API ──────────────────────────────────────────────────────────
+
+    async def acquire(
+        self, run_id: str, entry_point: Optional[str] = None, timeout: Optional[float] = None
+    ) -> AcquireResult:
+        self.ensure_started()
+        timeout = timeout or self._acquire_timeout
+        try:
+            await asyncio.wait_for(self._connected.wait(), timeout=timeout)
+        except asyncio.TimeoutError:
+            return AcquireResult(UNAVAILABLE)
+
+        fut: asyncio.Future = asyncio.get_event_loop().create_future()
+        self._waiters.setdefault(run_id, []).append(fut)
+        try:
+            await self._send({"op": "acquire", "run_id": run_id, "entry_point": entry_point})
+        except Exception:
+            self._discard_waiter(run_id, fut)
+            return AcquireResult(UNAVAILABLE)
+
+        try:
+            outcome = await asyncio.wait_for(fut, timeout=timeout)
+        except asyncio.TimeoutError:
+            self._discard_waiter(run_id, fut)
+            return AcquireResult(UNAVAILABLE)
+
+        if outcome == ACQUIRED:
+            self._held.add(run_id)
+        return AcquireResult(outcome)
+
+    async def release(self, run_id: str, status: str = SUCCEEDED) -> None:
+        self._held.discard(run_id)
+        # Best-effort: if the send fails, the broker frees the lease via the
+        # grace window when the connection drops anyway.
+        await self._send_safe({"op": "release", "run_id": run_id, "status": status})
+
+    def _discard_waiter(self, run_id: str, fut: asyncio.Future) -> None:
+        waiters = self._waiters.get(run_id)
+        if waiters and fut in waiters:
+            waiters.remove(fut)
+            if not waiters:
+                self._waiters.pop(run_id, None)
+
+
+# ── singleton wiring ─────────────────────────────────────────────────────────
+
+_client: Optional[LeaseBrokerClient] = None
+_resolved = False
+
+
+def _derive_broker_url() -> Optional[str]:
+    explicit = os.getenv("LEASE_BROKER_URL")
+    if explicit:
+        return explicit
+    api = os.getenv("FLOWSTASH_API_URL") or os.getenv("MANAGED_API_URL")
+    if not api:
+        return None
+    api = api.rstrip("/")
+    if api.startswith("https://"):
+        return "wss://" + api[len("https://") :] + "/ws/leases"
+    if api.startswith("http://"):
+        return "ws://" + api[len("http://") :] + "/ws/leases"
+    return None
+
+
+def get_lease_client() -> Optional[LeaseBrokerClient]:
+    """Return the process-singleton lease client, or None when the guard is
+    disabled (broker not configured, explicitly turned off, or ``websockets``
+    not installed). Callers treat None as 'run without the guard'."""
+    global _client, _resolved
+    if _resolved:
+        return _client
+    _resolved = True
+
+    if os.getenv("LEASE_BROKER_ENABLED", "true").lower() in ("0", "false", "no"):
+        logger.info("lease broker disabled via LEASE_BROKER_ENABLED")
+        return None
+    if websockets is None:
+        logger.warning("lease broker disabled: 'websockets' not installed")
+        return None
+    url = _derive_broker_url()
+    token = os.getenv("MANAGED_AUTH_TOKEN", "")
+    if not url or not token:
+        logger.info("lease broker disabled: URL/token not configured")
+        return None
+
+    _client = LeaseBrokerClient(url, token)
+    logger.info("lease broker enabled: %s", url)
+    return _client

{flowstash_runtime-0.9.0 → flowstash_runtime-0.9.2}/src/flowstash/runtime/worker/backends/managed/managed_consumer.py

@@ -14,6 +14,7 @@ import logging
 import os
 import sys
 import traceback
+import uuid
 from typing import Any
 
 from collections import defaultdict
@@ -32,6 +33,14 @@ from flowstash.observability.ingestion import (
 from flowstash.queue.consumer import TaskConsumer
 
 from .task_resolver import _invoke_task_callable, resolve_function
+from .lease_client import (
+    get_lease_client,
+    ACQUIRED,
+    BUSY,
+    COMPLETED,
+    SUCCEEDED,
+    FAILED,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -146,13 +155,16 @@ async def _record_missing_task(task_name: str) -> None:
         )
 
 
-async def _run_job_task(func: Any, task_name: str, args: list, kwargs: dict) -> bool:
+async def _run_job_task(
+    func: Any, task_name: str, args: list, kwargs: dict, run_id: str | None = None
+) -> bool:
     """Execute the resolved task within an integration_context. Returns True on success."""
     normalized_args = normalize_arguments(func, args, kwargs)
 
     with integration_context(
         integration="managed-job",
         integration_pipeline=task_name,
+        run_id=run_id,
         record_lifecycle=False,
     ) as ctx:
         await record_run_started(
@@ -388,7 +400,11 @@ class ManagedConsumer(TaskConsumer):
             sys.exit(1)
 
         task_name, args, kwargs = _parse_cli_args(argv[1:])
-        logger.info(f"[run-task] Resolving task: {task_name!r}")
+        # Stable run_id injected by the platform at schedule time (FLOWSTASH_RUN_ID).
+        # It is identical across Cloud Run Job retries of one execution, so the lease
+        # guard collapses duplicate executions of the same logical run.
+        run_id = os.getenv("FLOWSTASH_RUN_ID") or str(uuid.uuid4())
+        logger.info(f"[run-task] Resolving task: {task_name!r} (run_id={run_id})")
 
         try:
             func = resolve_function(task_name)
@@ -398,10 +414,41 @@ class ManagedConsumer(TaskConsumer):
             await asyncio.to_thread(AsyncManager.get_instance().flush, 15.0)
             sys.exit(1)
 
+        # Idempotency guard via the lease broker. Unlike HTTP tasks, a job that
+        # exits has no automatic redelivery, so on BUSY/COMPLETED we exit 0 (another
+        # execution owns it / it is already done) and on UNAVAILABLE we fail OPEN
+        # (run anyway) rather than silently drop the work.
+        lease = get_lease_client()
+        lease_held = False
+        if lease is not None:
+            res = await lease.acquire(run_id, entry_point=task_name)
+            if res.outcome in (COMPLETED, BUSY):
+                logger.info(
+                    "[run-task] run %s not started (%s) — another execution owns it",
+                    run_id,
+                    res.outcome,
+                )
+                await asyncio.to_thread(AsyncManager.get_instance().flush, 15.0)
+                sys.exit(0)
+            elif res.outcome == ACQUIRED:
+                lease_held = True
+            else:
+                # UNAVAILABLE (broker unreachable) or RECOVERING (broker restarting).
+                # A job has no automatic redelivery, so fail OPEN (run unguarded)
+                # rather than silently drop the work; duplicates here are rare.
+                logger.warning(
+                    "[run-task] lease not granted (%s) — running run %s without guard",
+                    res.outcome,
+                    run_id,
+                )
+
         logger.info(
             f"[run-task] Executing task: {task_name!r}, args={args}, kwargs={kwargs}"
         )
-        success = await _run_job_task(func, task_name, args, kwargs)
+        success = await _run_job_task(func, task_name, args, kwargs, run_id=run_id)
+
+        if lease_held:
+            await lease.release(run_id, SUCCEEDED if success else FAILED)
 
         await asyncio.to_thread(AsyncManager.get_instance().flush, 15.0)
         sys.exit(0 if success else 1)

{flowstash_runtime-0.9.0 → flowstash_runtime-0.9.2}/src/flowstash/runtime/worker/backends/managed/task_resolver.py

@@ -4,6 +4,7 @@ Task registry, function resolution, and callable dispatch.
 Shared by the managed job consumer.
 """
 
+import asyncio
 import importlib
 import inspect
 import logging
@@ -48,14 +49,23 @@ def resolve_function(func_ref: str) -> Any:
 
 
 async def _invoke_task_callable(func: Any, args: list, kwargs: dict) -> Any:
-    """Dispatch a task callable regardless of its wrapper type."""
+    """Dispatch a task callable regardless of its wrapper type.
+
+    Sync callables are run via ``asyncio.to_thread`` rather than inline: a
+    multi-minute sync task would otherwise block the single worker event loop and
+    starve all other concurrent requests on the instance — the very condition that
+    triggers Cloud Tasks' connection resets/retries — and would also freeze the
+    lease WebSocket. ``asyncio.to_thread`` copies the current ``contextvars``
+    context into the worker thread, so the ``integration_context`` / ``run_id``
+    remain visible to the task and its observability.
+    """
     if hasattr(func, "run"):
         return await func.run(*args, **kwargs)
     if hasattr(func, "func"):
         underlying = func.func
         if inspect.iscoroutinefunction(underlying):
             return await underlying(*args, **kwargs)
-        return underlying(*args, **kwargs)
+        return await asyncio.to_thread(underlying, *args, **kwargs)
     if inspect.iscoroutinefunction(func):
         return await func(*args, **kwargs)
-    return func(*args, **kwargs)
+    return await asyncio.to_thread(func, *args, **kwargs)