Cron-like recurring jobs

Author: sandhose

Cargo.lock

@@ -102,6 +102,10 @@ features = ["serde", "clock"]
 version = "4.5.21"
 features = ["derive"]
 
+# Cron expressions
+[workspace.dependencies.cron]
+version = "0.13.0"
+
 # Elliptic curve cryptography
 [workspace.dependencies.elliptic-curve]
 version = "0.13.8"

Cargo.toml

@@ -0,0 +1,29 @@
+-- Copyright 2024 New Vector Ltd.
+--
+-- SPDX-License-Identifier: AGPL-3.0-only
+-- Please see LICENSE in the repository root for full details.
+
+-- Add a table to track the state of scheduled recurring jobs.
+CREATE TABLE queue_schedules (
+    -- A unique name for the schedule
+    schedule_name TEXT PRIMARY KEY,
+
+    -- The cron expression to use to schedule the job. This is there just for
+    -- convenience, as this is defined by the backend
+    schedule_expression TEXT NOT NULL,
+
+    -- The last time the job was scheduled. If NULL, it means that the job was
+    -- never scheduled.
+    last_scheduled_at TIMESTAMP WITH TIME ZONE,
+
+    -- The job that was scheduled last time. If NULL, it means that either the
+    -- job was never scheduled, or the job cleaned up from the database
+    last_scheduled_job_id UUID
+        REFERENCES queue_jobs (queue_job_id)
+);
+
+-- When a job is scheduled from a recurreing schedule, we keep a column
+-- referencing the name of the schedule
+ALTER TABLE queue_jobs
+    ADD COLUMN schedule_name TEXT
+        REFERENCES queue_schedules (schedule_name);

crates/storage-pg/migrations/20241125110803_queue_job_recurrent.sql

@@ -12,8 +12,10 @@ use mas_storage::{
     queue::{Job, QueueJobRepository, Worker},
     Clock,
 };
+use opentelemetry_semantic_conventions::trace::DB_QUERY_TEXT;
 use rand::RngCore;
 use sqlx::PgConnection;
+use tracing::Instrument;
 use ulid::Ulid;
 use uuid::Uuid;
 
@@ -137,6 +139,7 @@ impl QueueJobRepository for PgQueueJobRepository<'_> {
         payload: serde_json::Value,
         metadata: serde_json::Value,
         scheduled_at: DateTime<Utc>,
+        schedule_name: Option<&str>,
     ) -> Result<(), Self::Error> {
         let created_at = clock.now();
         let id = Ulid::from_datetime_with_source(created_at.into(), rng);
@@ -145,20 +148,47 @@ impl QueueJobRepository for PgQueueJobRepository<'_> {
         sqlx::query!(
             r#"
                 INSERT INTO queue_jobs
-                    (queue_job_id, queue_name, payload, metadata, created_at, scheduled_at, status)
-                VALUES ($1, $2, $3, $4, $5, $6, 'scheduled')
+                    (queue_job_id, queue_name, payload, metadata, created_at, scheduled_at, schedule_name, status)
+                VALUES ($1, $2, $3, $4, $5, $6, $7, 'scheduled')
             "#,
             Uuid::from(id),
             queue_name,
             payload,
             metadata,
             created_at,
             scheduled_at,
+            schedule_name,
         )
         .traced()
         .execute(&mut *self.conn)
         .await?;
 
+        // If there was a schedule name supplied, update the queue_schedules table
+        if let Some(schedule_name) = schedule_name {
+            let span = tracing::info_span!(
+                "db.queue_job.schedule_later.update_schedules",
+                { DB_QUERY_TEXT } = tracing::field::Empty,
+            );
+
+            let res = sqlx::query!(
+                r#"
+                    UPDATE queue_schedules
+                    SET last_scheduled_at = $1,
+                        last_scheduled_job_id = $2
+                    WHERE schedule_name = $3
+                "#,
+                scheduled_at,
+                Uuid::from(id),
+                schedule_name,
+            )
+            .record(&span)
+            .execute(&mut *self.conn)
+            .instrument(span)
+            .await?;
+
+            DatabaseError::ensure_affected_rows(&res, 1)?;
+        }
+
         Ok(())
     }
 
@@ -315,14 +345,19 @@ impl QueueJobRepository for PgQueueJobRepository<'_> {
         let scheduled_at = now + delay;
         let new_id = Ulid::from_datetime_with_source(now.into(), rng);
 
+        let span = tracing::info_span!(
+            "db.queue_job.retry.insert_job",
+            { DB_QUERY_TEXT } = tracing::field::Empty
+        );
         // Create a new job with the same payload and metadata, but a new ID and
         // increment the attempt
         // We make sure we do this only for 'failed' jobs
         let res = sqlx::query!(
             r#"
                 INSERT INTO queue_jobs
-                    (queue_job_id, queue_name, payload, metadata, created_at, attempt, scheduled_at, status)
-                SELECT $1, queue_name, payload, metadata, $2, attempt + 1, $3, 'scheduled'
+                    (queue_job_id, queue_name, payload, metadata, created_at,
+                     attempt, scheduled_at, schedule_name, status)
+                SELECT $1, queue_name, payload, metadata, $2, attempt + 1, $3, schedule_name, 'scheduled'
                 FROM queue_jobs
                 WHERE queue_job_id = $4
                   AND status = 'failed'
@@ -332,13 +367,39 @@ impl QueueJobRepository for PgQueueJobRepository<'_> {
             scheduled_at,
             Uuid::from(id),
         )
-        .traced()
+        .record(&span)
         .execute(&mut *self.conn)
+        .instrument(span)
         .await?;
 
         DatabaseError::ensure_affected_rows(&res, 1)?;
 
+        // If that job was referenced by a schedule, update the schedule
+        let span = tracing::info_span!(
+            "db.queue_job.retry.update_schedule",
+            { DB_QUERY_TEXT } = tracing::field::Empty
+        );
+        sqlx::query!(
+            r#"
+                UPDATE queue_schedules
+                SET last_scheduled_at = $1,
+                    last_scheduled_job_id = $2
+                WHERE last_scheduled_job_id = $3
+            "#,
+            scheduled_at,
+            Uuid::from(new_id),
+            Uuid::from(id),
+        )
+        .record(&span)
+        .execute(&mut *self.conn)
+        .instrument(span)
+        .await?;
+
         // Update the old job to point to the new attempt
+        let span = tracing::info_span!(
+            "db.queue_job.retry.update_old_job",
+            { DB_QUERY_TEXT } = tracing::field::Empty
+        );
         let res = sqlx::query!(
             r#"
                 UPDATE queue_jobs
@@ -348,8 +409,9 @@ impl QueueJobRepository for PgQueueJobRepository<'_> {
             Uuid::from(new_id),
             Uuid::from(id),
         )
-        .traced()
+        .record(&span)
         .execute(&mut *self.conn)
+        .instrument(span)
         .await?;
 
         DatabaseError::ensure_affected_rows(&res, 1)?;

crates/storage-pg/src/queue/job.rs

@@ -6,4 +6,5 @@
 //! A module containing the PostgreSQL implementation of the job queue
 
 pub mod job;
+pub mod schedule;
 pub mod worker;

crates/storage-pg/src/queue/mod.rs

@@ -0,0 +1,91 @@
+// Copyright 2024 New Vector Ltd.
+//
+// SPDX-License-Identifier: AGPL-3.0-only
+// Please see LICENSE in the repository root for full details.
+
+//! A module containing the PostgreSQL implementation of the
+//! [`QueueScheduleRepository`].
+
+use async_trait::async_trait;
+use chrono::{DateTime, Utc};
+use mas_storage::queue::{QueueScheduleRepository, Schedule, ScheduleStatus};
+use sqlx::PgConnection;
+
+use crate::{DatabaseError, ExecuteExt};
+
+/// An implementation of [`QueueScheduleRepository`] for a PostgreSQL
+/// connection.
+pub struct PgQueueScheduleRepository<'c> {
+    conn: &'c mut PgConnection,
+}
+
+impl<'c> PgQueueScheduleRepository<'c> {
+    /// Create a new [`PgQueueScheduleRepository`] from an active PostgreSQL
+    /// connection.
+    #[must_use]
+    pub fn new(conn: &'c mut PgConnection) -> Self {
+        Self { conn }
+    }
+}
+
+struct ScheduleLookup {
+    schedule_name: String,
+    last_scheduled_at: Option<DateTime<Utc>>,
+    last_scheduled_job_completed: Option<bool>,
+}
+
+impl From<ScheduleLookup> for ScheduleStatus {
+    fn from(value: ScheduleLookup) -> Self {
+        ScheduleStatus {
+            schedule_name: value.schedule_name,
+            last_scheduled_at: value.last_scheduled_at,
+            last_scheduled_job_completed: value.last_scheduled_job_completed,
+        }
+    }
+}
+
+#[async_trait]
+impl<'c> QueueScheduleRepository for PgQueueScheduleRepository<'c> {
+    type Error = DatabaseError;
+
+    async fn setup(&mut self, schedules: &[(&'static str, Schedule)]) -> Result<(), Self::Error> {
+        sqlx::query!(
+            r#"
+                INSERT INTO queue_schedules (schedule_name, schedule_expression)
+                SELECT * FROM UNNEST($1::text[], $2::text[]) AS t (schedule_name, schedule_expression)
+                ON CONFLICT (schedule_name) DO UPDATE
+                    SET schedule_expression = EXCLUDED.schedule_expression
+            "#,
+            &schedules.iter().map(|(name, _)| (*name).to_owned()).collect::<Vec<_>>(),
+            &schedules
+                .iter()
+                .map(|(_, schedule)| schedule.source().to_owned())
+                .collect::<Vec<_>>()
+        )
+        .traced()
+        .execute(&mut *self.conn)
+        .await?;
+
+        Ok(())
+    }
+
+    async fn list(&mut self) -> Result<Vec<ScheduleStatus>, Self::Error> {
+        let res = sqlx::query_as!(
+            ScheduleLookup,
+            r#"
+                SELECT
+                    queue_schedules.schedule_name as "schedule_name!",
+                    queue_schedules.last_scheduled_at,
+                    queue_jobs.status IN ('completed', 'failed') as last_scheduled_job_completed
+                FROM queue_schedules
+                LEFT JOIN queue_jobs
+                    ON queue_jobs.queue_job_id = queue_schedules.last_scheduled_job_id
+            "#
+        )
+        .traced()
+        .fetch_all(&mut *self.conn)
+        .await?;
+
+        Ok(res.into_iter().map(Into::into).collect())
+    }
+}

crates/storage-pg/src/queue/schedule.rs

@@ -17,6 +17,7 @@ use mas_storage::{
         OAuth2AccessTokenRepository, OAuth2AuthorizationGrantRepository, OAuth2ClientRepository,
         OAuth2DeviceCodeGrantRepository, OAuth2RefreshTokenRepository, OAuth2SessionRepository,
     },
+    queue::{QueueJobRepository, QueueScheduleRepository, QueueWorkerRepository},
     upstream_oauth2::{
         UpstreamOAuthLinkRepository, UpstreamOAuthProviderRepository,
         UpstreamOAuthSessionRepository,
@@ -38,7 +39,10 @@ use crate::{
         PgOAuth2ClientRepository, PgOAuth2DeviceCodeGrantRepository,
         PgOAuth2RefreshTokenRepository, PgOAuth2SessionRepository,
     },
-    queue::{job::PgQueueJobRepository, worker::PgQueueWorkerRepository},
+    queue::{
+        job::PgQueueJobRepository, schedule::PgQueueScheduleRepository,
+        worker::PgQueueWorkerRepository,
+    },
     upstream_oauth2::{
         PgUpstreamOAuthLinkRepository, PgUpstreamOAuthProviderRepository,
         PgUpstreamOAuthSessionRepository,
@@ -259,15 +263,17 @@ where
         Box::new(PgCompatRefreshTokenRepository::new(self.conn.as_mut()))
     }
 
-    fn queue_worker<'c>(
-        &'c mut self,
-    ) -> Box<dyn mas_storage::queue::QueueWorkerRepository<Error = Self::Error> + 'c> {
+    fn queue_worker<'c>(&'c mut self) -> Box<dyn QueueWorkerRepository<Error = Self::Error> + 'c> {
         Box::new(PgQueueWorkerRepository::new(self.conn.as_mut()))
     }
 
-    fn queue_job<'c>(
-        &'c mut self,
-    ) -> Box<dyn mas_storage::queue::QueueJobRepository<Error = Self::Error> + 'c> {
+    fn queue_job<'c>(&'c mut self) -> Box<dyn QueueJobRepository<Error = Self::Error> + 'c> {
         Box::new(PgQueueJobRepository::new(self.conn.as_mut()))
     }
+
+    fn queue_schedule<'c>(
+        &'c mut self,
+    ) -> Box<dyn QueueScheduleRepository<Error = Self::Error> + 'c> {
+        Box::new(PgQueueScheduleRepository::new(self.conn.as_mut()))
+    }
 }

crates/storage-pg/src/repository.rs

@@ -14,6 +14,7 @@ workspace = true
 [dependencies]
 async-trait.workspace = true
 chrono.workspace = true
+cron.workspace = true
 futures-util.workspace = true
 opentelemetry.workspace = true
 rand_core = "0.6.4"