Modrinth-plus/apps/labrinth/src/queue/session.rs

use crate::database::models::pat_item::PersonalAccessToken;
use crate::database::models::session_item::Session;
use crate::database::models::{
    DatabaseError, OAuthAccessTokenId, PatId, SessionId, UserId,
};
use crate::database::redis::RedisPool;
use crate::routes::internal::session::SessionMetadata;
use chrono::Utc;
use itertools::Itertools;
use sqlx::PgPool;
use std::collections::{HashMap, HashSet};
use tokio::sync::Mutex;

pub struct AuthQueue {
    session_queue: Mutex<HashMap<SessionId, SessionMetadata>>,
    pat_queue: Mutex<HashSet<PatId>>,
    oauth_access_token_queue: Mutex<HashSet<OAuthAccessTokenId>>,
}

impl Default for AuthQueue {
    fn default() -> Self {
        Self::new()
    }
}

// Batches session accessing transactions every 30 seconds
impl AuthQueue {
    pub fn new() -> Self {
        AuthQueue {
            session_queue: Mutex::new(HashMap::with_capacity(1000)),
            pat_queue: Mutex::new(HashSet::with_capacity(1000)),
            oauth_access_token_queue: Mutex::new(HashSet::with_capacity(1000)),
        }
    }
    pub async fn add_session(&self, id: SessionId, metadata: SessionMetadata) {
        self.session_queue.lock().await.insert(id, metadata);
    }

    pub async fn add_pat(&self, id: PatId) {
        self.pat_queue.lock().await.insert(id);
    }

    pub async fn add_oauth_access_token(
        &self,
        id: crate::database::models::OAuthAccessTokenId,
    ) {
        self.oauth_access_token_queue.lock().await.insert(id);
    }

    pub async fn take_sessions(&self) -> HashMap<SessionId, SessionMetadata> {
        let mut queue = self.session_queue.lock().await;
        let len = queue.len();

        std::mem::replace(&mut queue, HashMap::with_capacity(len))
    }

    pub async fn take_hashset<T>(queue: &Mutex<HashSet<T>>) -> HashSet<T> {
        let mut queue = queue.lock().await;
        let len = queue.len();

        std::mem::replace(&mut queue, HashSet::with_capacity(len))
    }

    pub async fn index(
        &self,
        pool: &PgPool,
        redis: &RedisPool,
    ) -> Result<(), DatabaseError> {
        let session_queue = self.take_sessions().await;
        let pat_queue = Self::take_hashset(&self.pat_queue).await;
        let oauth_access_token_queue =
            Self::take_hashset(&self.oauth_access_token_queue).await;

        if !session_queue.is_empty()
            || !pat_queue.is_empty()
            || !oauth_access_token_queue.is_empty()
        {
            let mut transaction = pool.begin().await?;
            let mut clear_cache_sessions = Vec::new();

            for (id, metadata) in session_queue {
                clear_cache_sessions.push((Some(id), None, None));

                sqlx::query!(
                    "
                    UPDATE sessions
                    SET last_login = $2, city = $3, country = $4, ip = $5, os = $6, platform = $7, user_agent = $8
                    WHERE (id = $1)
                    ",
                    id as SessionId,
                    Utc::now(),
                    metadata.city,
                    metadata.country,
                    metadata.ip,
                    metadata.os,
                    metadata.platform,
                    metadata.user_agent,
                )
                .execute(&mut *transaction)
                .await?;
            }

            use futures::TryStreamExt;
            let expired_ids = sqlx::query!(
                "
                SELECT id, session, user_id
                FROM sessions
                WHERE refresh_expires <= NOW()
                "
            )
            .fetch(&mut *transaction)
            .map_ok(|x| (SessionId(x.id), x.session, UserId(x.user_id)))
            .try_collect::<Vec<(SessionId, String, UserId)>>()
            .await?;

            for (id, session, user_id) in expired_ids {
                clear_cache_sessions.push((
                    Some(id),
                    Some(session),
                    Some(user_id),
                ));
                Session::remove(id, &mut transaction).await?;
            }

            Session::clear_cache(clear_cache_sessions, redis).await?;

            let ids = pat_queue.iter().map(|id| id.0).collect_vec();
            let clear_cache_pats = pat_queue
                .into_iter()
                .map(|id| (Some(id), None, None))
                .collect_vec();
            sqlx::query!(
                "
                UPDATE pats
                SET last_used = $2
                WHERE id IN
                (SELECT * FROM UNNEST($1::bigint[]))
                ",
                &ids[..],
                Utc::now(),
            )
            .execute(&mut *transaction)
            .await?;

            update_oauth_access_token_last_used(
                oauth_access_token_queue,
                &mut transaction,
            )
            .await?;

            transaction.commit().await?;
            PersonalAccessToken::clear_cache(clear_cache_pats, redis).await?;
        }

        Ok(())
    }
}

async fn update_oauth_access_token_last_used(
    oauth_access_token_queue: HashSet<OAuthAccessTokenId>,
    transaction: &mut sqlx::Transaction<'_, sqlx::Postgres>,
) -> Result<(), DatabaseError> {
    let ids = oauth_access_token_queue.iter().map(|id| id.0).collect_vec();
    sqlx::query!(
        "
        UPDATE oauth_access_tokens
        SET last_used = $2
        WHERE id IN
        (SELECT * FROM UNNEST($1::bigint[]))
        ",
        &ids[..],
        Utc::now()
    )
    .execute(&mut **transaction)
    .await?;
    Ok(())
}