Lock on upsert (#95)

* Lock on upsert

* Release notes

Author: lsbardel

docs/release-notes.md

@@ -1,5 +1,27 @@
 # Release Notes
 
+## v2.2.6
+
+Adds tag filtering for task listings and fixes a race in the task database plugin.
+
+- Task listings can be filtered by tag: the `GET /tasks` endpoint and the `ls`
+  command of the [task CLI](https://fluid.quantmind.com/reference/task_cli/)
+  accept a repeatable `tags` option that returns only tasks carrying at least
+  one of the given tags, and `TaskInfo` now reports each task's tags.
+  ([#94](https://github.com/quantmind/aio-fluid/pull/94))
+- The [task database plugin](https://fluid.quantmind.com/reference/task_plugin/#fluid.scheduler.db.TaskDbPlugin)
+  now serialises its per-run lifecycle writes with a dedicated task-run lock.
+  `CrudDB.db_upsert` is not atomic — it issues an `UPDATE` and only `INSERT`s
+  when nothing matched — so when the scheduler wrote the `queued` row and a
+  consumer wrote the `running` row a few milliseconds later, the consumer's
+  `UPDATE` could miss the not-yet-committed `INSERT`, fall through to its own
+  `INSERT` and violate the task-runs primary key. Holding the lock around the
+  upsert removes the race.
+  ([#95](https://github.com/quantmind/aio-fluid/pull/95))
+- [TaskRun.lock](https://fluid.quantmind.com/reference/task_run/#fluid.scheduler.TaskRun.lock)
+  accepts an optional `name` to acquire a named sub-lock for the task run, and
+  `timeout` now defaults to `None`. ([#95](https://github.com/quantmind/aio-fluid/pull/95))
+
 ## v2.2.5
 
 - The [task decorator](https://fluid.quantmind.com/reference/task/#fluid.scheduler.task)

fluid/scheduler/db.py

@@ -125,19 +125,20 @@ async def get_run(self, run_id: str) -> TaskRunHistory:
     async def _handle_update(self, task_run: TaskRun) -> None:
         if self.skip_db_tag in task_run.task.tags:
             return
-        await self.db.db_upsert(
-            self.db.tables[self.table_name],
-            dict(id=task_run.id),
-            dict(
-                state=task_run.state,
-                name=task_run.name,
-                priority=task_run.priority,
-                queued=task_run.queued,
-                start=task_run.start,
-                end=task_run.end,
-                params=task_run.params.model_dump(mode="json"),
-            ),
-        )
+        async with task_run.lock(timeout=5, name="db_upsert"):
+            await self.db.db_upsert(
+                self.db.tables[self.table_name],
+                dict(id=task_run.id),
+                dict(
+                    state=task_run.state,
+                    name=task_run.name,
+                    priority=task_run.priority,
+                    queued=task_run.queued,
+                    start=task_run.start,
+                    end=task_run.end,
+                    params=task_run.params.model_dump(mode="json"),
+                ),
+            )
 
 
 def task_meta(meta: sa.MetaData, table_name: str = "tasks") -> None:

fluid/scheduler/models.py

@@ -469,9 +469,12 @@ def set_state(
             case _:
                 raise TaskRunError(f"invalid state transition {self.state} -> {state}")
 
-    def lock(self, timeout: float | None) -> Lock:
+    def lock(self, timeout: float | None = None, name: str | None = None) -> Lock:
         """Get a lock for this task run"""
-        return self.task_manager.broker.lock(f"tasks:{self.name}", timeout=timeout)
+        lock_name = f"tasks:{self.name}"
+        if name:
+            lock_name = f"{lock_name}:{name}"
+        return self.task_manager.broker.lock(lock_name, timeout=timeout)
 
     async def _execute(self) -> None:
         try:

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "aio-fluid"
-version = "2.2.5"
+version = "2.2.6"
 description = "Tools for backend python services"
 authors = [
     { name = "Luca Sbardella", email = "luca@quantmind.com" },