ydb 0.17.1

Crate contains generated low-level grpc code from YDB API protobuf, used as base for ydb crate
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131

use tracing::trace;

use crate::client::TimeoutSettings;
use crate::errors::YdbResult;
use crate::grpc_connection_manager::GrpcConnectionManager;
use crate::retry_budget::RetryControl;
use crate::session::{NodePinnedTableClient, TableSession};

use super::pool::{SessionPool, spawn_pool_release};

/// Table-side adapter over the driver session pool.
#[derive(Clone)]
pub(crate) struct TableSessionPool {
    pool: SessionPool,
    connection_manager: GrpcConnectionManager,
    retry_control: std::sync::Arc<RetryControl>,
}

impl TableSessionPool {
    pub(crate) fn from_shared(
        pool: SessionPool,
        connection_manager: GrpcConnectionManager,
        retry_control: std::sync::Arc<RetryControl>,
    ) -> Self {
        Self {
            pool,
            connection_manager,
            retry_control,
        }
    }

    pub(crate) fn connection_manager(&self) -> &GrpcConnectionManager {
        &self.connection_manager
    }

    pub(crate) fn retry_control(&self) -> &RetryControl {
        &self.retry_control
    }

    pub(crate) async fn session(&self) -> YdbResult<TableSession> {
        let mut lease = self.pool.acquire_explicit().await?;
        lease.ensure_alive()?;
        lease.begin_use();
        let session_id = lease.session_id().to_string();
        let node_uri = lease.node_uri().clone();

        let mut session = TableSession::new(
            session_id,
            NodePinnedTableClient::new(self.connection_manager.clone(), node_uri),
            TimeoutSettings::default(),
        );

        session.on_drop(Box::new(move |s: &mut TableSession| {
            if !s.can_pooled {
                lease.invalidate_session();
            }
            // Drop the in-use guard synchronously so AttachSession close can drain in-flight
            // table RPCs before async return_to_pool (return_to_pool also calls end_use).
            lease.end_use();
            spawn_pool_release(async move {
                lease.return_to_pool().await;
            });
        }));

        trace!("leased table session: {}", session.id);
        Ok(session)
    }
}

#[cfg(test)]
mod test {
    use super::TableSessionPool;
    use crate::GrpcOptions;
    use crate::errors::YdbResult;
    use crate::grpc_connection_manager::GrpcConnectionManager;
    use crate::grpc_wrapper::runtime_interceptors::MultiInterceptor;
    use crate::load_balancer::{SharedLoadBalancer, StaticLoadBalancer};
    use crate::session_pool::{SessionPool, SessionPoolSettings};
    use http::Uri;
    use std::time::Duration;
    use tokio::sync::oneshot;

    fn bench_pool() -> SessionPool {
        SessionPool::new_explicit_bench(SessionPoolSettings::new().with_limit(1).with_warm_up(1))
    }

    fn bench_connection_manager() -> GrpcConnectionManager {
        GrpcConnectionManager::new(
            SharedLoadBalancer::new_with_balancer(Box::new(StaticLoadBalancer::new(
                Uri::from_static("http://127.0.0.1/bench"),
            ))),
            "bench".to_string(),
            MultiInterceptor::new(),
            GrpcOptions::default(),
        )
    }

    use crate::retry_budget::RetryControl;

    #[tokio::test]
    async fn max_active_session() -> YdbResult<()> {
        let pool = TableSessionPool::from_shared(
            bench_pool(),
            bench_connection_manager(),
            std::sync::Arc::new(RetryControl::default()),
        );
        let first_session = pool.session().await?;

        let (thread_started_sender, thread_started_receiver) = oneshot::channel();
        let (second_session_got_sender, mut second_session_got_receiver) = oneshot::channel();
        let cloned_pool = pool.clone();

        tokio::spawn(async move {
            thread_started_sender.send(true).unwrap();
            cloned_pool.session().await.unwrap();
            second_session_got_sender.send(true).unwrap();
        });

        thread_started_receiver.await?;

        tokio::time::sleep(Duration::from_millis(100)).await;

        assert!(second_session_got_receiver.try_recv().is_err());

        drop(first_session);

        second_session_got_receiver.await?;

        Ok(())
    }
}