Lock jobs before starting processing

Author: augustuswm

rfd-model/migrations/2023-11-02-161428_lock_job/down.sql

@@ -0,0 +1 @@
+ALTER TABLE job DROP COLUMN started_at;

rfd-model/migrations/2023-11-02-161428_lock_job/up.sql

@@ -0,0 +1 @@
+ALTER TABLE job ADD COLUMN started_at TIMESTAMPTZ;

rfd-model/src/db.rs

@@ -72,6 +72,7 @@ pub struct JobModel {
     pub committed_at: DateTime<Utc>,
     pub processed: bool,
     pub created_at: DateTime<Utc>,
+    pub started_at: Option<DateTime<Utc>>,
 }
 
 #[derive(Debug, Deserialize, Serialize, Queryable, Insertable)]

rfd-model/src/lib.rs

@@ -145,6 +145,8 @@ pub struct Job {
     pub processed: bool,
     #[partial(NewJob(skip))]
     pub created_at: DateTime<Utc>,
+    #[partial(NewJob(skip))]
+    pub started_at: Option<DateTime<Utc>>,
 }
 
 impl From<JobModel> for Job {
@@ -160,6 +162,7 @@ impl From<JobModel> for Job {
             committed_at: value.committed_at,
             processed: value.processed,
             created_at: value.created_at,
+            started_at: value.started_at,
         }
     }
 }

rfd-model/src/schema.rs

@@ -88,6 +88,7 @@ diesel::table! {
         committed_at -> Timestamptz,
         processed -> Bool,
         created_at -> Timestamptz,
+        started_at -> Nullable<Timestamptz>,
     }
 }
 

rfd-model/src/storage/mod.rs

@@ -221,6 +221,7 @@ pub struct JobFilter {
     pub id: Option<Vec<i32>>,
     pub sha: Option<Vec<String>>,
     pub processed: Option<bool>,
+    pub started: Option<bool>,
 }
 
 impl JobFilter {
@@ -238,6 +239,11 @@ impl JobFilter {
         self.processed = processed;
         self
     }
+
+    pub fn started(mut self, started: Option<bool>) -> Self {
+        self.started = started;
+        self
+    }
 }
 
 #[cfg_attr(feature = "mock", automock)]
@@ -250,6 +256,7 @@ pub trait JobStore {
         pagination: &ListPagination,
     ) -> Result<Vec<Job>, StoreError>;
     async fn upsert(&self, new_job: NewJob) -> Result<Job, StoreError>;
+    async fn start(&self, id: i32) -> Result<Option<Job>, StoreError>;
     async fn complete(&self, id: i32) -> Result<Option<Job>, StoreError>;
 }
 

rfd-model/src/storage/postgres.rs

@@ -430,7 +430,7 @@ impl JobStore for PostgresStore {
     ) -> Result<Vec<Job>, StoreError> {
         let mut query = job::dsl::job.into_boxed();
 
-        let JobFilter { id, sha, processed } = filter;
+        let JobFilter { id, sha, processed, started } = filter;
 
         if let Some(id) = id {
             query = query.filter(job::id.eq_any(id));
@@ -444,6 +444,14 @@ impl JobStore for PostgresStore {
             query = query.filter(job::processed.eq(processed));
         }
 
+        if let Some(started) = started {
+            if started {
+                query = query.filter(job::started_at.is_null());
+            } else {
+                query = query.filter(job::started_at.is_not_null());
+            }
+        }
+
         let results = query
             .offset(pagination.offset)
             .limit(pagination.limit)
@@ -474,6 +482,17 @@ impl JobStore for PostgresStore {
         Ok(rfd.into())
     }
 
+    async fn start(&self, id: i32) -> Result<Option<Job>, StoreError> {
+        let _ = update(job::dsl::job)
+            .filter(job::id.eq(id))
+            .filter(job::started_at.is_null())
+            .set(job::started_at.eq(Utc::now()))
+            .execute_async(&*self.pool.get().await?)
+            .await?;
+
+        JobStore::get(self, id).await
+    }
+
     async fn complete(&self, id: i32) -> Result<Option<Job>, StoreError> {
         let _ = update(job::dsl::job)
             .filter(job::id.eq(id))

rfd-processor/src/processor.rs

@@ -25,42 +25,54 @@ pub async fn processor(ctx: Arc<Context>) -> Result<(), JobError> {
     loop {
         let jobs = JobStore::list(
             &ctx.db.storage,
-            JobFilter::default().processed(Some(false)),
+            JobFilter::default().processed(Some(false)).started(Some(false)),
             &pagination,
         )
         .await?;
 
         for job in jobs {
             let ctx = ctx.clone();
             tokio::spawn(async move {
-                let location = GitHubRfdLocation {
-                    owner: job.owner.clone(),
-                    repo: job.repository.clone(),
-                    branch: job.branch.clone(),
-                    commit: job.sha.clone(),
-                    default_branch: ctx.github.repository.default_branch.clone(),
-                };
+                // Make the job as started
+                match JobStore::start(&ctx.db.storage, job.id).await {
+                    Ok(Some(job)) => {
 
-                let update = GitHubRfdUpdate {
-                    location,
-                    number: job.rfd.into(),
-                    committed_at: job.committed_at,
-                };
+                        let location = GitHubRfdLocation {
+                            owner: job.owner.clone(),
+                            repo: job.repository.clone(),
+                            branch: job.branch.clone(),
+                            commit: job.sha.clone(),
+                            default_branch: ctx.github.repository.default_branch.clone(),
+                        };
 
-                let updater = RfdUpdater::new(&ctx.actions, ctx.processor.update_mode);
+                        let update = GitHubRfdUpdate {
+                            location,
+                            number: job.rfd.into(),
+                            committed_at: job.committed_at,
+                        };
 
-                match updater.handle(&ctx, &[update]).await {
-                    Ok(_) => {
-                        let _ = JobStore::complete(&ctx.db.storage, job.id)
-                            .await
-                            .tap_err(|err| {
-                                tracing::error!(?err, "Failed to mark job as completed")
-                            });
+                        let updater = RfdUpdater::new(&ctx.actions, ctx.processor.update_mode);
+
+                        match updater.handle(&ctx, &[update]).await {
+                            Ok(_) => {
+                                let _ = JobStore::complete(&ctx.db.storage, job.id)
+                                    .await
+                                    .tap_err(|err| {
+                                        tracing::error!(?err, "Failed to mark job as completed")
+                                    });
+                            }
+                            Err(err) => {
+                                tracing::error!(?err, "RFD update failed");
+
+                                // TODO: Mark job as failed or retry?
+                            }
+                        }
+                    }
+                    Ok(None) => {
+                        tracing::error!(?job, "Job that was scheduled to run has gone missing! Was it started by a different task?");
                     }
                     Err(err) => {
-                        tracing::error!(?err, "RFD update failed");
-
-                        // TODO: Mark job as failed or retry?
+                        tracing::warn!(?job, ?err, "Failed to start job. Was it previously started?");
                     }
                 }