diff --git a/src/backend/base/langflow/services/jobs/service.py b/src/backend/base/langflow/services/jobs/service.py index b535fda534..166c837f2b 100644 --- a/src/backend/base/langflow/services/jobs/service.py +++ b/src/backend/base/langflow/services/jobs/service.py @@ -343,6 +343,73 @@ class JobService(Service): await session.flush() return len(rows) + async def heartbeat(self, job_id: UUID, owner: str) -> None: + """Stamp the running owner + a fresh heartbeat on the job row. + + This is the liveness signal a reconciler reads to tell a live in-flight + run from a genuinely orphaned one: only the running owner refreshes it, + and a reconciler only fails/requeues a job whose heartbeat is STALE + (``is_lease_stale``). Stored in ``job_metadata`` (no new column, matching + the attempt-accounting decision in the design) via a shallow merge so the + persisted request / retry flags are preserved. + """ + now = datetime.now(timezone.utc).isoformat() + await self.update_job_metadata(job_id, {"owner": owner, "heartbeat_at": now}) + + @staticmethod + def is_lease_stale(job: Job, *, lease_ttl_s: float) -> bool: + """True when a job's heartbeat is older than ``lease_ttl_s`` (or absent). + + An absent heartbeat means the owner never recorded liveness (it died in + the QUEUED->IN_PROGRESS window, or this is a legacy row), so it is + treated as stale and reconcilable. A fresh heartbeat (within the TTL) + means a live owner is running the job and a reconciler must NOT touch it. + """ + meta = job.job_metadata or {} + raw = meta.get("heartbeat_at") + if not raw: + return True + try: + hb = datetime.fromisoformat(raw) + except (TypeError, ValueError): + return True + if hb.tzinfo is None: + hb = hb.replace(tzinfo=timezone.utc) + age = (datetime.now(timezone.utc) - hb).total_seconds() + return age > lease_ttl_s + + async def increment_attempt_if(self, job_id: UUID, *, expected: int, new: int) -> bool: + """Atomically bump ``job_metadata.attempt`` from ``expected`` to ``new``. + + Returns True only for the single caller whose read-modify-write observed + ``attempt == expected`` and committed. Concurrent reconcilers reading the + same ``expected`` race on the row, but UNIQUE row identity + a re-read + under the write lock means only one transaction's conditional UPDATE + matches: every other sees ``rowcount == 0`` and returns False. This + closes the lost-update window where two reconcilers both bumped 1->2 and + each requeued, pushing a job past ``max_attempts``. + + Portable across SQLite and Postgres: the conditional UPDATE matches the + JSON-extracted attempt cast to integer, the same single-row-conditional + primitive ``claim_queued_job`` relies on. + """ + from sqlalchemy import Integer + from sqlalchemy import cast as sa_cast + from sqlmodel import update + + attempt_expr = sa_cast(col(Job.job_metadata)["attempt"].as_string(), Integer) + async with session_scope() as session: + # Read the current metadata so we can write back a full merged blob + # (JSON columns are replaced wholesale, not patched in place). + job = await session.get(Job, job_id) + if job is None: + return False + merged = {**(job.job_metadata or {}), "attempt": new} + stmt = update(Job).where(Job.job_id == job_id, attempt_expr == expected).values(job_metadata=merged) + result = await session.exec(stmt) # type: ignore[call-overload] + await session.flush() + return result.rowcount == 1 + async def claim_queued_job(self, job_id: UUID) -> bool: """Atomically claim a QUEUED job for execution. Returns True if we won. diff --git a/src/backend/tests/unit/services/jobs/test_jobs_liveness.py b/src/backend/tests/unit/services/jobs/test_jobs_liveness.py new file mode 100644 index 0000000000..d7d3e50c37 --- /dev/null +++ b/src/backend/tests/unit/services/jobs/test_jobs_liveness.py @@ -0,0 +1,112 @@ +"""Store-layer liveness primitives: heartbeat, stale predicate, atomic attempt. + +These back the lease+heartbeat reconciliation model. A reconciler must only +fail/requeue a job whose heartbeat is STALE (older than the lease TTL), never a +fresh one, and the retry-safe attempt bump must be atomic so two concurrent +reconcilers cannot push a job past max_attempts. +""" + +from __future__ import annotations + +import asyncio +from datetime import datetime, timedelta, timezone +from uuid import uuid4 + +import pytest +from langflow.services.jobs.service import JobService + + +@pytest.mark.usefixtures("client") +async def test_heartbeat_records_owner_and_timestamp(): + service = JobService() + job_id = uuid4() + await service.create_job(job_id=job_id, flow_id=uuid4(), user_id=uuid4()) + + before = datetime.now(timezone.utc) + await service.heartbeat(job_id, owner="worker-A") + after = datetime.now(timezone.utc) + + job = await service.get_job_by_job_id(job_id) + meta = job.job_metadata or {} + assert meta.get("owner") == "worker-A" + hb = datetime.fromisoformat(meta["heartbeat_at"]) + assert before <= hb <= after + + +@pytest.mark.usefixtures("client") +async def test_heartbeat_preserves_other_metadata(): + service = JobService() + job_id = uuid4() + await service.create_job(job_id=job_id, flow_id=uuid4(), user_id=uuid4()) + await service.update_job_metadata(job_id, {"request": {"input_value": "hi"}, "retry_safe": True}) + + await service.heartbeat(job_id, owner="worker-A") + + job = await service.get_job_by_job_id(job_id) + meta = job.job_metadata or {} + # Heartbeat must not clobber the persisted request / retry flag. + assert meta["request"] == {"input_value": "hi"} + assert meta["retry_safe"] is True + assert meta["owner"] == "worker-A" + + +@pytest.mark.usefixtures("client") +async def test_is_lease_stale_fresh_vs_stale(): + service = JobService() + fresh = uuid4() + stale = uuid4() + absent = uuid4() + await service.create_job(job_id=fresh, flow_id=uuid4(), user_id=uuid4()) + await service.create_job(job_id=stale, flow_id=uuid4(), user_id=uuid4()) + await service.create_job(job_id=absent, flow_id=uuid4(), user_id=uuid4()) + + await service.heartbeat(fresh, owner="w1") + # Backdate the stale job's heartbeat well past the TTL. + old = (datetime.now(timezone.utc) - timedelta(seconds=120)).isoformat() + await service.update_job_metadata(stale, {"owner": "w1", "heartbeat_at": old}) + + fresh_job = await service.get_job_by_job_id(fresh) + stale_job = await service.get_job_by_job_id(stale) + absent_job = await service.get_job_by_job_id(absent) + + assert service.is_lease_stale(fresh_job, lease_ttl_s=30.0) is False + assert service.is_lease_stale(stale_job, lease_ttl_s=30.0) is True + # No heartbeat ever recorded -> treated as stale (truly orphaned / never ran). + assert service.is_lease_stale(absent_job, lease_ttl_s=30.0) is True + + +@pytest.mark.usefixtures("client") +async def test_increment_attempt_atomic_is_conditional(): + """Atomic attempt bump only succeeds when it observes the expected value. + + This is the optimistic guard that prevents two concurrent reconcilers from + both reading attempt=1 and both writing attempt=2 (a lost update). + """ + service = JobService() + job_id = uuid4() + await service.create_job(job_id=job_id, flow_id=uuid4(), user_id=uuid4()) + await service.update_job_metadata(job_id, {"attempt": 1, "max_attempts": 3}) + + # First bump from 1 -> 2 wins. + assert await service.increment_attempt_if(job_id, expected=1, new=2) is True + job = await service.get_job_by_job_id(job_id) + assert job.job_metadata["attempt"] == 2 + + # A racer that still thinks attempt==1 must NOT bump it again. + assert await service.increment_attempt_if(job_id, expected=1, new=2) is False + job = await service.get_job_by_job_id(job_id) + assert job.job_metadata["attempt"] == 2 + + +@pytest.mark.usefixtures("client") +async def test_increment_attempt_concurrent_only_one_wins(): + """Under concurrent bumps from the same expected value exactly one wins.""" + service = JobService() + job_id = uuid4() + await service.create_job(job_id=job_id, flow_id=uuid4(), user_id=uuid4()) + await service.update_job_metadata(job_id, {"attempt": 1, "max_attempts": 5}) + + results = await asyncio.gather(*(service.increment_attempt_if(job_id, expected=1, new=2) for _ in range(8))) + assert results.count(True) == 1 + job = await service.get_job_by_job_id(job_id) + assert job.job_metadata["attempt"] == 2