TiDB 内部系统会话管理增强基于 WithSession 与所有权机制的会话池重构解析【免费下载链接】tidbTiDB is built for agentic workloads that grow unpredictably, with ACID guarantees and native support for transactions, analytics, and vector search. No data silos. No noisy neighbors. No infrastructure ceiling.项目地址: https://gitcode.com/GitHub_Trending/ti/tidb本文深入解析 TiDB 针对内部系统会话Internal System Session管理的增强设计。TTL、统计分析、异步加载等后台模块在运行时会频繁创建、复用内部会话本文将从原设计文档出发结合仓库中pkg/session/syssession的实际源码与测试完整阐述新会话池的WithSession编程模型、Session 包装层、所有权owner转移机制、状态重置策略以及未来的并发检测方向帮助开发者理解并正确使用这套会话生命周期管理方案。背景与动机会话池为什么会泄漏内存在 TiDB 中TTL 任务、统计信息收集stats handle、Workload Learning 等大量后台模块都需要执行 SQL 来读写系统表与业务数据。为了降低创建、销毁会话的开销TiDB 很早就引入了会话池Session Pool来复用内部会话。旧实现位于 pkg/util/session_pool.go其核心接口如下type SessionPool interface { Get() (pools.Resource, error) Put(pools.Resource) Close() } type DestroyableSessionPool interface { SessionPool Destroy(pools.Resource) }旧池采用Get/Put的借还模式并额外引入了Destroy方法用于销毁脏会话见 pkg/util/session_pool.go#L32-L38 的注释If the caller meets an error when using the session, it can destroy the session。但从实际运行情况看这种依赖调用方自觉归还的设计带来了几个典型问题会话未归还导致内存泄漏TTL 模块曾出现从池中取出会话后未归还的场景对应 issue #56934会话被长期持有且其内部引用例如注册到 session manager 的 map无法释放。Close 与 map 引用冲突PR #32726 之后每个从池中取出的会话都会被作为引用存储在一个 map 中该引用只有在会话归还时才被移除。如果调用方直接调用Close而不是归还map 中仍然残留引用造成泄漏。API 歧义虽然 PR #59546 通过引入Destroy方法试图缓解上述问题但开发者仍然难以判断何时该调用Put、何时该调用Destroy、何时该调用Close错误调用的风险并未消除。设计文档给出的三个核心目标由此展开docs/design/2025-03-17-sys-session-manage-enhancement.md阻止跨模块并发使用会话一旦归还池中就不应再被任何调用方使用简化生命周期管理保证会话用后必定归还或关闭消除内存泄漏风险强制状态重置会话归还时强制重置内部状态避免脏状态影响后续复用。新会话池设计以 WithSession 为核心新设计引入了增强型会话池位于 pkg/session/syssession/pool.go对外暴露的接口如下// Pool is an interface for system internal session pool. type Pool interface { // Get gets a session from the session pool. Get() (*Session, error) // Put puts the session back to the pool. Put(*Session) // WithForceBlockGCSession executes the input function with the session and ensures the session is registered to // the session manager so GC can be blocked safely. WithForceBlockGCSession(ctx context.Context, fn func(*Session) error) error // WithSession executes the input function with the session. // After the function called, the session will be returned to the pool automatically. WithSession(func(*Session) error) error }其中最重要的方法是WithSession它接收一个func(*Session) error回调池负责从池中取出会话、执行回调、并在回调结束后自动把会话归还。调用方无需再关心借和还的配对问题。WithSession 的实现defer 保证必归还AdvancedSessionPool.WithSession的实际实现pkg/session/syssession/pool.go#L257-L279与设计文档中的示例一致func (p *AdvancedSessionPool) WithSession(fn func(*Session) error) error { se, err : p.Get() if err ! nil { return err } success : false defer func() { if success { p.Put(se) } else { se.Close() } }() if err fn(se); err ! nil { return err } success true return nil }这段代码的精妙之处在于回调正常返回时success true会话通过p.Put(se)归还池中供后续复用回调返回错误或发生 panic时success仍为false会话通过se.Close()被销毁而不是归还。原因在于错误或 panic 场景下会话内部可能处于未定义状态例如事务未提交、游标未关闭直接归还反而会把脏状态传染给下一次使用者销毁是更安全的选择。在 pkg/session/syssession/pool_test.go#L318-L381 的TestSessionPoolWithSession中对成功、返回错误、panic、Get失败四种场景都有明确的断言成功路径池中会话数变为 1已归还错误与 panic 路径池中会话数为 0已销毁。池的容量与创建NewAdvancedSessionPoolAdvancedSessionPool的内部结构pkg/session/syssession/pool.go#L51-L81如下type AdvancedSessionPool struct { noopOwnerHook ctx context.Context pool chan *session factory Factory mu struct { sync.RWMutex closed bool cancel context.CancelFunc } } func NewAdvancedSessionPool(capacity int, factory Factory) *AdvancedSessionPool几个值得注意的细节池底层是一个带缓冲的 channelpool chan *session容量即池大小PoolMaxSize为1024 * 1024 * 1024pkg/session/syssession/pool.go#L31-L32传入非法容量 0或超上限时会被重置为PoolMaxSizefactory是一个Factory类型签名func() (SessionContext, error)用于在池为空时创建新的内部会话上下文池创建时会注入内部事务来源类型kv.InternalTxnOthers确保池内会话执行的事务被识别为内部事务getInternal采用先尝试从 channel 取、取不到再新建的策略若新建失败会在 defer 中关闭已创建的sctx避免半初始化资源泄漏。在 Domain 层面TiDB 在初始化时通过syssession.NewAdvancedSessionPool(systemSessionPoolSize, factory)创建了全局的增强池并提供了AdvancedSysSessionPool()访问入口见 pkg/domain/domain.go#L598-L606 与 pkg/domain/domain.go#L1301-L1304后者明确标注Deprecated: Use AdvancedSysSessionPool instead.建议新代码统一使用增强池。池的关闭Close 的并发安全处理Closepkg/session/syssession/pool.go#L331-L347会关闭底层 channel、取消 context并逐个关闭池内残留的会话。值得注意的是Put与Close并发时的安全性处理Put在向p.poolchannel 发送前会先持有p.mu.RLock检查closed标志防止出现检查通过后 channel 已被关闭导致 panic的竞态见代码注释 pkg/session/syssession/pool.go#L245-L248。Session 包装层公开包装器与私有内部会话设计文档明确指出公开的Session类型并非真实会话而是内部会话的包装器// Session is public to callers. type Session struct { internal *session }这一设计对应源码中的两个类型pkg/session/syssession/session.gosession私有真正的会话载体封装了sctx SessionContext并附带所有权、并发计数等元信息只有Pool与Session两个类型可以访问它Session公开调用方拿到的包装器实现sqlexec.SQLExecutor与sqlexec.RestrictedSQLExecutor接口代理访问内部会话。type session struct { mu sync.Mutex // sctx is the raw state of the session. sctx SessionContext // owner is the one that holds the current session. // In most cases, only the owner can do operations on the session. // If owner nil, it means the session is closed and should not be used. owner sessionOwner // seq is a monotone increasing uint64 to provide the unique sequence number for each operation seq uint64 // inUse is the counter to record how many operations are on going on the session. inUse uint64 // unsafe is used to detect race conditions in thread-unsafe operations. unsafe uint64 // avoidReuse will be marked as true when some panics detected. avoidReuse bool }包装层的三大优势设计文档给出了选择包装层方案的三个理由均能在源码中得到印证方法级安全检查包装方法在访问内部会话前先校验调用方是否是当前 owner。例如ExecuteInternal的实现func (s *Session) ExecuteInternal(ctx context.Context, sql string, args ...any) (sqlexec.RecordSet, error) { sctx, exit, err : s.internal.EnterOperation(s, false) if err ! nil { return nil, err } defer exit() return sctx.GetSQLExecutor().ExecuteInternal(ctx, sql, args...) }会话被归还或关闭后EnterOperation会因caller is not the owner或session is closed返回错误从而拦截对失效会话的误用对应 pkg/session/syssession/session.go#L534-L541 与checkOwnerWithoutLock的判定逻辑。接口瘦身内部SessionContext包含大量方法包装层只暴露Execute、ExecuteInternal、ExecuteStmt、ExecRestrictedSQL、ParseWithParams等调用方真正需要的方法简化使用面。可扩展性Session后续可以方便地增加新方法例如WithSessionContext、IsOwner、AvoidReuse等便利能力。内部会话的并发保护与状态字段session内部通过mu sync.Mutex保护并发访问并维护了几个关键计数器注释见 pkg/session/syssession/session.go#L82-L149owner当前持有者。owner nil表示已关闭只有 owner 才能对会话执行操作所有权通过TransferOwner转移inUse正在进行的操作数。EnterOperation进入时加一、退出时减一inUse 0时会话空闲才允许转移所有权或关闭inUse 0时TransferOwner会失败从而阻止不同 owner 的并发访问unsafe用于检测非线程安全方法的并发竞争。unsafe 0/1为正常unsafe 1表示检测到并发竞争只有第一个操作被允许执行其余操作返回错误avoidReuse当 panic 发生时被标记为 true此后该会话不再复用seq单调递增的uint64序号用于日志与调试帮助追踪操作顺序。这些字段共同实现了设计文档中的owner 机制池本身是池内会话的 ownerGet后包装器Session成为新 owner归还时所有权转回池公开的Session一旦失去所有权就永远无法再次获得因此一次使用后即失效从机制上杜绝了归还后继续使用的问题。所有权转移与状态重置Put 的完整流程AdvancedSessionPool.Putpkg/session/syssession/pool.go#L142-L255是整条归还链路上逻辑最复杂的部分它依次执行空值/失效判断se nil或se.internal nil直接返回若se.internal.Owner() ! se例如重复Put、会话已关闭、内部会话已被其他Session取走直接忽略保证幂等所有权转回池internal.TransferOwner(se, p)这一步先校验输入Session的有效性若转移失败典型场景是并发Put与Get竞争则调用se.Close()关闭会话并记录错误日志。源码注释特意说明了为何必须用se.Close()而非se.internal.Close()——后者可能误关掉已被其他 goroutine 取走的会话避免复用检查internal.IsAvoidReuse()为 true此前发生过 panic时直接关闭不再放回池未结束事务检查internal.CheckNoPendingTxn()检查是否存在待 TSO 的PreparedTxnFuture或仍然有效的Txn若存在则以错误日志记录并关闭会话避免把未确定状态的事务带回池中状态重置internal.OwnerResetState(p.ctx, p)对会话执行sctx.RollbackTxn(ctx)强制回滚、重置内部状态保证下一次复用时是干净的会话放回池在p.mu.RLock保护下通过select向p.poolchannel 发送若池已满则走default分支由 defer 中的internal.Close()兜底关闭。对应单元测试TestSessionPoolPutpkg/session/syssession/pool_test.go#L123-L316覆盖了全部关键分支重复Put、归还非 owner 的旧Session、归还正在使用的会话、归还标记 avoid-reuse 的会话、归还含 pending txn 的会话、归还已关闭会话、池满、池关闭等十余种场景。会话注册与反注册session manager 的联动设计文档中的Pool.get/Pool.put示意代码展示了会话与 session manager 的注册联动。在实际实现中这一联动被放置在所有权挂钩owner hook中完成Session.onBecameOwner调用infosync.StoreInternalSession(sctx)将会话注册进 session managerSession.onResignOwner调用infosync.DeleteInternalSession(sctx)将会话从 session manager 注销。见 pkg/session/syssession/session.go#L595-L602。这样一来无论会话是归还池中还是被关闭只要发生所有权交接注册表都会随之更新调用方无需感知。infosync层提供了三个入口pkg/domain/infosync/info.go#L1149-L1190StoreInternalSession、DeleteInternalSession、ContainsInternalSession。最终由 server 层的SessionManager落地到内存 mapfunc (s *Server) StoreInternalSession(se any) { s.sessionMapMutex.Lock() s.internalSessions[se] struct{}{} metrics.InternalSessions.Set(float64(len(s.internalSessions))) s.sessionMapMutex.Unlock() }见 pkg/server/server.go#L1224-L1231DeleteInternalSession与之对称。server 还通过GetInternalSessionStartTSList收集内部会话的 startTS 列表pkg/server/server.go#L1257-L1271供 GC 安全点计算时阻塞相关 TS防止内部会话持有的事务被 GC 清理。值得说明的是这一注册机制在服务端无 session manager 时会被安全跳过StoreInternalSession在 manager 为 nil 时返回 false不产生副作用并配套了WithForceBlockGCSession这类需要确保会话已注册、GC 可被安全阻塞的场景专用入口pkg/session/syssession/pool.go#L283-L329。集成测试TestDomainAdvancedSessionPoolInternalSessionRegistrypkg/session/syssession/session_integration_test.go#L31-L75验证了WithSession回调执行期间会话存在于 session manager 中回调结束会话归还后即被注销使用Get取出后手动Close的场景下会话同样会从 session manager 注销。这正对应了设计文档要解决的第二个问题——会话被 Close 后仍残留在 map 中导致内存泄漏。实际应用TTL 模块中的会话池使用方式TTL 是内部会话的主要使用者之一。TTL 的 JobManager 持有syssession.Poolpkg/ttl/ttlworker/job_manager.go#L111并通过WithSession驱动整个任务循环return withSession(m.sessPool, m.jobLoopWithSession)见 pkg/ttl/ttlworker/job_manager.go#L212。删除任务的工作线程ttlDeleteWorker同样以syssession.Pool作为会话来源pkg/ttl/ttlworker/del.go#L326-L330。这种一次回调持有一个会话、用后自动归还的模式天然适配 TTL 表数量大、任务并发高的场景——这正是设计文档测试计划中1000 TTL 表、小tidb_ttl_job_internal设置下频繁创建与销毁会话所要验证的目标。在 pkg/ttl/ttlworker/job_manager_integration_test.go#L1266 等集成测试中测试代码直接通过dom.AdvancedSysSessionPool()获取 Domain 级增强池来驱动 TTL 任务验证了从全局池取会话的完整链路。测试策略设计文档规划的测试分为单元测试与集成测试两类仓库中均已落地单元测试pkg/session/syssession/pool_test.go 使用 mock 工厂与 mock 会话上下文覆盖池的创建、Get含新建、复用、工厂失败、池已关闭、Put含十余种异常分支、WithSession成功/错误/panic/Get 失败、Close含重复关闭等全部路径集成测试pkg/session/syssession/session_integration_test.go 在真实 mock store 与 Domain 之上验证 session manager 注册联动以及归还脏会话包括归还已关闭会话、回调返回错误、回调内开启乐观/悲观事务、回调内执行 DML 等场景见TestDomainAdvancedSessionPoolPutBackDirtySessionpkg/session/syssession/session_integration_test.go#L77-L230设计文档还提出在 TTL 高频建/销会话、统计信息损坏导致 async load 内存泄漏等场景下补充更多集成验证以确保持续运行无泄漏。未来可扩展方向设计文档在末尾提出了两个可后续支持的特性其中并发检测在现有源码中已具备雏形非线程安全方法的并发使用检测内部会话的unsafe计数器已能检测到多个 goroutine 并发执行非线程安全方法的情况unsafe 1时仅放行第一个操作并记录错误/竞争日志测试环境直接 panic 使测试失败见 pkg/session/syssession/session.go#L299-L311。未来可以将其扩展为对所有方法统一的并发使用检测能力更多会话状态的检测与重置目前归还时通过RollbackTxn重置事务状态未来可以检测系统变量等更多会话状态的变化并在归还时一并恢复进一步降低脏状态复用的风险。总结TiDB 的本次内部系统会话管理增强核心是用「回调式生命周期托管 所有权机制」取代「手动 Get/Put」WithSession通过 defer 保证会话必归还、错误与 panic 时必销毁Session包装层结合 owner 校验、inUse/unsafe计数与avoidReuse标记在机制层面杜绝了会话归还后被继续使用、跨模块并发使用、以及脏会话被复用等隐患session manager 的注册联动则彻底解决了旧实现中会话关闭但 map 引用残留的内存泄漏问题。对于 TiDB 内部模块的开发者而言接入增强池只需遵循一条规则——在WithSession的回调内使用会话不要跨回调持有会话引用。【免费下载链接】tidbTiDB is built for agentic workloads that grow unpredictably, with ACID guarantees and native support for transactions, analytics, and vector search. No data silos. No noisy neighbors. No infrastructure ceiling.项目地址: https://gitcode.com/GitHub_Trending/ti/tidb创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考