fixing a race condition in celery task wrapper. could randomly blow up any task. (#2321)

backend/danswer/background/task_utils.py

@@ -93,9 +93,16 @@ def build_apply_async_wrapper(build_name_fn: Callable[..., str]) -> Callable[[AA
             kwargs_for_build_name = kwargs or {}
             task_name = build_name_fn(*args_for_build_name, **kwargs_for_build_name)
             with Session(get_sqlalchemy_engine()) as db_session:
-                # mark the task as started
+                # register_task must come before fn = apply_async or else the task
+                # might run mark_task_start (and crash) before the task row exists
+                db_task = register_task(task_name, db_session)
+
                 task = fn(args, kwargs, *other_args, **other_kwargs)
-                register_task(task.id, task_name, db_session)
+
+                # we update the celery task id for diagnostic purposes
+                # but it isn't currently used by any code
+                db_task.task_id = task.id
+                db_session.commit()
 
             return task
 

backend/danswer/db/models.py

@@ -1409,7 +1409,7 @@ class TaskQueueState(Base):
     __tablename__ = "task_queue_jobs"
 
     id: Mapped[int] = mapped_column(primary_key=True)
-    # Celery task id
+    # Celery task id. currently only for readability/diagnostics
     task_id: Mapped[str] = mapped_column(String)
     # For any job type, this would be the same
     task_name: Mapped[str] = mapped_column(String)

backend/danswer/db/tasks.py

@@ -44,12 +44,11 @@ def get_latest_task_by_type(
 
 
 def register_task(
-    task_id: str,
     task_name: str,
     db_session: Session,
 ) -> TaskQueueState:
     new_task = TaskQueueState(
-        task_id=task_id, task_name=task_name, status=TaskStatus.PENDING
+        task_id="", task_name=task_name, status=TaskStatus.PENDING
     )
 
     db_session.add(new_task)