feat: adding dead letter processing and publishing timeout (#208)

* feat: adding dead letter processing and publishing timeout

* chore: adding give up check, removing unnecesary error

* chore(tests): adding test_dead_letters

* chore(tests): adding listening for the pg_notify to the test

* chore: changing timeouts to the Duration

* chore: resolving merge conflict on a new updates

* chore: choring comments, removing unnecessary match

Author: geekbrother

src/services/publisher_service/helpers.rs

@@ -3,6 +3,7 @@ use {
     crate::{metrics::Metrics, model::types::AccountId, types::Notification},
     relay_rpc::domain::{ProjectId, Topic},
     sqlx::{FromRow, PgPool, Postgres},
+    std::time::Duration,
     tracing::{error, instrument},
     uuid::Uuid,
     wc::metrics::otel::Context,
@@ -226,3 +227,43 @@ pub async fn update_metrics_on_queue_stats(metrics: &Metrics, postgres: &PgPool)
         }
     }
 }
+
+/// Checks for messages in the `processing` state for more than threshold
+/// and put it back in a `queued` state for processing
+#[instrument(skip(postgres))]
+pub async fn dead_letters_check(
+    threshold: Duration,
+    postgres: &PgPool,
+) -> std::result::Result<(), sqlx::error::Error> {
+    let update_status_query = "
+        UPDATE subscriber_notification
+        SET status = 'queued'
+        WHERE status = 'processing'
+        AND EXTRACT(EPOCH FROM (NOW() - updated_at)) > $1::INTEGER
+    ";
+    sqlx::query::<Postgres>(update_status_query)
+        .bind(threshold.as_secs() as i64)
+        .execute(postgres)
+        .await?;
+    Ok(())
+}
+
+/// Checks for message is created more than threshold
+#[instrument(skip(postgres))]
+pub async fn dead_letter_give_up_check(
+    notification: Uuid,
+    threshold: Duration,
+    postgres: &PgPool,
+) -> std::result::Result<bool, sqlx::error::Error> {
+    let query_to_check = "
+        SELECT now() - created_at > interval '$1 seconds' 
+        FROM subscriber_notification 
+        WHERE id = $2
+    ";
+    let row: (bool,) = sqlx::query_as(query_to_check)
+        .bind(threshold.as_secs() as i64)
+        .bind(notification)
+        .fetch_one(postgres)
+        .await?;
+    Ok(row.0)
+}

src/services/publisher_service/mod.rs

@@ -11,18 +11,22 @@ use {
         types::{Envelope, EnvelopeType0},
     },
     base64::Engine,
-    helpers::update_message_processing_status,
+    helpers::{dead_letter_give_up_check, update_message_processing_status},
     relay_client::http::Client,
     relay_rpc::{
         domain::{ClientId, DecodedClientId, Topic},
         rpc::{msg_id::MsgId, Publish, JSON_RPC_VERSION_STR},
     },
     sqlx::{postgres::PgListener, PgPool},
-    std::sync::{
-        atomic::{AtomicUsize, Ordering},
-        Arc,
+    std::{
+        sync::{
+            atomic::{AtomicUsize, Ordering},
+            Arc,
+        },
+        time::Duration,
     },
-    tracing::{info, instrument, warn},
+    tokio::time::{interval, timeout},
+    tracing::{error, info, instrument, warn},
     types::SubscriberNotificationStatus,
     wc::metrics::otel::Context,
 };
@@ -31,12 +35,19 @@ pub mod helpers;
 pub mod types;
 
 // TODO: These should be configurable, add to the config
-// Maximum of the parallel messages processing workers
+/// Maximum of the parallel messages processing workers
 const MAX_WORKERS: usize = 10;
-// Number of workers to be spawned on the service start to clean the queue
+/// Number of workers to be spawned on the service start to clean the queue
 const START_WORKERS: usize = 10;
 // Messages queue stats observing database polling interval
 const QUEUE_STATS_POLLING_INTERVAL: std::time::Duration = std::time::Duration::from_secs(60);
+/// Maximum publishing time before the publish will be considered as failed
+/// and the messages in queue with the `processing` state will be returned to the queue
+const PUBLISHING_TIMEOUT: Duration = Duration::from_secs(60 * 5); // 5 minutes
+/// Interval to check for dead letters
+const DEAD_LETTER_POLL_INTERVAL: Duration = Duration::from_secs(60);
+/// Total maximum time to process the message before it will be considered as failed
+const PUBLISHING_GIVE_UP_TIMEOUT: Duration = Duration::from_secs(60 * 60 * 24); // One day
 
 #[instrument(skip_all)]
 pub async fn start(
@@ -59,7 +70,20 @@ pub async fn start(
         });
     }
 
-    // TODO: Spawned tasks counter should be exported to metrics
+    // Spawning dead letters check task
+    tokio::spawn({
+        let postgres = postgres.clone();
+        async move {
+            let mut poll_interval = interval(DEAD_LETTER_POLL_INTERVAL);
+            loop {
+                poll_interval.tick().await;
+                if let Err(e) = helpers::dead_letters_check(PUBLISHING_TIMEOUT, &postgres).await {
+                    warn!("Error on dead letters check: {:?}", e);
+                }
+            }
+        }
+    });
+
     let spawned_tasks_counter = Arc::new(AtomicUsize::new(0));
 
     // Spawning initial workers to process messages from the queue in case
@@ -176,16 +200,36 @@ async fn process_queued_messages(
         if let Some(notification) = result {
             let notification_id = notification.subscriber_notification;
             info!("Got a notification with id: {}", notification_id);
-            process_notification(notification, relay_http_client.clone(), metrics, analytics)
-                .await?;
 
-            update_message_processing_status(
-                notification_id,
-                SubscriberNotificationStatus::Published,
-                postgres,
+            let process_result = process_with_timeout(
+                PUBLISHING_TIMEOUT,
+                notification,
+                relay_http_client.clone(),
                 metrics,
-            )
-            .await?;
+                analytics,
+            );
+
+            match process_result.await {
+                Ok(_) => {
+                    update_message_processing_status(
+                        notification_id,
+                        SubscriberNotificationStatus::Published,
+                        postgres,
+                        metrics,
+                    )
+                    .await?;
+                }
+                Err(e) => {
+                    warn!("Error on `process_notification`: {:?}", e);
+                    update_message_status_queued_or_failed(
+                        notification_id,
+                        PUBLISHING_GIVE_UP_TIMEOUT,
+                        postgres,
+                        metrics,
+                    )
+                    .await?;
+                }
+            }
         } else {
             info!("No more notifications to process, stopping the loop");
             break;
@@ -194,6 +238,31 @@ async fn process_queued_messages(
     Ok(())
 }
 
+/// Process publishing with the threshold timeout
+#[instrument(skip(relay_http_client, metrics, analytics, notification))]
+async fn process_with_timeout(
+    execution_threshold: Duration,
+    notification: NotificationToProcess,
+    relay_http_client: Arc<Client>,
+    metrics: Option<&Metrics>,
+    analytics: &NotifyAnalytics,
+) -> crate::error::Result<()> {
+    match timeout(
+        execution_threshold,
+        process_notification(notification, relay_http_client.clone(), metrics, analytics),
+    )
+    .await
+    {
+        Ok(result) => {
+            result?;
+        }
+        Err(e) => {
+            return Err(crate::error::Error::TokioTimeElapsed(e));
+        }
+    };
+    Ok(())
+}
+
 #[instrument(skip_all, fields(notification = ?notification))]
 async fn process_notification(
     notification: NotificationToProcess,
@@ -263,3 +332,34 @@ async fn process_notification(
 
     Ok(())
 }
+
+/// Updates message status back to `Queued` or mark as `Failed`
+/// depending on the `giveup_threshold` messsage creation time
+#[instrument(skip(postgres, metrics))]
+async fn update_message_status_queued_or_failed(
+    notification_id: uuid::Uuid,
+    giveup_threshold: Duration,
+    postgres: &PgPool,
+    metrics: Option<&Metrics>,
+) -> crate::error::Result<()> {
+    if dead_letter_give_up_check(notification_id, giveup_threshold, postgres).await? {
+        error!("Message was not processed during the giving up threshold, marking it as failed");
+        update_message_processing_status(
+            notification_id,
+            SubscriberNotificationStatus::Failed,
+            postgres,
+            metrics,
+        )
+        .await?;
+    } else {
+        update_message_processing_status(
+            notification_id,
+            SubscriberNotificationStatus::Queued,
+            postgres,
+            metrics,
+        )
+        .await?;
+    }
+
+    Ok(())
+}

src/services/publisher_service/types.rs

@@ -8,6 +8,7 @@ pub enum SubscriberNotificationStatus {
     Queued,
     Processing,
     Published,
+    Failed,
     NotSubscribed,
     WrongScope,
     RateLimited,
@@ -19,6 +20,7 @@ impl fmt::Display for SubscriberNotificationStatus {
             SubscriberNotificationStatus::Queued => write!(f, "queued"),
             SubscriberNotificationStatus::Processing => write!(f, "processing"),
             SubscriberNotificationStatus::Published => write!(f, "published"),
+            SubscriberNotificationStatus::Failed => write!(f, "failed"),
             SubscriberNotificationStatus::NotSubscribed => write!(f, "not-subscribed"),
             SubscriberNotificationStatus::WrongScope => write!(f, "wrong-scope"),
             SubscriberNotificationStatus::RateLimited => write!(f, "rate-limited"),
@@ -33,6 +35,7 @@ impl FromStr for SubscriberNotificationStatus {
             "queued" => Ok(SubscriberNotificationStatus::Queued),
             "processing" => Ok(SubscriberNotificationStatus::Processing),
             "published" => Ok(SubscriberNotificationStatus::Published),
+            "failed" => Ok(SubscriberNotificationStatus::Failed),
             "not-subscribed" => Ok(SubscriberNotificationStatus::NotSubscribed),
             "wrong-scope" => Ok(SubscriberNotificationStatus::WrongScope),
             "rate-limited" => Ok(SubscriberNotificationStatus::RateLimited),