ydb 0.15.0

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

use tracing::trace;

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

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

/// Table-side adapter over the driver session pool.
#[derive(Clone)]
pub(crate) struct TableSessionPool {
    pool: SessionPool,
    connection_manager: GrpcConnectionManager,
    timeouts: TimeoutSettings,
}

impl TableSessionPool {
    pub(crate) fn from_shared(
        pool: SessionPool,
        connection_manager: GrpcConnectionManager,
        timeouts: TimeoutSettings,
    ) -> Self {
        Self {
            pool,
            connection_manager,
            timeouts,
        }
    }

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

    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),
            self.timeouts,
        );

        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::client::TimeoutSettings;
    use crate::errors::YdbResult;
    use crate::grpc_connection_manager::GrpcConnectionManager;
    use crate::grpc_wrapper::grpc_limits::DEFAULT_GRPC_MESSAGE_SIZE_LIMIT_BYTES;
    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(),
            None,
            DEFAULT_GRPC_MESSAGE_SIZE_LIMIT_BYTES,
        )
    }

    #[tokio::test]
    async fn max_active_session() -> YdbResult<()> {
        let pool = TableSessionPool::from_shared(
            bench_pool(),
            bench_connection_manager(),
            TimeoutSettings::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(())
    }
}