fix(db): 会话复用与生命周期改为按 session 同步,修掉会话池三处并发缺陷

按 review 意见修复上一版(a044b63d3)留下的三处问题:

1. try_acquire_session 在持有会话表写锁时逐条 await 候选会话的 config()
   (原 manager.rs:633),一个正在跑长语句的会话会连带堵住其它页签与该连接的
   元数据路径 —— 本次要修的串行问题在复用/建会话路径上并没有修干净。
   现在会话表锁内只做「候选筛选 + 占用」,连接锁与 ping 一律在锁外;候选身份
   改用缓存的 DatabaseIdentity(由 verify_and_sync_database 同步刷新),校验
   失败或会话已不在池内就换下一个候选,不再持锁等待。

2. release_session_internal 用 bool 覆盖 in_use,释放与下一条语句的占用交错时
   会把正在执行的会话标成空闲,随后被 idle 清理回收(表现为 session not found)。
   改为引用计数占用 + 预留位:create_session/try_acquire_session 只做一次预留,
   由紧接着取连接的语句消费,保证一条语句恰好占用一次、释放一次,且两个并发
   create_session 不会共用同一会话。

3. release_session_internal 即使 detach 未成功也会 close_session,与并发的
   close_session/remove_all_sessions 重复 disconnect。现在只有真正把会话从池中
   摘出的调用方才负责断开;try_acquire_session 丢弃失效会话也走同一条路径。

附带:get_session_connection 在取得连接后复查会话是否仍在池内,避免把已摘除并
关闭的连接交出去;list_sessions 改读缓存身份,不再等待忙碌会话的连接锁;
get_session_config 保持读取实时配置,但不再在持有会话表锁时等待。

验证(macOS 本地,crates/db edition 2021):
- cargo test -p db --lib:1308 passed / 0 failed
- cargo test -p db:全量通过(tests/real_postgres.rs 4、tests/real_sqlite.rs 4,
  doc-test 1 ignored)
- cargo test -p db_view --lib:717 passed / 0 failed
- cargo check -p main:通过
- cargo clippy -p db --all-targets:manager.rs 仅剩 3 条既有告警(行号后移)
- rustfmt --edition 2021 crates/db/src/manager.rs

新增 5 个回归测试,均先复现失败再修复通过:过期释放覆盖新占用、复用扫描堵住
其它连接、并发 release/close 只 disconnect 一次、一次语句只占用一次占用单位、
复用扫描遇到忙碌会话立即返回。未跑 story 测试(需真实数据库与交互环境),
未在 Windows/Linux 上验证。
This commit is contained in:
胡飞
2026-09-21 14:51:08 +08:00
parent a044b63d33
commit 24e87f8627
+494 -78
View File
@@ -27,7 +27,7 @@ use gpui::{AppContext, AsyncApp, Global, Task};
use one_core::connection_notifier::{ConnectionDataEvent, GlobalConnectionNotifier};
use one_core::gpui_tokio::Tokio;
use one_core::storage::{ConnectionRepository, DatabaseType, DbConnectionConfig};
use std::collections::HashMap;
use std::collections::{HashMap, HashSet};
use std::sync::{Arc, Mutex as StdMutex};
type ExternalRegistryReloader = dyn Fn() -> IpcDriverRegistry + Send + Sync;
@@ -283,27 +283,76 @@ impl Clone for DbManager {
///
/// The physical connection lives behind its own mutex, so a running statement only
/// locks the connection it uses. Bookkeeping lives in [`SessionState`] instead of the
/// session itself so status reads never wait for an in-flight statement.
/// session itself so status reads never wait for an in-flight statement. Reuse checks
/// read the cached [`DatabaseIdentity`] for the same reason: picking a candidate must
/// not wait for a session that is busy with a statement.
struct ConnectionSession {
connection: Arc<AsyncMutex<Box<dyn DbConnection + Send + Sync>>>,
close_on_release: bool,
created_at: Instant,
session_id: String,
state: StdMutex<SessionState>,
identity: StdMutex<DatabaseIdentity>,
}
/// Mutable per-session bookkeeping kept separate from the connection lock.
struct SessionState {
in_use: bool,
/// Statements currently holding this session.
///
/// A bool would lose a claim: a release that finishes after the next statement
/// already took the session would mark a busy session idle, and the cleanup task
/// could then recycle the connection of a running statement.
in_flight: usize,
/// Set when a session is handed to a caller that has not taken the connection yet
/// (`create_session`). The matching `get_session_connection` consumes it, so one
/// statement is never counted twice while the hand-off stays protected against
/// concurrent reuse.
reserved: bool,
last_active: Instant,
}
/// Identity fields that decide whether a session can serve a request.
///
/// Cached on the session so the pool can pick a reuse candidate without waiting for
/// the connection lock. `verify_and_sync_database` is the only place that rewrites
/// `DbConnectionConfig::database`, and it refreshes this snapshot, so the cache cannot
/// silently drift from the live connection.
struct DatabaseIdentity {
database_type: DatabaseType,
database: Option<String>,
sid: Option<String>,
service_name: Option<String>,
}
impl DatabaseIdentity {
fn from_config(config: &DbConnectionConfig) -> Self {
Self {
database_type: config.database_type.clone(),
database: config.database.clone(),
sid: config.sid.clone(),
service_name: config.service_name.clone(),
}
}
/// Whether a session with this identity can serve `config`.
fn matches(&self, config: &DbConnectionConfig) -> bool {
match self.database_type {
DatabaseType::Oracle => {
(self.sid.is_some() && self.sid == config.sid)
|| (self.service_name.is_some() && self.service_name == config.service_name)
}
_ => self.database.is_some() && self.database == config.database,
}
}
}
impl ConnectionSession {
fn new(
connection: Box<dyn DbConnection + Send + Sync>,
session_id: String,
close_on_release: bool,
) -> Self {
let identity = DatabaseIdentity::from_config(connection.config());
let now = Instant::now();
Self {
connection: Arc::new(AsyncMutex::new(connection)),
@@ -311,9 +360,11 @@ impl ConnectionSession {
created_at: now,
session_id,
state: StdMutex::new(SessionState {
in_use: false,
in_flight: 0,
reserved: false,
last_active: now,
}),
identity: StdMutex::new(identity),
}
}
@@ -323,29 +374,66 @@ impl ConnectionSession {
.unwrap_or_else(|poisoned| poisoned.into_inner())
}
fn mark_in_use(&self) {
/// Cached identity used to pick reuse candidates without the connection lock.
fn identity(&self) -> std::sync::MutexGuard<'_, DatabaseIdentity> {
self.identity
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner())
}
/// Refresh the cached identity from the live connection config.
fn refresh_identity(&self, connection: &(dyn DbConnection + Send + Sync)) {
*self.identity() = DatabaseIdentity::from_config(connection.config());
}
/// Whether the cached identity says this session can serve `config`.
fn can_serve(&self, config: &DbConnectionConfig) -> bool {
self.identity().matches(config)
}
/// Hand the session to a statement that is about to take its connection.
fn reserve(&self) {
let mut state = self.state();
state.in_use = true;
state.reserved = true;
state.last_active = Instant::now();
}
/// Claim the session for the statement that is taking its connection now.
///
/// Consumes a reservation made by [`Self::reserve`], so create_session and the
/// following statement do not count the same use twice.
fn mark_in_use(&self) {
let mut state = self.state();
state.reserved = false;
state.in_flight += 1;
state.last_active = Instant::now();
}
/// Release one claim (or undo one reservation).
fn release(&self) {
let mut state = self.state();
state.in_use = false;
state.reserved = false;
state.in_flight = state.in_flight.saturating_sub(1);
state.last_active = Instant::now();
}
fn in_use(&self) -> bool {
self.state().in_use
let state = self.state();
state.in_flight > 0 || state.reserved
}
fn last_active(&self) -> Instant {
self.state().last_active
}
/// Idle long enough to be recycled.
///
/// Ignores `reserved` on purpose: a reservation that is never consumed (its caller
/// died in between) must not pin the session forever. A live reservation is young,
/// so it is still protected by the timeout.
fn is_expired(&self, timeout: Duration) -> bool {
let state = self.state();
!state.in_use && state.last_active.elapsed() > timeout
state.in_flight == 0 && state.last_active.elapsed() > timeout
}
fn is_lifetime_expired(&self, max_lifetime: Duration) -> bool {
@@ -357,13 +445,19 @@ impl ConnectionSession {
Arc::clone(&self.connection).lock_owned().await
}
async fn config(&self) -> DbConnectionConfig {
/// Validate that this session can serve `config` before it is reused.
///
/// Takes the connection lock, so the caller must not hold the session table lock.
/// Returns `Ok(false)` when the live database no longer matches `config`, in which
/// case the refreshed identity is kept so the next loop can pick another candidate.
async fn validate_reuse(&self, config: &DbConnectionConfig) -> Result<bool, DbError> {
let connection = self.lock_connection().await;
connection.config().clone()
}
async fn ping(&self) -> Result<(), DbError> {
self.lock_connection().await.ping().await
self.refresh_identity(&**connection);
if !self.can_serve(config) {
return Ok(false);
}
connection.ping().await?;
Ok(true)
}
/// Check if current database matches config database
@@ -381,8 +475,9 @@ impl ConnectionSession {
if config_db == current_db {
Ok(true)
} else {
// Database changed, update config
// Database changed: update the config together with the cached identity
connection.set_config_database(current_db.clone());
self.refresh_identity(&**connection);
info!(
"Session {} database changed: {:?} -> {:?}",
self.session_id, config_db, current_db
@@ -510,10 +605,13 @@ impl ConnectionManager {
session_id, config.database
);
// Store session
// Store session, reserved for the caller that is about to take its connection.
// The matching `get_session_connection` consumes the reservation, so every
// statement adds exactly one claim and its release brings the counter back to
// zero, while two concurrent create_session calls never share a session.
let session =
ConnectionSession::new(connection, session_id.clone(), lifecycle.close_on_release);
session.mark_in_use();
session.reserve();
let mut sessions = self.sessions.write().await;
sessions
@@ -567,9 +665,22 @@ impl ConnectionManager {
session
};
Ok(SessionConnectionGuard {
connection: session.lock_connection().await,
})
// The session may have been detached (and its connection closed) while we queued
// behind a running statement; never hand out a connection the pool lost.
//
// This nests the connection lock and the session table lock. It cannot deadlock:
// no path holds the session table lock while awaiting a connection lock
// (`try_acquire_session` releases it before validating a candidate).
let connection = session.lock_connection().await;
if !self.is_pooled(session_id).await {
session.release();
return Err(DbError::Internal(format!(
"session not found: {}",
session_id
)));
}
Ok(SessionConnectionGuard { connection })
}
/// Clone the shared handle of one session without holding the pool lock.
@@ -603,83 +714,118 @@ impl ConnectionManager {
detached
}
fn db_equals(db1: &DbConnectionConfig, db2: &DbConnectionConfig) -> bool {
match db1.database_type {
DatabaseType::Oracle => {
(db1.sid.is_some() && db1.sid == db2.sid)
|| (db1.service_name.is_some() && db1.service_name == db2.service_name)
}
_ => db1.database.is_some() && db1.database == db2.database,
}
}
/// Try to acquire an existing idle session with matching database
///
/// The session table lock is held only to pick and claim a candidate: every wait —
/// the live database check and the ping — happens after it is released. A session
/// busy with a long statement therefore never stalls the pool for other connections.
async fn try_acquire_session(
&self,
config: &DbConnectionConfig,
) -> Result<Option<String>, DbError> {
loop {
let mut sessions = self.sessions.write().await;
let mut has_busy_close_on_release_session = false;
// Candidates already rejected by the out-of-lock checks below.
let mut rejected: HashSet<String> = HashSet::new();
let remove_config_entry = {
loop {
let (candidate, waiting_for_busy_session) = {
let mut sessions = self.sessions.write().await;
let Some(session_list) = sessions.get_mut(&config.id) else {
return Ok(None);
};
let mut index = 0;
while index < session_list.len() {
let session = &session_list[index];
let matches_database = Self::db_equals(&session.config().await, config);
if matches_database && session.in_use() && session.close_on_release {
has_busy_close_on_release_session = true;
index += 1;
let mut candidate: Option<Arc<ConnectionSession>> = None;
let mut waiting_for_busy_session = false;
for session in session_list.iter() {
if rejected.contains(&session.session_id) || !session.can_serve(config) {
continue;
}
let matches_config = !session.in_use() && matches_database;
if !matches_config {
index += 1;
if session.in_use() {
waiting_for_busy_session |= session.close_on_release;
continue;
}
if let Err(error) = session.ping().await {
let session = session_list.remove(index);
warn!(
"Discarding stale session {} before reuse: {}",
session.session_id, error
);
session.close().await;
continue;
}
let session = &session_list[index];
session.mark_in_use();
debug!(
"Reusing session: {} (database: {:?})",
session.session_id, config.database
);
return Ok(Some(session.session_id.clone()));
candidate = Some(Arc::clone(session));
break;
}
session_list.is_empty()
if let Some(session) = &candidate {
session.reserve();
}
let empty = session_list.is_empty();
if empty {
sessions.remove(&config.id);
}
(candidate, waiting_for_busy_session)
};
if remove_config_entry {
sessions.remove(&config.id);
let Some(session) = candidate else {
if waiting_for_busy_session {
sleep(BUSY_CLOSE_ON_RELEASE_RETRY_DELAY).await;
continue;
}
return Ok(None);
};
let session_id = session.session_id.clone();
match session.validate_reuse(config).await {
Ok(true) => {}
Ok(false) => {
// The live database drifted from the request; the refreshed identity
// stays cached, so look for another candidate.
session.release();
rejected.insert(session_id);
continue;
}
Err(error) => {
warn!(
"Discarding stale session {} before reuse: {}",
session_id, error
);
session.release();
self.detach_and_close(&session_id).await;
continue;
}
}
drop(sessions);
if has_busy_close_on_release_session {
sleep(BUSY_CLOSE_ON_RELEASE_RETRY_DELAY).await;
if !self.is_pooled(&session_id).await {
// Removed while we were waiting for the connection; never hand it out.
session.release();
continue;
}
return Ok(None);
debug!(
"Reusing session: {} (database: {:?})",
session_id, config.database
);
return Ok(Some(session_id));
}
}
/// Whether a session is still part of the pool.
///
/// Borrows the session table only, so a caller holding a connection lock can use it
/// to re-check ownership before handing the connection out.
async fn is_pooled(&self, session_id: &str) -> bool {
let sessions = self.sessions.read().await;
sessions
.values()
.any(|list| list.iter().any(|session| session.session_id == session_id))
}
/// Detach a session from the pool and close it.
///
/// Returns `false` when another task detached it first: that task owns the close, so
/// the physical connection is disconnected exactly once.
async fn detach_and_close(&self, session_id: &str) -> bool {
let Some(session) = self.detach_session(session_id).await else {
return false;
};
session.close().await;
true
}
}
/// Guard that locks one session's connection for the duration of a statement
@@ -697,7 +843,8 @@ impl SessionConnectionGuard {
impl ConnectionManager {
/// Get session config
pub async fn get_session_config(&self, session_id: &str) -> Option<DbConnectionConfig> {
self.find_session(session_id).await?.config().await.into()
let session = self.find_session(session_id).await?;
Some(session.lock_connection().await.config().clone())
}
pub async fn release_session(&self, session_id: &str) -> Result<(), DbError> {
@@ -732,14 +879,21 @@ impl ConnectionManager {
}
};
if should_close {
self.detach_session(session_id).await;
let detached = should_close && self.detach_session(session_id).await.is_some();
if detached {
// This task removed the session from the pool, so it owns the disconnect.
session.release();
session.close().await;
return Ok(());
}
session.release();
debug!("Session {} released", session_id);
if should_close {
debug!("Session {} was already detached", session_id);
} else {
debug!("Session {} released", session_id);
}
Ok(())
}
@@ -852,7 +1006,9 @@ impl ConnectionManager {
let last_active = session.last_active();
infos.push(SessionInfo {
session_id: session.session_id.clone(),
database: session.config().await.database.clone(),
// Cached identity instead of the live config: listing sessions must not
// wait for a session that is busy with a statement.
database: session.identity().database.clone(),
in_use: session.in_use(),
idle_time: last_active.elapsed(),
lifetime: session.created_at.elapsed(),
@@ -5288,6 +5444,266 @@ mod tests {
assert!(execution_dropped.load(Ordering::SeqCst));
}
/// 一个语句的释放不能清掉下一条语句对同一个 session 的占用。
/// 占用必须按计数维护:释放与下一次获取会交错,布尔量会把正在执行的
/// session 记为空闲,进而被清理任务回收。
#[tokio::test]
async fn overlapping_release_keeps_the_next_statement_claim() {
let session = ConnectionSession::new(
Box::new(MockConnection::new(test_config("claim"), true)),
"claim:session:1".to_string(),
false,
);
session.mark_in_use();
session.mark_in_use();
session.release();
assert!(
session.in_use(),
"finishing the first statement must not clear the claim of the next one"
);
session.release();
assert!(!session.in_use());
}
/// 一次语句只占用一个单位:取连接时占用,释放后退回空闲并重新可回收。
#[tokio::test]
async fn a_statement_claims_and_releases_a_session_exactly_once() {
let manager = ConnectionManager::new();
let config = test_config("balanced-claim");
let session_id = "balanced-claim:session:1";
let session = Arc::new(ConnectionSession::new(
Box::new(MockConnection::new(config.clone(), true)),
session_id.to_string(),
false,
));
manager
.sessions
.write()
.await
.entry(config.id.clone())
.or_default()
.push(Arc::clone(&session));
let guard = manager.get_session_connection(session_id).await.unwrap();
assert!(session.in_use());
drop(guard);
manager.release_session(session_id).await.unwrap();
assert!(
!session.in_use(),
"a released session must go back to idle so it can be reused"
);
assert!(
session.is_expired(Duration::ZERO),
"a released session must be eligible for idle cleanup"
);
// 复用成功后再走一遍语句的占用/释放,占用计数必须重新回到 0。
let acquired = manager.try_acquire_session(&config).await.unwrap();
assert_eq!(Some(session_id.to_string()), acquired);
let guard = manager.get_session_connection(session_id).await.unwrap();
assert!(session.in_use());
drop(guard);
manager.release_session(session_id).await.unwrap();
assert!(
!session.in_use(),
"reusing a session must not leave a claim behind"
);
}
/// 释放任务先排队等连接锁、下一条语句随后拿到连接时,session 不能被记为 idle。
#[tokio::test]
async fn releasing_a_session_does_not_report_the_next_statement_as_idle() {
let manager = ConnectionManager::new();
let config = test_config("claim-race");
let session_id = "claim-race:session:1";
let session = Arc::new(ConnectionSession::new(
Box::new(MockConnection::new(config.clone(), true)),
session_id.to_string(),
false,
));
manager
.sessions
.write()
.await
.entry(config.id.clone())
.or_default()
.push(Arc::clone(&session));
let mut first = manager.get_session_connection(session_id).await.unwrap();
assert!(first.connection().is_some());
// 释放任务先进连接锁队列;连接锁是 FIFO,sleep 只用来固定队列顺序。
let release = {
let manager = manager.clone();
let session_id = session_id.to_string();
tokio::spawn(async move { manager.release_session(&session_id).await })
};
sleep(Duration::from_millis(20)).await;
// 下一条语句随后取同一 session:先标记占用,然后排在释放之后拿到连接。
let next = {
let manager = manager.clone();
let session_id = session_id.to_string();
tokio::spawn(async move { manager.get_session_connection(&session_id).await })
};
sleep(Duration::from_millis(20)).await;
drop(first);
let mut next = tokio::time::timeout(Duration::from_secs(2), next)
.await
.expect("next statement should not hang")
.expect("join")
.expect("next statement should get the connection");
assert!(next.connection().is_some());
tokio::time::timeout(Duration::from_secs(2), release)
.await
.expect("release should not hang")
.expect("join")
.expect("release should succeed");
// 直接读 session 状态:读会话表需要写锁,旧实现里会被 guard 挡住而挂住。
assert!(
session.in_use(),
"a session holding a connection for a statement must not be reported as idle"
);
drop(next);
}
/// 复用会话时不得在持有会话表锁的情况下等待某条会话的连接锁,
/// 否则一个正在跑长查询的会话会重新堵住所有页签(含其它连接)。
#[tokio::test]
async fn try_acquire_session_does_not_block_the_pool_while_a_session_is_busy() {
let manager = ConnectionManager::new();
let busy = test_config("busy-config");
let other = test_config("other-config");
let busy_id = "busy-config:session:1";
let other_id = "other-config:session:1";
insert_test_session(
&manager,
&busy,
busy_id,
MockConnection::new(busy.clone(), true),
)
.await;
insert_test_session(
&manager,
&other,
other_id,
MockConnection::new(other.clone(), true),
)
.await;
// busy 会话正在执行:它的连接锁被长期持有。
let mut busy_guard = manager.get_session_connection(busy_id).await.unwrap();
assert!(busy_guard.connection().is_some());
let acquire = {
let manager = manager.clone();
let config = busy.clone();
tokio::spawn(async move { manager.try_acquire_session(&config).await })
};
sleep(Duration::from_millis(20)).await;
// 复用扫描进行中,其它连接仍必须能取到自己的会话。
tokio::time::timeout(
Duration::from_millis(300),
manager.get_session_connection(other_id),
)
.await
.expect("a busy session must not block the pool for other connections")
.expect("other session should be found");
let acquired = tokio::time::timeout(Duration::from_secs(2), acquire)
.await
.expect("acquire should not block on a busy session's connection lock")
.expect("join")
.expect("acquire");
assert_eq!(
None, acquired,
"a session that is executing a statement must not be handed out again"
);
drop(busy_guard);
manager.release_session(busy_id).await.unwrap();
let acquired =
tokio::time::timeout(Duration::from_secs(2), manager.try_acquire_session(&busy))
.await
.expect("acquire should not hang")
.expect("acquire");
assert_eq!(Some(busy_id.to_string()), acquired);
}
/// 并发释放与关闭同一个 session 时,物理连接只能断开一次。
#[tokio::test]
async fn concurrent_release_and_close_disconnect_the_session_once() {
let manager = ConnectionManager::new();
let config = test_config("double-close");
let session_id = "double-close:session:1";
let disconnect_count = Arc::new(AtomicUsize::new(0));
manager
.sessions
.write()
.await
.entry(config.id.clone())
.or_default()
.push(Arc::new(ConnectionSession::new(
Box::new(MockConnection::with_disconnect_count(
config.clone(),
true,
Arc::clone(&disconnect_count),
)),
session_id.to_string(),
true,
)));
let mut guard = manager.get_session_connection(session_id).await.unwrap();
assert!(guard.connection().is_some());
// 释放任务先排队等连接锁,随后 close_session 把同一个 session 摘走。
let release = {
let manager = manager.clone();
let session_id = session_id.to_string();
tokio::spawn(async move { manager.release_session(&session_id).await })
};
sleep(Duration::from_millis(20)).await;
let close = {
let manager = manager.clone();
let session_id = session_id.to_string();
tokio::spawn(async move { manager.close_session(&session_id).await })
};
sleep(Duration::from_millis(20)).await;
drop(guard);
tokio::time::timeout(Duration::from_secs(2), release)
.await
.expect("release should not hang")
.expect("join")
.expect("release should succeed");
tokio::time::timeout(Duration::from_secs(2), close)
.await
.expect("close should not hang")
.expect("join")
.expect("close should succeed");
assert_eq!(
1,
disconnect_count.load(Ordering::SeqCst),
"the physical connection must be disconnected exactly once"
);
assert!(manager.list_sessions(&config.id).await.is_empty());
}
#[tokio::test]
async fn closing_an_executing_session_waits_for_cancellation_and_disconnects_once() {
let state = GlobalDbState::new();