Struct google_cloud_spanner::apiv1::conn_pool::ConnectionManager

pub struct ConnectionManager { /* private fields */ }

Implementations

impl ConnectionManager

pub async fn new(
    pool_size: usize,
    environment: &Environment,
    domain: &str
) -> Result<Self, Error>

Examples found in repository

src/client.rs (line 199)

    pub async fn new_with_config(
        database: impl Into<String>,
        config: ClientConfig,
    ) -> Result<Self, InitializationError> {
        if config.session_config.max_opened > config.channel_config.num_channels * 100 {
            return Err(InitializationError::InvalidConfig(format!(
                "max session size is {} because max session size is 100 per gRPC connection",
                config.channel_config.num_channels * 100
            )));
        }

        let environment = Environment::from_project(config.project).await?;
        let pool_size = config.channel_config.num_channels;
        let conn_pool = ConnectionManager::new(pool_size, &environment, config.endpoint.as_str()).await?;
        let session_manager = SessionManager::new(database, conn_pool, config.session_config).await?;

        Ok(Client {
            sessions: Arc::new(session_manager),
        })
    }

pub fn num(&self) -> usize

Examples found in repository

src/session.rs (line 172)

    async fn init_pool(
        database: String,
        conn_pool: &ConnectionManager,
        min_opened: usize,
    ) -> Result<VecDeque<SessionHandle>, Status> {
        let channel_num = conn_pool.num();
        let creation_count_per_channel = min_opened / channel_num;

        let mut sessions = Vec::<SessionHandle>::new();
        for _ in 0..channel_num {
            let next_client = conn_pool.conn();
            match batch_create_session(next_client, database.clone(), creation_count_per_channel).await {
                Ok(r) => {
                    for i in r {
                        sessions.push(i);
                    }
                }
                Err(e) => return Err(e),
            }
        }
        tracing::debug!("initial session created count = {}", sessions.len());
        Ok(sessions.into())
    }

pub fn conn(&self) -> Client

Examples found in repository

src/session.rs (line 177)

    async fn init_pool(
        database: String,
        conn_pool: &ConnectionManager,
        min_opened: usize,
    ) -> Result<VecDeque<SessionHandle>, Status> {
        let channel_num = conn_pool.num();
        let creation_count_per_channel = min_opened / channel_num;

        let mut sessions = Vec::<SessionHandle>::new();
        for _ in 0..channel_num {
            let next_client = conn_pool.conn();
            match batch_create_session(next_client, database.clone(), creation_count_per_channel).await {
                Ok(r) => {
                    for i in r {
                        sessions.push(i);
                    }
                }
                Err(e) => return Err(e),
            }
        }
        tracing::debug!("initial session created count = {}", sessions.len());
        Ok(sessions.into())
    }

    fn request(&self) -> oneshot::Receiver<SessionHandle> {
        let (sender, receiver) = oneshot::channel();
        {
            self.waiters.lock().push_back(sender);
        }
        let _ = self.allocation_request_sender.send(true);
        receiver
    }

    fn num_opened(&self) -> usize {
        self.inner.lock().num_opened()
    }

    fn num_waiting(&self) -> usize {
        self.waiters.lock().len()
    }

    fn grow(&self, mut sessions: Vec<SessionHandle>) {
        while let Some(session) = sessions.pop() {
            match { self.waiters.lock().pop_front() } {
                Some(c) => {
                    let mut inner = self.inner.lock();
                    match c.send(session) {
                        Err(session) => inner.grow(session),
                        _ => {
                            // Mark as using when notify to waiter directory.
                            inner.inuse += 1
                        }
                    };
                }
                None => self.inner.lock().grow(session),
            };
        }
    }

    fn recycle(&self, session: SessionHandle) {
        if session.valid {
            tracing::trace!("recycled name={}", session.session.name);
            match { self.waiters.lock().pop_front() } {
                Some(c) => {
                    if let Err(session) = c.send(session) {
                        self.inner.lock().release(session)
                    }
                }
                None => self.inner.lock().release(session),
            };
        } else {
            self.inner.lock().release(session);

            // request session creation
            let _ = self.allocation_request_sender.send(true);
        }
    }
}

impl Clone for SessionPool {
    fn clone(&self) -> Self {
        SessionPool {
            inner: Arc::clone(&self.inner),
            waiters: Arc::clone(&self.waiters),
            allocation_request_sender: self.allocation_request_sender.clone(),
        }
    }
}

#[derive(Clone, Debug)]
pub struct SessionConfig {
    /// max_opened is the maximum number of opened sessions allowed by the session
    /// pool. If the client tries to open a session and there are already
    /// max_opened sessions, it will block until one becomes available or the
    /// context passed to the client method is canceled or times out.
    pub max_opened: usize,

    /// min_opened is the minimum number of opened sessions that the session pool
    /// tries to maintain. Session pool won't continue to expire sessions if
    /// number of opened connections drops below min_opened. However, if a session
    /// is found to be broken, it will still be evicted from the session pool,
    /// therefore it is posssible that the number of opened sessions drops below
    /// min_opened.
    pub min_opened: usize,

    /// max_idle is the maximum number of idle sessions, pool is allowed to keep.
    pub max_idle: usize,

    /// idle_timeout is the wait time before discarding an idle session.
    /// Sessions older than this value since they were last used will be discarded.
    /// However, if the number of sessions is less than or equal to min_opened, it will not be discarded.
    pub idle_timeout: std::time::Duration,

    pub session_alive_trust_duration: std::time::Duration,

    /// session_get_timeout is the maximum value of the waiting time that occurs when retrieving from the connection pool when there is no idle session.
    pub session_get_timeout: std::time::Duration,

    /// refresh_interval is the interval of cleanup and health check functions.
    pub refresh_interval: std::time::Duration,

    /// incStep is the number of sessions to create in one batch when at least
    /// one more session is needed.
    inc_step: usize,
}

impl Default for SessionConfig {
    fn default() -> Self {
        SessionConfig {
            max_opened: 400,
            min_opened: 10,
            max_idle: 300,
            inc_step: 25,
            idle_timeout: std::time::Duration::from_secs(30 * 60),
            session_alive_trust_duration: std::time::Duration::from_secs(55 * 60),
            session_get_timeout: std::time::Duration::from_secs(1),
            refresh_interval: std::time::Duration::from_secs(5 * 60),
        }
    }
}

pub struct SessionManager {
    session_pool: SessionPool,
    session_get_timeout: Duration,
    cancel: CancellationToken,
    tasks: Vec<JoinHandle<()>>,
}

#[derive(thiserror::Error, Debug)]
pub enum SessionError {
    #[error("session get time out")]
    SessionGetTimeout,
    #[error("failed to create session")]
    FailedToCreateSession,
    #[error(transparent)]
    GRPC(#[from] Status),
}

impl TryAs<Status> for SessionError {
    fn try_as(&self) -> Option<&Status> {
        match self {
            SessionError::GRPC(e) => Some(e),
            _ => None,
        }
    }
}

impl SessionManager {
    pub async fn new(
        database: impl Into<String>,
        conn_pool: ConnectionManager,
        config: SessionConfig,
    ) -> Result<SessionManager, Status> {
        let database = database.into();
        let (sender, receiver) = broadcast::channel(1);
        let session_pool = SessionPool::new(database.clone(), &conn_pool, config.min_opened, sender).await?;

        let cancel = CancellationToken::new();
        let session_get_timeout = config.session_get_timeout;
        let task_cleaner = schedule_refresh(config.clone(), session_pool.clone(), cancel.clone());
        let task_listener = listen_session_creation_request(
            config,
            session_pool.clone(),
            database,
            conn_pool,
            receiver,
            cancel.clone(),
        );

        let sm = SessionManager {
            session_get_timeout,
            session_pool,
            cancel,
            tasks: vec![task_cleaner, task_listener],
        };
        Ok(sm)
    }

    pub fn num_opened(&self) -> usize {
        self.session_pool.num_opened()
    }

    pub fn session_waiters(&self) -> usize {
        self.session_pool.num_waiting()
    }

    pub async fn get(&self) -> Result<ManagedSession, SessionError> {
        if let Some(mut s) = self.session_pool.inner.lock().take() {
            s.last_used_at = Instant::now();
            return Ok(ManagedSession::new(self.session_pool.clone(), s));
        }

        // Wait for the session creation.
        match timeout(self.session_get_timeout, self.session_pool.request()).await {
            Ok(Ok(mut session)) => {
                session.last_used_at = Instant::now();
                Ok(ManagedSession {
                    session_pool: self.session_pool.clone(),
                    session: Some(session),
                })
            }
            _ => Err(SessionError::SessionGetTimeout),
        }
    }

    pub(crate) async fn close(&self) {
        if self.cancel.is_cancelled() {
            return;
        }
        self.cancel.cancel();
        sleep(Duration::from_secs(1)).await;
        for task in &self.tasks {
            task.abort();
        }
        let deleting_sessions = {
            let mut lock = self.session_pool.inner.lock();
            let mut deleting_sessions = Vec::with_capacity(lock.sessions.len());
            while let Some(session) = lock.sessions.pop_front() {
                deleting_sessions.push(session);
            }
            deleting_sessions
        };
        for mut session in deleting_sessions {
            delete_session(&mut session).await;
        }
    }
}

fn listen_session_creation_request(
    config: SessionConfig,
    session_pool: SessionPool,
    database: String,
    conn_pool: ConnectionManager,
    mut rx: broadcast::Receiver<bool>,
    cancel: CancellationToken,
) -> JoinHandle<()> {
    tokio::spawn(async move {
        let mut allocation_request_size = 0;
        loop {
            select! {
                _ = rx.recv() => {},
                _ = cancel.cancelled() => break
            }
            let num_opened = session_pool.num_opened();
            if num_opened >= config.min_opened && allocation_request_size >= session_pool.num_waiting() {
                continue;
            }

            let mut creation_count = config.max_opened - num_opened;
            if creation_count > config.inc_step {
                creation_count = config.inc_step;
            }
            if creation_count == 0 {
                continue;
            }
            allocation_request_size += creation_count;

            let database = database.clone();
            let next_client = conn_pool.conn();

            match batch_create_session(next_client, database, creation_count).await {
                Ok(fresh_sessions) => {
                    allocation_request_size -= creation_count;
                    session_pool.grow(fresh_sessions)
                }
                Err(e) => {
                    allocation_request_size -= creation_count;
                    tracing::error!("failed to create new sessions {:?}", e)
                }
            };
        }
        tracing::trace!("stop session creating listener")
    })
}