Merge pull request #1658 from mcass19/feature/scheduled-jobs

[FEATURE] Scheduled Jobs

Author: Mark-Simulacrum

Cargo.lock

@@ -23,13 +23,13 @@ hyper = { version = "0.14.4", features = ["server", "stream"]}
 tokio = { version = "1.7.1", features = ["macros", "time", "rt"] }
 futures = { version = "0.3", default-features = false, features = ["std"] }
 async-trait = "0.1.31"
-uuid = { version = "0.8", features = ["v4"] }
+uuid = { version = "0.8", features = ["v4", "serde"] }
 tracing = "0.1"
 tracing-subscriber = { version = "0.3", features = ["env-filter"] }
 url = "2.1.0"
 once_cell = "1"
 chrono = { version = "0.4", features = ["serde"] }
-tokio-postgres = { version = "0.7.2", features = ["with-chrono-0_4", "with-serde_json-1"] }
+tokio-postgres = { version = "0.7.2", features = ["with-chrono-0_4", "with-serde_json-1", "with-uuid-0_8"] }
 postgres-native-tls = "0.5.0"
 native-tls = "0.2"
 serde_path_to_error = "0.1.2"
@@ -42,6 +42,8 @@ tower = { version = "0.4.13", features = ["util", "limit", "buffer", "load-shed"
 github-graphql = { path = "github-graphql" }
 rand = "0.8.5"
 ignore = "0.4.18"
+postgres-types = { version = "0.2.4", features = ["derive"] }
+cron = { version = "0.12.0" }
 
 [dependencies.serde]
 version = "1"

Cargo.toml

@@ -1,11 +1,15 @@
+use crate::db::jobs::*;
+use crate::handlers::jobs::handle_job;
 use anyhow::Context as _;
+use chrono::Utc;
 use native_tls::{Certificate, TlsConnector};
 use postgres_native_tls::MakeTlsConnector;
 use std::sync::{Arc, Mutex};
 use tokio::sync::{OwnedSemaphorePermit, Semaphore};
 use tokio_postgres::Client as DbClient;
 
 pub mod issue_data;
+pub mod jobs;
 pub mod notifications;
 pub mod rustc_commits;
 
@@ -179,6 +183,43 @@ pub async fn run_migrations(client: &DbClient) -> anyhow::Result<()> {
     Ok(())
 }
 
+pub async fn schedule_jobs(db: &DbClient, jobs: Vec<JobSchedule>) -> anyhow::Result<()> {
+    for job in jobs {
+        let mut upcoming = job.schedule.upcoming(Utc).take(1);
+
+        if let Some(scheduled_at) = upcoming.next() {
+            if let Err(_) = get_job_by_name_and_scheduled_at(&db, &job.name, &scheduled_at).await {
+                // mean there's no job already in the db with that name and scheduled_at
+                insert_job(&db, &job.name, &scheduled_at, &job.metadata).await?;
+            }
+        }
+    }
+
+    Ok(())
+}
+
+pub async fn run_scheduled_jobs(db: &DbClient) -> anyhow::Result<()> {
+    let jobs = get_jobs_to_execute(&db).await.unwrap();
+    tracing::trace!("jobs to execute: {:#?}", jobs);
+
+    for job in jobs.iter() {
+        update_job_executed_at(&db, &job.id).await?;
+
+        match handle_job(&job.name, &job.metadata).await {
+            Ok(_) => {
+                tracing::trace!("job successfully executed (id={})", job.id);
+                delete_job(&db, &job.id).await?;
+            }
+            Err(e) => {
+                tracing::trace!("job failed on execution (id={:?}, error={:?})", job.id, e);
+                update_job_error_message(&db, &job.id, &e.to_string()).await?;
+            }
+        }
+    }
+
+    Ok(())
+}
+
 static MIGRATIONS: &[&str] = &[
     "
 CREATE TABLE notifications (
@@ -215,5 +256,21 @@ CREATE TABLE issue_data (
     data JSONB,
     PRIMARY KEY (repo, issue_number, key)
 );
+",
+    "
+CREATE TABLE jobs (
+    id UUID DEFAULT gen_random_uuid() PRIMARY KEY,
+    name TEXT NOT NULL,
+    scheduled_at TIMESTAMP WITH TIME ZONE NOT NULL,
+    metadata JSONB,
+    executed_at TIMESTAMP WITH TIME ZONE,
+    error_message TEXT
+);
+",
+    "
+CREATE UNIQUE INDEX jobs_name_scheduled_at_unique_index 
+    ON jobs (
+        name, scheduled_at
+    );
 ",
 ];

src/db.rs

@@ -0,0 +1,137 @@
+//! The `jobs` table provides a way to have scheduled jobs
+use anyhow::{Context as _, Result};
+use chrono::{DateTime, Utc};
+use cron::Schedule;
+use serde::{Deserialize, Serialize};
+use tokio_postgres::Client as DbClient;
+use uuid::Uuid;
+
+pub struct JobSchedule {
+    pub name: String,
+    pub schedule: Schedule,
+    pub metadata: serde_json::Value,
+}
+
+#[derive(Serialize, Deserialize, Debug)]
+pub struct Job {
+    pub id: Uuid,
+    pub name: String,
+    pub scheduled_at: DateTime<Utc>,
+    pub metadata: serde_json::Value,
+    pub executed_at: Option<DateTime<Utc>>,
+    pub error_message: Option<String>,
+}
+
+pub async fn insert_job(
+    db: &DbClient,
+    name: &String,
+    scheduled_at: &DateTime<Utc>,
+    metadata: &serde_json::Value,
+) -> Result<()> {
+    tracing::trace!("insert_job(name={})", name);
+
+    db.execute(
+        "INSERT INTO jobs (name, scheduled_at, metadata) VALUES ($1, $2, $3) 
+            ON CONFLICT (name, scheduled_at) DO UPDATE SET metadata = EXCLUDED.metadata",
+        &[&name, &scheduled_at, &metadata],
+    )
+    .await
+    .context("Inserting job")?;
+
+    Ok(())
+}
+
+pub async fn delete_job(db: &DbClient, id: &Uuid) -> Result<()> {
+    tracing::trace!("delete_job(id={})", id);
+
+    db.execute("DELETE FROM jobs WHERE id = $1", &[&id])
+        .await
+        .context("Deleting job")?;
+
+    Ok(())
+}
+
+pub async fn update_job_error_message(db: &DbClient, id: &Uuid, message: &String) -> Result<()> {
+    tracing::trace!("update_job_error_message(id={})", id);
+
+    db.execute(
+        "UPDATE jobs SET error_message = $2 WHERE id = $1",
+        &[&id, &message],
+    )
+    .await
+    .context("Updating job error message")?;
+
+    Ok(())
+}
+
+pub async fn update_job_executed_at(db: &DbClient, id: &Uuid) -> Result<()> {
+    tracing::trace!("update_job_executed_at(id={})", id);
+
+    db.execute("UPDATE jobs SET executed_at = now() WHERE id = $1", &[&id])
+        .await
+        .context("Updating job executed at")?;
+
+    Ok(())
+}
+
+pub async fn get_job_by_name_and_scheduled_at(
+    db: &DbClient,
+    name: &String,
+    scheduled_at: &DateTime<Utc>,
+) -> Result<Job> {
+    tracing::trace!(
+        "get_job_by_name_and_scheduled_at(name={}, scheduled_at={})",
+        name,
+        scheduled_at
+    );
+
+    let job = db
+        .query_one(
+            "SELECT * FROM jobs WHERE name = $1 AND scheduled_at = $2",
+            &[&name, &scheduled_at],
+        )
+        .await
+        .context("Select job by name and scheduled at")?;
+
+    deserialize_job(&job)
+}
+
+// Selects all jobs with:
+//  - scheduled_at in the past
+//  - error_message is null or executed_at is at least 60 minutes ago (intended to make repeat executions rare enough)
+pub async fn get_jobs_to_execute(db: &DbClient) -> Result<Vec<Job>> {
+    let jobs = db
+        .query(
+            "
+        SELECT * FROM jobs WHERE scheduled_at <= now() AND (error_message IS NULL OR executed_at <= now() - INTERVAL '60 minutes')",
+            &[],
+        )
+        .await
+        .context("Getting jobs data")?;
+
+    let mut data = Vec::with_capacity(jobs.len());
+    for job in jobs {
+        let serialized_job = deserialize_job(&job);
+        data.push(serialized_job.unwrap());
+    }
+
+    Ok(data)
+}
+
+fn deserialize_job(row: &tokio_postgres::row::Row) -> Result<Job> {
+    let id: Uuid = row.try_get(0)?;
+    let name: String = row.try_get(1)?;
+    let scheduled_at: DateTime<Utc> = row.try_get(2)?;
+    let metadata: serde_json::Value = row.try_get(3)?;
+    let executed_at: Option<DateTime<Utc>> = row.try_get(4)?;
+    let error_message: Option<String> = row.try_get(5)?;
+
+    Ok(Job {
+        id,
+        name,
+        scheduled_at,
+        metadata,
+        executed_at,
+        error_message,
+    })
+}

src/db/jobs.rs

@@ -28,6 +28,7 @@ mod autolabel;
 mod close;
 mod github_releases;
 mod glacier;
+pub mod jobs;
 mod major_change;
 mod mentions;
 mod milestone_prs;

src/handlers.rs

@@ -0,0 +1,21 @@
+// Function to match the scheduled job function with its corresponding handler.
+// In case you want to add a new one, just add a new clause to the match with
+// the job name and the corresponding function.
+
+// Further info could be find in src/jobs.rs
+
+pub async fn handle_job(name: &String, metadata: &serde_json::Value) -> anyhow::Result<()> {
+    match name {
+        _ => default(&name, &metadata),
+    }
+}
+
+fn default(name: &String, metadata: &serde_json::Value) -> anyhow::Result<()> {
+    tracing::trace!(
+        "handle_job fell into default case: (name={:?}, metadata={:?})",
+        name,
+        metadata
+    );
+
+    Ok(())
+}