tina-core 0.0.2

#![cfg(feature = "redis")]
//! Redis封装
pub mod cache;
pub(crate) mod lock;
pub(crate) mod throttle;

use crate::async_trait;
use crate::redis::cluster::{ClusterClient, ClusterConnection};
use crate::redis::{Client, Cmd, ConnectionLike, Pipeline, RedisFuture, Value};
use crate::tina::data::app_error::AppError;
use crate::tina::data::AppResult;
use crate::tina::server::application::AppConfig;
use crate::tina::server::session::Session;
use deadpool::managed::{Manager, Object, Pool, RecycleError, RecycleResult};
use futures::FutureExt;
use futures_util::future::BoxFuture;
use std::{
    fmt::Debug,
    ops::{Deref, DerefMut},
};

/// Redis客户端
#[async_trait]
pub trait IRedisClient {
    /// 获取Redis客户端
    fn get_redis_client(&self) -> AppResult<PooledRedisClient>;
    /// 设置Redis客户端
    fn set_redis_client(&mut self, client: PooledRedisClient);
    /// 获取Redis连接
    async fn get_redis_connection(&self) -> AppResult<PooledRedisConnection>;
}

/// Session的Redis客户端
pub trait ISessionRedisClient {
    /// 获取Redis客户端
    fn get_redis_client(&self) -> AppResult<PooledRedisClient>;
    /// 获取Redis连接
    fn get_redis_connection(&self) -> BoxFuture<'static, AppResult<PooledRedisConnection>>;
}

/// Redis客户端实例
pub enum RedisClient {
    /// 单点
    Single(Client),
    /// 集群
    Cluster(ClusterClient),
}

/// 池化的Redis客户端
#[derive(Clone, Debug)]
pub struct PooledRedisClient {
    pool: Pool<RedisClient>,
}

/// Redis连接
pub enum RedisConnection {
    /// 单点
    Single(crate::redis::Connection),
    /// 集群
    Cluster(crate::redis::cluster::ClusterConnection),
}

/// 池化的Redis连接
pub struct PooledRedisConnection {
    conn: Object<RedisClient>,
}

impl Debug for RedisClient {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Single(_) => f.debug_tuple("Single").finish(),
            Self::Cluster(_) => f.debug_tuple("Cluster").finish(),
        }
    }
}

impl Debug for RedisConnection {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Single(_) => f.debug_tuple("Single").finish(),
            Self::Cluster(_) => f.debug_tuple("Cluster").finish(),
        }
    }
}

impl Deref for PooledRedisClient {
    type Target = Pool<RedisClient>;
    fn deref(&self) -> &Pool<RedisClient> {
        &self.pool
    }
}

impl DerefMut for PooledRedisClient {
    fn deref_mut(&mut self) -> &mut Pool<RedisClient> {
        &mut self.pool
    }
}

impl AsRef<Pool<RedisClient>> for PooledRedisClient {
    fn as_ref(&self) -> &Pool<RedisClient> {
        &self.pool
    }
}

impl AsMut<Pool<RedisClient>> for PooledRedisClient {
    fn as_mut(&mut self) -> &mut Pool<RedisClient> {
        &mut self.pool
    }
}

impl Deref for PooledRedisConnection {
    type Target = RedisConnection;
    fn deref(&self) -> &RedisConnection {
        self.conn.deref()
    }
}

impl DerefMut for PooledRedisConnection {
    fn deref_mut(&mut self) -> &mut RedisConnection {
        self.conn.deref_mut()
    }
}

impl AsRef<RedisConnection> for PooledRedisConnection {
    fn as_ref(&self) -> &RedisConnection {
        self.conn.deref()
    }
}

impl AsMut<RedisConnection> for PooledRedisConnection {
    fn as_mut(&mut self) -> &mut RedisConnection {
        self.conn.deref_mut()
    }
}

impl From<Client> for RedisClient {
    fn from(client: Client) -> Self {
        RedisClient::Single(client)
    }
}

impl From<ClusterClient> for RedisClient {
    fn from(client: ClusterClient) -> Self {
        RedisClient::Cluster(client)
    }
}

impl From<crate::redis::Connection> for RedisConnection {
    fn from(conn: crate::redis::Connection) -> Self {
        RedisConnection::Single(conn)
    }
}

impl From<crate::redis::cluster::ClusterConnection> for RedisConnection {
    fn from(conn: ClusterConnection) -> Self {
        RedisConnection::Cluster(conn)
    }
}

#[async_trait]
impl Manager for RedisClient {
    type Type = RedisConnection;
    type Error = AppError;

    /// Creates a new instance of [`Manager::Type`].
    async fn create(&self) -> Result<Self::Type, Self::Error> {
        match self {
            RedisClient::Single(client) => match client.get_connection() {
                Ok(conn) => Ok(RedisConnection::from(conn)),
                Err(err) => Err(crate::app_system_error!("{}", err)),
            },
            RedisClient::Cluster(client) => match client.get_connection() {
                Ok(conn) => Ok(RedisConnection::from(conn)),
                Err(err) => Err(crate::app_system_error!("{}", err)),
            },
        }
    }

    /// Tries to recycle an instance of [`Manager::Type`].
    ///
    /// # Errors
    ///
    /// Returns [`Manager::Error`] if the instance couldn't be recycled.
    async fn recycle(&self, obj: &mut Self::Type) -> RecycleResult<Self::Error> {
        let is_open = match obj {
            RedisConnection::Single(v) => match v.is_open() {
                true => v.check_connection(),
                false => false,
            },
            RedisConnection::Cluster(v) => match v.is_open() {
                true => v.check_connection(),
                false => false,
            },
        };
        if !is_open {
            let new_conn = {
                match obj {
                    RedisConnection::Single(_) => match self.create().await {
                        Ok(conn) => Some(conn),
                        Err(err) => return Err(RecycleError::Backend(err)),
                    },
                    RedisConnection::Cluster(_) => match self.create().await {
                        Ok(conn) => Some(conn),
                        Err(err) => return Err(RecycleError::Backend(err)),
                    },
                }
            };
            if let Some(conn) = new_conn {
                (*obj) = conn;
            }
        }
        Ok(())
    }

    /// Detaches an instance of [`Manager::Type`] from this [`Manager`].
    ///
    /// This method is called when using the [`Object::take()`] method for
    /// removing an [`Object`] from a [`Pool`]. If the [`Manager`] doesn't hold
    /// any references to the handed out [`Object`]s then the default
    /// implementation can be used which does nothing.
    fn detach(&self, _obj: &mut Self::Type) {}
}

impl crate::redis::aio::ConnectionLike for PooledRedisConnection {
    fn req_packed_command<'a>(&'a mut self, cmd: &'a Cmd) -> RedisFuture<'a, Value> {
        self.deref_mut().req_packed_command(cmd)
    }

    fn req_packed_commands<'a>(&'a mut self, cmd: &'a Pipeline, offset: usize, count: usize) -> RedisFuture<'a, Vec<Value>> {
        self.deref_mut().req_packed_commands(cmd, offset, count)
    }

    fn get_db(&self) -> i64 {
        self.deref().get_db()
    }
}

impl crate::redis::aio::ConnectionLike for RedisConnection {
    fn req_packed_command<'a>(&'a mut self, cmd: &'a Cmd) -> RedisFuture<'a, Value> {
        match self {
            RedisConnection::Single(conn) => async move {
                let cmd = cmd.get_packed_command();
                conn.req_packed_command(cmd.as_ref())
            }
            .boxed(),
            RedisConnection::Cluster(conn) => async move {
                let cmd = cmd.get_packed_command();
                conn.req_packed_command(cmd.as_ref())
            }
            .boxed(),
        }
    }

    fn req_packed_commands<'a>(&'a mut self, cmd: &'a Pipeline, offset: usize, count: usize) -> RedisFuture<'a, Vec<Value>> {
        match self {
            RedisConnection::Single(conn) => async move {
                let cmd = cmd.get_packed_pipeline();
                conn.req_packed_commands(cmd.as_ref(), offset, count)
            }
            .boxed(),
            RedisConnection::Cluster(conn) => async move {
                let cmd = cmd.get_packed_pipeline();
                conn.req_packed_commands(cmd.as_ref(), offset, count)
            }
            .boxed(),
        }
    }

    fn get_db(&self) -> i64 {
        match self {
            RedisConnection::Single(conn) => conn.get_db(),
            RedisConnection::Cluster(conn) => conn.get_db(),
        }
    }
}

impl From<Pool<RedisClient>> for PooledRedisClient {
    fn from(pool: Pool<RedisClient>) -> Self {
        PooledRedisClient {
            pool,
        }
    }
}

impl From<Object<RedisClient>> for PooledRedisConnection {
    fn from(conn: Object<RedisClient>) -> Self {
        PooledRedisConnection {
            conn,
        }
    }
}

#[async_trait]
impl IRedisClient for AppConfig {
    fn get_redis_client(&self) -> AppResult<PooledRedisClient> {
        Ok(self.extension::<PooledRedisClient>()?.deref().clone())
    }

    fn set_redis_client(&mut self, client: PooledRedisClient) {
        self.add_extension(client);
    }

    async fn get_redis_connection(&self) -> AppResult<PooledRedisConnection> {
        let client = self.get_redis_client()?;
        let conn = client.get().await.map_err(crate::app_error_from!())?;
        Ok(PooledRedisConnection::from(conn))
    }
}

impl ISessionRedisClient for Session {
    fn get_redis_client(&self) -> AppResult<PooledRedisClient> {
        self.get_application().get_redis_client()
    }
    fn get_redis_connection(&self) -> BoxFuture<'static, AppResult<PooledRedisConnection>> {
        let app = self.get_application().clone();
        async move { app.get_redis_connection().await }.boxed()
    }
}