Setup a job to expire compatibility sessions

Author: sandhose

crates/storage/src/queue/tasks.rs

@@ -4,7 +4,9 @@
 // Please see LICENSE in the repository root for full details.
 
 use chrono::{DateTime, Utc};
-use mas_data_model::{Device, Session, User, UserEmailAuthentication, UserRecoverySession};
+use mas_data_model::{
+    CompatSession, Device, Session, User, UserEmailAuthentication, UserRecoverySession,
+};
 use serde::{Deserialize, Serialize};
 use ulid::Ulid;
 
@@ -378,3 +380,60 @@ impl ExpireInactiveOAuthSessionsJob {
 impl InsertableJob for ExpireInactiveOAuthSessionsJob {
     const QUEUE_NAME: &'static str = "expire-inactive-oauth-sessions";
 }
+
+/// Expire inactive compatibility sessions
+#[derive(Serialize, Deserialize, Debug, Clone)]
+pub struct ExpireInactiveCompatSessionsJob {
+    threshold: DateTime<Utc>,
+    after: Option<Ulid>,
+}
+
+impl ExpireInactiveCompatSessionsJob {
+    /// Create a new job to expire inactive compatibility sessions
+    ///
+    /// # Parameters
+    ///
+    /// * `threshold` - The threshold to expire sessions at
+    #[must_use]
+    pub fn new(threshold: DateTime<Utc>) -> Self {
+        Self {
+            threshold,
+            after: None,
+        }
+    }
+
+    /// Get the threshold to expire sessions at
+    #[must_use]
+    pub fn threshold(&self) -> DateTime<Utc> {
+        self.threshold
+    }
+
+    /// Get the pagination cursor
+    #[must_use]
+    pub fn pagination(&self, batch_size: usize) -> Pagination {
+        let pagination = Pagination::first(batch_size);
+        if let Some(after) = self.after {
+            pagination.after(after)
+        } else {
+            pagination
+        }
+    }
+
+    /// Get the next job given the page returned by the database
+    #[must_use]
+    pub fn next(&self, page: &Page<CompatSession>) -> Option<Self> {
+        if !page.has_next_page {
+            return None;
+        }
+
+        let last_edge = page.edges.last()?;
+        Some(Self {
+            threshold: self.threshold,
+            after: Some(last_edge.id),
+        })
+    }
+}
+
+impl InsertableJob for ExpireInactiveCompatSessionsJob {
+    const QUEUE_NAME: &'static str = "expire-inactive-compat-sessions";
+}

crates/tasks/src/lib.rs

@@ -129,6 +129,7 @@ pub async fn init(
         .register_handler::<mas_storage::queue::SendEmailAuthenticationCodeJob>()
         .register_handler::<mas_storage::queue::SyncDevicesJob>()
         .register_handler::<mas_storage::queue::VerifyEmailJob>()
+        .register_handler::<mas_storage::queue::ExpireInactiveCompatSessionsJob>()
         .register_handler::<mas_storage::queue::ExpireInactiveOAuthSessionsJob>()
         .add_schedule(
             "cleanup-expired-tokens",

crates/tasks/src/sessions.rs

@@ -8,8 +8,12 @@ use std::collections::HashSet;
 use async_trait::async_trait;
 use chrono::Duration;
 use mas_storage::{
+    compat::CompatSessionFilter,
     oauth2::OAuth2SessionFilter,
-    queue::{ExpireInactiveOAuthSessionsJob, QueueJobRepositoryExt, SyncDevicesJob},
+    queue::{
+        ExpireInactiveCompatSessionsJob, ExpireInactiveOAuthSessionsJob, QueueJobRepositoryExt,
+        SyncDevicesJob,
+    },
 };
 
 use crate::{
@@ -80,3 +84,65 @@ impl RunnableJob for ExpireInactiveOAuthSessionsJob {
         Ok(())
     }
 }
+
+#[async_trait]
+impl RunnableJob for ExpireInactiveCompatSessionsJob {
+    async fn run(&self, state: &State, _context: JobContext) -> Result<(), JobError> {
+        let mut repo = state.repository().await.map_err(JobError::retry)?;
+        let clock = state.clock();
+        let mut rng = state.rng();
+        let mut users_synced = HashSet::new();
+
+        // This delay is used to space out the device sync jobs
+        // We add 10 seconds between each device sync, meaning that it will spread out
+        // the syncs over ~16 minutes max if we get a full batch of 100 users
+        let mut delay = Duration::minutes(1);
+
+        let filter = CompatSessionFilter::new()
+            .with_last_active_before(self.threshold())
+            .active_only();
+
+        let pagination = self.pagination(100);
+
+        let page = repo
+            .compat_session()
+            .list(filter, pagination)
+            .await
+            .map_err(JobError::retry)?
+            .map(|(c, _)| c);
+
+        if let Some(job) = self.next(&page) {
+            tracing::info!("Scheduling job to expire the next batch of inactive sessions");
+            repo.queue_job()
+                .schedule_job(&mut rng, &clock, job)
+                .await
+                .map_err(JobError::retry)?;
+        }
+
+        for edge in page.edges {
+            let inserted = users_synced.insert(edge.user_id);
+            if inserted {
+                tracing::info!(user.id = %edge.user_id, "Scheduling devices sync for user");
+                repo.queue_job()
+                    .schedule_job_later(
+                        &mut rng,
+                        &clock,
+                        SyncDevicesJob::new_for_id(edge.user_id),
+                        clock.now() + delay,
+                    )
+                    .await
+                    .map_err(JobError::retry)?;
+                delay += Duration::seconds(10);
+            }
+
+            repo.compat_session()
+                .finish(&clock, edge)
+                .await
+                .map_err(JobError::retry)?;
+        }
+
+        repo.save().await.map_err(JobError::retry)?;
+
+        Ok(())
+    }
+}