gitdataai/libs/session/storage/redis_cluster.rs

use std::sync::Arc;

use actix_web::cookie::time::Duration;
use redis::Commands;
use redis::cluster::ClusterClient;
use tokio::task;

use super::SessionKey;
use crate::storage::{
    SessionStore,
    format::{deserialize_session_state, serialize_session_state},
    interface::{LoadError, SaveError, SessionState, UpdateError},
    utils::generate_session_key,
};

#[derive(Clone)]
pub struct RedisClusterSessionStore {
    configuration: CacheConfiguration,
    client: ClusterClient,
}

#[derive(Clone)]
struct CacheConfiguration {
    cache_keygen: Arc<dyn Fn(&str) -> String + Send + Sync>,
}

impl Default for CacheConfiguration {
    fn default() -> Self {
        Self {
            cache_keygen: Arc::new(str::to_owned),
        }
    }
}

impl RedisClusterSessionStore {
    pub fn builder(connection_strings: Vec<String>) -> RedisClusterSessionStoreBuilder {
        RedisClusterSessionStoreBuilder {
            configuration: CacheConfiguration::default(),
            connection_strings,
        }
    }

    pub async fn new(connection_strings: Vec<String>) -> anyhow::Result<RedisClusterSessionStore> {
        Self::builder(connection_strings).build().await
    }

    fn get_connection(&self) -> anyhow::Result<redis::cluster::ClusterConnection> {
        self.client.get_connection().map_err(|e| anyhow::anyhow!(e))
    }
}

#[must_use]
pub struct RedisClusterSessionStoreBuilder {
    configuration: CacheConfiguration,
    connection_strings: Vec<String>,
}

impl RedisClusterSessionStoreBuilder {
    pub fn cache_keygen<F>(mut self, keygen: F) -> Self
    where
        F: Fn(&str) -> String + 'static + Send + Sync,
    {
        self.configuration.cache_keygen = Arc::new(keygen);
        self
    }

    pub async fn build(self) -> anyhow::Result<RedisClusterSessionStore> {
        let client = ClusterClient::new(self.connection_strings)?;
        Ok(RedisClusterSessionStore {
            configuration: self.configuration,
            client,
        })
    }
}

impl SessionStore for RedisClusterSessionStore {
    async fn load(&self, session_key: &SessionKey) -> Result<Option<SessionState>, LoadError> {
        let cache_key = self.configuration.cache_keygen.as_ref()(session_key.as_ref());
        let conn = self.get_connection().map_err(LoadError::Other)?;

        let value: Option<String> = task::spawn_blocking(move || {
            let mut conn = conn;
            conn.get::<_, Option<String>>(&cache_key)
        })
        .await
        .map_err(|_| LoadError::Other(anyhow::anyhow!("Task panicked")))?
        .map_err(|e: redis::RedisError| LoadError::Other(anyhow::anyhow!(e)))?;

        match value {
            None => Ok(None),
            Some(value) => Ok(Some(
                deserialize_session_state(&value).map_err(LoadError::Deserialization)?,
            )),
        }
    }

    async fn save(
        &self,
        session_state: SessionState,
        ttl: &Duration,
    ) -> Result<SessionKey, SaveError> {
        let body = serialize_session_state(&session_state).map_err(SaveError::Serialization)?;
        let session_key = generate_session_key();
        let cache_key = self.configuration.cache_keygen.as_ref()(session_key.as_ref());
        let ttl_secs = ttl.whole_seconds() as u64;
        let conn = self.get_connection().map_err(SaveError::Other)?;

        task::spawn_blocking(move || {
            let mut conn = conn;
            conn.set_ex::<_, _, ()>(&cache_key, &body, ttl_secs)
        })
        .await
        .map_err(|_| SaveError::Other(anyhow::anyhow!("Task panicked")))?
        .map_err(|e: redis::RedisError| SaveError::Other(anyhow::anyhow!(e)))?;

        Ok(session_key)
    }

    async fn update(
        &self,
        session_key: SessionKey,
        session_state: SessionState,
        ttl: &Duration,
    ) -> Result<SessionKey, UpdateError> {
        let body = serialize_session_state(&session_state).map_err(UpdateError::Serialization)?;
        let cache_key = self.configuration.cache_keygen.as_ref()(session_key.as_ref());
        let ttl_secs = ttl.whole_seconds() as u64;
        let conn = self.get_connection().map_err(UpdateError::Other)?;

        let existed: bool = task::spawn_blocking(move || {
            let mut conn = conn;
            conn.set_ex::<_, _, bool>(&cache_key, &body, ttl_secs)
        })
        .await
        .map_err(|_| UpdateError::Other(anyhow::anyhow!("Task panicked")))?
        .map_err(|e: redis::RedisError| UpdateError::Other(anyhow::anyhow!(e)))?;

        if !existed {
            self.save(session_state, ttl)
                .await
                .map_err(|err| match err {
                    SaveError::Serialization(err) => UpdateError::Serialization(err),
                    SaveError::Other(err) => UpdateError::Other(err),
                })
        } else {
            Ok(session_key)
        }
    }

    async fn update_ttl(&self, session_key: &SessionKey, ttl: &Duration) -> anyhow::Result<()> {
        let cache_key = self.configuration.cache_keygen.as_ref()(session_key.as_ref());
        let ttl_secs = ttl.whole_seconds() as i64;
        let conn = self.get_connection()?;

        task::spawn_blocking(move || {
            let mut conn = conn;
            conn.expire::<_, bool>(&cache_key, ttl_secs)
        })
        .await
        .map_err(|_| anyhow::anyhow!("Task panicked"))?
        .map_err(|e: redis::RedisError| anyhow::anyhow!(e))?;

        Ok(())
    }

    async fn delete(&self, session_key: &SessionKey) -> Result<(), anyhow::Error> {
        let cache_key = self.configuration.cache_keygen.as_ref()(session_key.as_ref());
        let conn = self.get_connection()?;

        task::spawn_blocking(move || {
            let mut conn = conn;
            conn.del::<_, i64>(&cache_key)
        })
        .await
        .map_err(|_| anyhow::anyhow!("Task panicked"))?
        .map_err(|e: redis::RedisError| anyhow::anyhow!(e))?;

        Ok(())
    }
}