Skip to content

feat: Graceful session pool shutdown #623

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
aral01 wants to merge 6 commits into from
Open

feat: Graceful session pool shutdown #623

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
db92b17
refactor(session-pool): route DeleteSession through worker
aral01 Aug 7, 2026
9a80c69
feat(client): await session cleanup during shutdown
aral01 Aug 13, 2026
188857a
feat(query): await rollback-on-drop during shutdown
aral01 Aug 13, 2026
8d733ee
fix(session-pool): drain recursively submitted cleanup
aral01 Aug 16, 2026
ce30175
fix(query): rollback transient transaction failures
aral01 Aug 16, 2026
2aaf693
fix(session-pool): bound cleanup connection acquisition
aral01 Aug 24, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 11 additions & 0 deletions ydb/src/client.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -115,6 +115,17 @@ impl Client {
self.session_pool.stats()
}

/// Stop the shared session pool and wait for accepted session cleanup attempts to finish.
///
/// This consumes the driver and stops new session acquisition. Existing session leases may
/// finish; shutdown waits for them before deleting idle sessions. Do not use clients,
/// sessions, transactions, or streams derived from this driver after shutdown begins.
/// Shutdown must run while the Tokio runtime that created the driver is still alive.
#[instrument(name = "ydb.Driver.Shutdown", skip_all, fields(db.system.name = "ydb", db.namespace = %self.credentials.database), err)]
pub async fn shutdown(self) -> YdbResult<()> {
self.session_pool.shutdown().await

@rekby rekby Aug 17, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P2] This only shuts down the currently configured pool. with_session_pool replaces self.session_pool without shutting down or retaining the previous pool, while query/table clients created before replacement keep clones of that old pool. Consequently, after reconfiguration this method can return while old leases and cleanup tasks are still active, and those older clients can still acquire sessions. That defeats the documented shutdown guarantee and can still lose cleanup when the runtime is dropped. Please include superseded pools in the shutdown lifecycle or shut down the old pool during replacement.

— 🤖 AI review on behalf of @rekby

@rekby rekby Aug 24, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This still applies at the current head: with_session_pool replaces the field via Self { session_pool, ..self }, while shutdown only awaits self.session_pool.shutdown(). A previously created QueryClient or TableClient therefore retains an untracked clone of the superseded pool and can keep acquiring sessions after this shutdown returns.

— 🤖 AI review on behalf of @rekby

@rekby rekby Aug 25, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Still unresolved at f60ecbf: the new commit only bounds cleanup connection acquisition; with_session_pool still replaces the pool, and shutdown still awaits only the replacement. Previously derived clients can therefore retain and use the superseded pool.

— 🤖 AI review on behalf of @rekby

@rekby rekby Aug 26, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Still unresolved at 97be54b3: with_session_pool still replaces the field with Self { session_pool, ..self }, while shutdown awaits only that replacement. A query/table client created before replacement therefore retains the superseded pool and can continue acquiring sessions after this shutdown completes.

— 🤖 AI review on behalf of @rekby

@rekby rekby Aug 26, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Still unresolved at ac216fe: this push only changes cleanup timeout coverage. with_session_pool still replaces the pool via Self { session_pool, ..self }, while shutdown awaits only self.session_pool. Query/table clients created before replacement therefore retain the original pool and can still acquire sessions after shutdown returns.

— 🤖 AI review on behalf of @rekby

@rekby rekby Aug 26, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Still unresolved at 2aaf6938: with_session_pool continues to replace the field via Self { session_pool, ..self }, while shutdown only awaits the replacement. A query/table client created before replacement therefore retains the superseded pool and can keep acquiring sessions after shutdown returns.

— 🤖 AI review on behalf of @rekby

}

pub fn database(&self) -> String {
self.credentials.database.clone()
}
Expand Down
97 changes: 21 additions & 76 deletions ydb/src/client_query/exec.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@ use crate::types::Value;
use crate::{TransactionOptions, TxMode, closure};
use tracing::instrument;

use crate::session_pool::{SessionPool, SessionPoolLease, spawn_pool_release};
use crate::session_pool::{SessionPool, SessionPoolLease};

use super::hooks::{QueryTxCommitStatus, QueryTxHook};

Expand Down Expand Up @@ -269,30 +269,19 @@ impl TxExecContext {
}

let ActiveTx {
mut client,
lease,
server_progress,
..
} = active;
let tx_id = match (server_progress.operation, server_progress.tx_id) {
match (server_progress.operation, server_progress.tx_id) {
(TxOperationState::Ready, None) => {
lease.return_to_pool();
return Ok(());
}
(TxOperationState::Ready, Some(tx_id)) | (TxOperationState::InFlight, Some(tx_id)) => {
tx_id
lease.schedule_rollback(tx_id);
}
(TxOperationState::InFlight, None) => return Ok(()),
};
let cleanup_timeout = lease.cleanup_timeout();
let session_id = lease.session_id().to_string();

spawn_pool_release(async move {
let rollback = client
.rollback_transaction(&session_id, tx_id.as_str())
.map_err(YdbError::from);
finish_rollback_cleanup(lease, cleanup_timeout, rollback).await;
});
(TxOperationState::InFlight, None) => {}
}
Ok(())
}

Expand Down Expand Up @@ -330,29 +319,15 @@ impl Drop for TxExecContext {

active.notify_hooks(QueryTxCommitStatus::Aborted);
let ActiveTx {
mut client,
lease,
server_progress,
..
} = active;
let tx_id = match (server_progress.operation, server_progress.tx_id) {
(TxOperationState::Ready, None) => {
lease.return_to_pool();
return;
}
(TxOperationState::Ready, Some(tx_id)) => tx_id,
(TxOperationState::InFlight, _) => return,
};

let cleanup_timeout = lease.cleanup_timeout();
let session_id = lease.session_id().to_string();

spawn_pool_release(async move {
let rollback = client
.rollback_transaction(&session_id, tx_id.as_str())
.map_err(YdbError::from);
finish_rollback_cleanup(lease, cleanup_timeout, rollback).await;
});
match (server_progress.operation, server_progress.tx_id) {
(TxOperationState::Ready, None) => lease.return_to_pool(),
(TxOperationState::Ready, Some(tx_id)) => lease.schedule_rollback(tx_id),
(TxOperationState::InFlight, _) => {}
}
}
}

Expand Down Expand Up @@ -894,15 +869,6 @@ pub(crate) async fn tx_rollback(tx: &mut TxExecContext) -> YdbResult<()> {
}
}

async fn finish_rollback_cleanup<F>(lease: SessionPoolLease, cleanup_timeout: Duration, rollback: F)
where
F: Future<Output = YdbResult<()>>,
{
if matches!(timeout(cleanup_timeout, rollback).await, Ok(Ok(()))) {
lease.return_to_pool();
}
}

pub(crate) fn tx_exec_context(
client: RawQueryClient,
lease: SessionPoolLease,
Expand Down Expand Up @@ -1013,8 +979,8 @@ mod unit_tests {
}

#[tokio::test]
async fn rejected_query_does_not_reuse_session_with_unfinished_tx() {
for tx_id in [None, Some("tx-1".to_string())] {
async fn rejected_query_reuses_session_only_after_known_tx_is_rolled_back() {
for (tx_id, expect_reuse) in [(None, false), (Some("tx-1".to_string()), true)] {
let pool = SessionPool::new_explicit_bench(SessionPoolSettings::new().with_limit(1));
let lease = pool.acquire_explicit().await.expect("acquire test session");
let session_id = lease.session_id().to_string();
Expand All @@ -1035,17 +1001,17 @@ mod unit_tests {
.expect("rejected operation must finish the transaction");

assert!(matches!(ctx.state, TxState::AttemptFailed(_)));
let replacement = pool
let acquired = pool
.acquire_explicit()
.await
.expect("session with an unfinished transaction must be replaced");
assert_ne!(replacement.session_id(), session_id);
replacement.return_to_pool();
.expect("a session must become available after failed-attempt cleanup");
assert_eq!(acquired.session_id() == session_id, expect_reuse);
acquired.return_to_pool();
}
}

#[tokio::test]
async fn transient_dispatched_error_discards_session_when_cleanup_fails() {
async fn transient_dispatched_error_cleans_up_possibly_active_transaction() {
for status in [StatusCode::Unavailable, StatusCode::Overloaded] {
let pool = SessionPool::new_explicit_bench(SessionPoolSettings::new().with_limit(1));
let lease = pool.acquire_explicit().await.expect("acquire test session");
Expand All @@ -1062,12 +1028,12 @@ mod unit_tests {
.expect("temporary failure must finish the local transaction attempt");

assert!(matches!(ctx.state, TxState::AttemptFailed(_)));
let replacement = pool
let reused = pool
.acquire_explicit()
.await
.expect("session with an unconfirmed transaction must be replaced");
assert_ne!(replacement.session_id(), session_id);
replacement.return_to_pool();
.expect("session must become available after cleanup");
assert_eq!(reused.session_id(), session_id);
reused.return_to_pool();
}
}

Expand All @@ -1093,27 +1059,6 @@ mod unit_tests {
replacement.return_to_pool();
}

#[tokio::test]
async fn rollback_cleanup_timeout_discards_session_and_releases_pool_permit() {
let pool = SessionPool::new_explicit_bench(SessionPoolSettings::new().with_limit(1));
let lease = pool.acquire_explicit().await.expect("acquire test session");
let session_id = lease.session_id().to_string();

finish_rollback_cleanup(
lease,
Duration::ZERO,
std::future::pending::<YdbResult<()>>(),
)
.await;

let replacement = pool
.acquire_explicit()
.await
.expect("timed-out rollback must release the pool permit");
assert_ne!(replacement.session_id(), session_id);
replacement.return_to_pool();
}

#[tokio::test]
async fn unhealthy_session_before_query_ends_attempt() {
let pool = SessionPool::new_explicit_bench(SessionPoolSettings::new().with_limit(1));
Expand Down
4 changes: 2 additions & 2 deletions ydb/src/client_query/explain_query.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -162,8 +162,8 @@ mod unit_tests {
}
}

#[test]
fn timeout_is_recorded_only_when_set() {
#[tokio::test]
async fn timeout_is_recorded_only_when_set() {
let ctx = unreachable_ctx();
assert_eq!(
ExplainQueryBuilder::new(&ctx, "SELECT 1".to_string()).configured_timeout(),
Expand Down
Loading
Loading