From ca422132c8f739f3ee873046a3f59d80f2462152 Mon Sep 17 00:00:00 2001
From: Vika <vika@fireburn.ru>
Date: Sun, 9 Jul 2023 18:25:58 +0300
Subject: webmentions/queue: give up on an item after 5 attempts

This also involves a crude "async drop" implementation that fires a
future incrementing an attempt if a Job has been dropped without
marking it as done.
---
 kittybox-rs/src/webmentions/queue.rs | 76 ++++++++++++++++++++++++++++--------
 1 file changed, 59 insertions(+), 17 deletions(-)

(limited to 'kittybox-rs/src/webmentions/queue.rs')

diff --git a/kittybox-rs/src/webmentions/queue.rs b/kittybox-rs/src/webmentions/queue.rs
index 77ad4ea..b585f58 100644
--- a/kittybox-rs/src/webmentions/queue.rs
+++ b/kittybox-rs/src/webmentions/queue.rs
@@ -9,7 +9,7 @@ use super::Webmention;
 static MIGRATOR: sqlx::migrate::Migrator = sqlx::migrate!();
 
 #[async_trait::async_trait]
-pub trait JobQueue<T: Send + Sync + Sized>: Send + Sync + Sized + Clone + 'static {
+pub trait JobQueue<T: JobItem>: Send + Sync + Sized + Clone + 'static {
     type Job: Job<T, Self>;
     type Error: std::error::Error + Send + Sync + Sized;
 
@@ -20,31 +20,70 @@ pub trait JobQueue<T: Send + Sync + Sized>: Send + Sync + Sized + Clone + 'stati
 }
 
 #[async_trait::async_trait]
-pub trait Job<T: Send + Sync + Sized, Q: JobQueue<T>>: Send + Sync + Sized {
+pub trait Job<T: JobItem, Q: JobQueue<T>>: Send + Sync + Sized {
     fn job(&self) -> &T;
     async fn done(self) -> Result<(), Q::Error>;
 }
 
+pub trait JobItem: Send + Sync + Sized + std::fmt::Debug {
+    const DATABASE_NAME: &'static str;
+}
+
 #[derive(Debug)]
-pub struct PostgresJobItem<'c, T> {
+pub struct PostgresJobItem<T: JobItem> {
     id: Uuid,
     job: T,
-    txn: sqlx::Transaction<'c, sqlx::Postgres>
+    // This will normally always be Some, except on drop
+    txn: Option<sqlx::Transaction<'static, sqlx::Postgres>>,
+    runtime_handle: tokio::runtime::Handle,
+}
+
+
+impl<T: JobItem> Drop for PostgresJobItem<T> {
+    // This is an emulation of "async drop" — the struct retains a
+    // runtime handle, which it uses to block on a future that does
+    // the actual cleanup.
+    //
+    // Of course, this is not portable between runtimes, but I don't
+    // care about that, since Kittybox is designed to work within the
+    // Tokio ecosystem.
+    fn drop(&mut self) {
+        tracing::error!("Job {:?} failed, incrementing attempts...", &self);
+        if let Some(mut txn) = self.txn.take() {
+            let id = self.id;
+            self.runtime_handle.spawn(async move {
+                tracing::debug!("Constructing query to increment attempts for job {}...", id);
+                // UPDATE "T::DATABASE_NAME" WHERE id = $1 SET attempts = attempts + 1
+                sqlx::query_builder::QueryBuilder::new("UPDATE ")
+                    // This is safe from a SQL injection standpoint, since it is a constant.
+                    .push(T::DATABASE_NAME)
+                    .push(" SET attempts = attempts + 1")
+                    .push(" WHERE id = ")
+                    .push_bind(id)
+                    .build()
+                    .execute(&mut txn)
+                    .await
+                    .unwrap();
+
+                txn.commit().await.unwrap();
+            });
+        }
+    }
 }
 
 #[async_trait::async_trait]
-impl Job<Webmention, PostgresJobQueue<Webmention>> for PostgresJobItem<'_, Webmention> {
+impl Job<Webmention, PostgresJobQueue<Webmention>> for PostgresJobItem<Webmention> {
     fn job(&self) -> &Webmention {
         &self.job
     }
     async fn done(mut self) -> Result<(), <PostgresJobQueue<Webmention> as JobQueue<Webmention>>::Error> {
-        println!("Deleting {} from the job queue", self.id);
-        sqlx::query("DELETE FROM incoming_webmention_queue WHERE id = $1")
+        tracing::debug!("Deleting {} from the job queue", self.id);
+        sqlx::query("DELETE FROM kittybox.incoming_webmention_queue WHERE id = $1")
             .bind(self.id)
-            .execute(&mut self.txn)
+            .execute(self.txn.as_mut().unwrap())
             .await?;
 
-        self.txn.commit().await
+        self.txn.take().unwrap().commit().await
     }
 }
 
@@ -87,29 +126,32 @@ impl PostgresJobQueue<Webmention> {
 
 #[async_trait::async_trait]
 impl JobQueue<Webmention> for PostgresJobQueue<Webmention> {
-    type Job = PostgresJobItem<'static, Webmention>;
+    type Job = PostgresJobItem<Webmention>;
     type Error = sqlx::Error;
 
     async fn get_one(&self) -> Result<Option<Self::Job>, Self::Error> {
         let mut txn = self.db.begin().await?;
 
         match sqlx::query_as::<_, (Uuid, String, String)>(
-            "SELECT id, source, target FROM incoming_webmention_queue FOR UPDATE SKIP LOCKED LIMIT 1"
+            "SELECT id, source, target FROM kittybox.incoming_webmention_queue WHERE attempts < 5 FOR UPDATE SKIP LOCKED LIMIT 1"
         )
             .fetch_optional(&mut txn)
             .await?
             .map(|(id, source, target)| (id, Webmention { source, target })) {
-                Some((id, webmention)) => Ok(Some(Self::Job {
-                    id,
-                    job: webmention,
-                    txn
-                })),
+                Some((id, webmention)) => {
+                    return Ok(Some(Self::Job {
+                        id,
+                        job: webmention,
+                        txn: Some(txn),
+                        runtime_handle: tokio::runtime::Handle::current(),
+                    }))
+                },
                 None => Ok(None)
             }
     }
 
     async fn put(&self, item: &Webmention) -> Result<Uuid, Self::Error> {
-        sqlx::query_scalar::<_, Uuid>("INSERT INTO incoming_webmention_queue (source, target) VALUES ($1, $2) RETURNING id")
+        sqlx::query_scalar::<_, Uuid>("INSERT INTO kittybox.incoming_webmention_queue (source, target) VALUES ($1, $2) RETURNING id")
             .bind(item.source.as_str())
             .bind(item.target.as_str())
             .fetch_one(&self.db)
-- 
cgit 1.4.1