Implement activity tracking for personal sessions

Author: reivilibre

crates/handlers/src/activity_tracker/mod.rs

@@ -10,7 +10,9 @@ mod worker;
 use std::net::IpAddr;
 
 use chrono::{DateTime, Utc};
-use mas_data_model::{BrowserSession, Clock, CompatSession, Session};
+use mas_data_model::{
+    BrowserSession, Clock, CompatSession, Session, personal::session::PersonalSession,
+};
 use mas_storage::BoxRepositoryFactory;
 use tokio_util::{sync::CancellationToken, task::TaskTracker};
 use ulid::Ulid;
@@ -115,7 +117,7 @@ impl ActivityTracker {
     pub async fn record_personal_access_token_session(
         &self,
         clock: &dyn Clock,
-        session: &Session,
+        session: &PersonalSession,
         ip: Option<IpAddr>,
     ) {
         let res = self

crates/handlers/src/activity_tracker/worker.rs

@@ -257,7 +257,9 @@ impl Worker {
         repo.compat_session()
             .record_batch_activity(compat_sessions)
             .await?;
-        // TODO: personal sessions: record
+        repo.personal_session()
+            .record_batch_activity(personal_sessions)
+            .await?;
 
         repo.save().await?;
         self.pending_records.clear();

crates/storage-pg/.sqlx/query-64b6e274e2bed6814f5ae41ddf57093589f7d1b2b8458521b635546b8012041e.json

@@ -361,6 +361,56 @@ impl PersonalSessionRepository for PgPersonalSessionRepository<'_> {
             .try_into()
             .map_err(DatabaseError::to_invalid_operation)
     }
+
+    #[tracing::instrument(
+        name = "db.personal_session.record_batch_activity",
+        skip_all,
+        fields(
+            db.query.text,
+        ),
+        err,
+    )]
+    async fn record_batch_activity(
+        &mut self,
+        mut activities: Vec<(Ulid, DateTime<Utc>, Option<IpAddr>)>,
+    ) -> Result<(), Self::Error> {
+        // Sort the activity by ID, so that when batching the updates, Postgres
+        // locks the rows in a stable order, preventing deadlocks
+        activities.sort_unstable();
+        let mut ids = Vec::with_capacity(activities.len());
+        let mut last_activities = Vec::with_capacity(activities.len());
+        let mut ips = Vec::with_capacity(activities.len());
+
+        for (id, last_activity, ip) in activities {
+            ids.push(Uuid::from(id));
+            last_activities.push(last_activity);
+            ips.push(ip);
+        }
+
+        let res = sqlx::query!(
+            r#"
+                UPDATE personal_sessions
+                SET last_active_at = GREATEST(t.last_active_at, personal_sessions.last_active_at)
+                  , last_active_ip = COALESCE(t.last_active_ip, personal_sessions.last_active_ip)
+                FROM (
+                    SELECT *
+                    FROM UNNEST($1::uuid[], $2::timestamptz[], $3::inet[])
+                        AS t(personal_session_id, last_active_at, last_active_ip)
+                ) AS t
+                WHERE personal_sessions.personal_session_id = t.personal_session_id
+            "#,
+            &ids,
+            &last_activities,
+            &ips as &[Option<IpAddr>],
+        )
+        .traced()
+        .execute(&mut *self.conn)
+        .await?;
+
+        DatabaseError::ensure_affected_rows(&res, ids.len().try_into().unwrap_or(u64::MAX))?;
+
+        Ok(())
+    }
 }
 
 impl Filter for PersonalSessionFilter<'_> {

crates/storage-pg/src/personal/session.rs

@@ -3,6 +3,8 @@
 // SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-Element-Commercial
 // Please see LICENSE files in the repository root for full details.
 
+use std::net::IpAddr;
+
 use async_trait::async_trait;
 use chrono::{DateTime, Utc};
 use mas_data_model::{
@@ -109,6 +111,21 @@ pub trait PersonalSessionRepository: Send + Sync {
     ///
     /// Returns [`Self::Error`] if the underlying repository fails
     async fn count(&mut self, filter: PersonalSessionFilter<'_>) -> Result<usize, Self::Error>;
+
+    /// Record a batch of [`PersonalSession`] activity
+    ///
+    /// # Parameters
+    ///
+    /// * `activity`: A list of tuples containing the session ID, the last
+    ///   activity timestamp and the IP address of the client
+    ///
+    /// # Errors
+    ///
+    /// Returns [`Self::Error`] if the underlying repository fails
+    async fn record_batch_activity(
+        &mut self,
+        activity: Vec<(Ulid, DateTime<Utc>, Option<IpAddr>)>,
+    ) -> Result<(), Self::Error>;
 }
 
 repository_impl!(PersonalSessionRepository:
@@ -137,6 +154,11 @@ repository_impl!(PersonalSessionRepository:
     ) -> Result<Page<PersonalSession>, Self::Error>;
 
     async fn count(&mut self, filter: PersonalSessionFilter<'_>) -> Result<usize, Self::Error>;
+
+    async fn record_batch_activity(
+        &mut self,
+        activity: Vec<(Ulid, DateTime<Utc>, Option<IpAddr>)>,
+    ) -> Result<(), Self::Error>;
 );
 
 /// Filter parameters for listing personal sessions