package db import ( "context" "crypto/rand" "database/sql" "encoding/hex" "fmt" "strconv" "sync" "time" "github.com/redis/go-redis/v9" ) const sessionTTL = 45 * time.Minute const maxSessions = 64 // maxPoolHandles 每会话缓存的 *sql.DB 句柄上限(按 database 分键),超出时关闭最久未用。 const maxPoolHandles = 8 type Session struct { ID string Conn ConnReq mu sync.Mutex sqlPool map[string]*sql.DB poolUsed map[string]time.Time // sqlPool 各键最近使用时间,供 LRU 轮换 redisClis map[int]redis.UniversalClient // 按 DB 序号缓存的 Redis 客户端(DB 是连接级参数) // agentHandles: agent 驱动(如达梦)按 database 缓存的 JDBC 句柄; // JVM 进程走全局池(agentPool),会话关闭只关句柄、不杀进程。 agentHandles map[string]string tunnel *SSHTunnel LastUsed time.Time } var ( sessionsMu sync.Mutex sessions = map[string]*Session{} ) func init() { go func() { t := time.NewTicker(2 * time.Minute) defer t.Stop() for range t.C { expireSessions() } }() } // newSessionID 生成 128bit 随机会话 ID;crypto/rand 失败时返回错误, // 不再静默降级为低熵时间戳。 func newSessionID() (string, error) { var b [16]byte if _, err := rand.Read(b[:]); err != nil { return "", fmt.Errorf("生成会话 ID 失败: %w", err) } return hex.EncodeToString(b[:]), nil } func OpenSession(c ConnReq) (string, error) { c.SQL = "" c.Table = "" c.Key = "" c.Pattern = "" c.SessionID = "" // agent 驱动(JDBC 桥)不走本地 SSH 端口转发(转发对 JVM 侧 TCP 无效) if isAgentDriver(c.Driver) && c.SSHEnabled { return "", fmt.Errorf("%s 走 JDBC agent,暂不支持 SSH 隧道", c.Driver) } dial, tun, err := ApplySSHTunnel(c) if err != nil { return "", err } if err := testDial(dial); err != nil { if tun != nil { tun.Close() } return "", fmt.Errorf("数据库: %w", err) } sessionsMu.Lock() if len(sessions) >= maxSessions { sessionsMu.Unlock() if tun != nil { tun.Close() } return "", fmt.Errorf("连接会话已满,请关闭不用的连接后重试") } id, err := newSessionID() if err != nil { sessionsMu.Unlock() if tun != nil { tun.Close() } return "", err } // Persist dialed endpoints (127.0.0.1:local) so later ops reuse the tunnel. stored := dial stored.SSHEnabled = false s := &Session{ ID: id, Conn: stored, sqlPool: map[string]*sql.DB{}, poolUsed: map[string]time.Time{}, agentHandles: map[string]string{}, tunnel: tun, LastUsed: time.Now(), } sessions[id] = s sessionsMu.Unlock() return id, nil } func GetSession(id string) (*Session, error) { if id == "" { return nil, fmt.Errorf("缺少 sessionId") } sessionsMu.Lock() s, ok := sessions[id] sessionsMu.Unlock() if !ok { return nil, fmt.Errorf("会话已失效,请重新连接") } s.mu.Lock() s.LastUsed = time.Now() s.mu.Unlock() return s, nil } func CloseSession(id string) { sessionsMu.Lock() s, ok := sessions[id] if ok { delete(sessions, id) } sessionsMu.Unlock() if ok { s.closeAll() } } func expireSessions() { cutoff := time.Now().Add(-sessionTTL) sessionsMu.Lock() var dead []*Session for id, s := range sessions { s.mu.Lock() stale := s.LastUsed.Before(cutoff) s.mu.Unlock() if stale { dead = append(dead, s) delete(sessions, id) } } sessionsMu.Unlock() for _, s := range dead { s.closeAll() } } func (s *Session) closeAll() { s.mu.Lock() defer s.mu.Unlock() for k, dbh := range s.sqlPool { _ = dbh.Close() delete(s.sqlPool, k) } s.poolUsed = map[string]time.Time{} for k, cli := range s.redisClis { _ = cli.Close() delete(s.redisClis, k) } // agent 句柄随会话释放;JVM 进程归全局池,进程生命周期不挂钩会话 if len(s.agentHandles) > 0 { if cli, err := getAgentClient(s.Conn.Driver); err == nil { for _, h := range s.agentHandles { ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) _ = cli.call(ctx, "close", map[string]any{"handle": h}, nil) cancel() } } s.agentHandles = map[string]string{} } if s.tunnel != nil { s.tunnel.Close() s.tunnel = nil } } // agentConn 按 database 懒建/复用 agent JDBC 句柄(模式同 sqlDB)。 // 调用方不得持有 s.mu(本函数内部自行加锁;getAgentClient 用独立的全局池锁)。 func (s *Session) agentConn(c ConnReq) (*AgentClient, string, error) { cli, err := getAgentClient(c.Driver) if err != nil { return nil, "", err } key := c.Database s.mu.Lock() if h, ok := s.agentHandles[key]; ok { s.mu.Unlock() return cli, h, nil } s.mu.Unlock() ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) defer cancel() h, err := agentOpen(ctx, cli, c) if err != nil { return nil, "", err } s.mu.Lock() defer s.mu.Unlock() if old, ok := s.agentHandles[key]; ok { // 双检:并发下已有别的 goroutine 建好同键句柄,回收多余句柄 go func() { ctx2, cancel2 := context.WithTimeout(context.Background(), 2*time.Second) defer cancel2() _ = cli.call(ctx2, "close", map[string]any{"handle": h}, nil) }() return cli, old, nil } s.agentHandles[key] = h return cli, h, nil } // touch 滚动续期(GetSession 已调用,池化路径同样经 GetSession 进入)。 func (s *Session) touch() { s.mu.Lock() s.LastUsed = time.Now() s.mu.Unlock() } // sqlDB 按 database 懒加载缓存 *sql.DB,会话内复用连接池; // 超过 maxPoolHandles 个句柄时关闭最久未用的池。 func (s *Session) sqlDB(database string) (*sql.DB, error) { // SQLite 的 database 是文件路径,对象树的 "main" 不能覆盖它 if driverKey(s.Conn.Driver) == "sqlite" { database = "" } s.mu.Lock() if dbh, ok := s.sqlPool[database]; ok { s.poolUsed[database] = time.Now() s.mu.Unlock() return dbh, nil } s.mu.Unlock() conn := s.Conn if database != "" { conn.Database = database } dbh, err := Open(conn) if err != nil { return nil, err } s.mu.Lock() defer s.mu.Unlock() // 双检:并发下可能已有别的 goroutine 建好同键池 if old, ok := s.sqlPool[database]; ok { s.poolUsed[database] = time.Now() _ = dbh.Close() return old, nil } // 超上限先淘汰最久未用 for len(s.sqlPool) >= maxPoolHandles { var lruKey string var lruTime time.Time for k := range s.sqlPool { t := s.poolUsed[k] if lruKey == "" || t.Before(lruTime) { lruKey, lruTime = k, t } } _ = s.sqlPool[lruKey].Close() delete(s.sqlPool, lruKey) delete(s.poolUsed, lruKey) } s.sqlPool[database] = dbh s.poolUsed[database] = time.Now() return dbh, nil } // redisClient 懒建会话级 Redis 单例,避免每次命令新建连接。 func (s *Session) redisClient() (redis.UniversalClient, error) { // 按连接配置缓存:DB 序号是 go-redis 的连接级参数(建连时钉死), // 对象树里切换 db0/db1 必须各建一个客户端,否则全部打到首个 db。 return s.redisClientFor(s.Conn) } func (s *Session) redisClientFor(c ConnReq) (redis.UniversalClient, error) { db := redisDBIndex(c) s.mu.Lock() defer s.mu.Unlock() if s.redisClis == nil { s.redisClis = map[int]redis.UniversalClient{} } if cli, ok := s.redisClis[db]; ok { return cli, nil } // 以会话凭据 + 目标 DB 建连(Sentinel/Cluster 模式 DB 恒为 0,天然只有一份) cli, err := openRedis(ConnReq{ Driver: s.Conn.Driver, Host: s.Conn.Host, Port: s.Conn.Port, User: s.Conn.User, Password: s.Conn.Password, Database: strconv.Itoa(db), RedisMode: s.Conn.RedisMode, RedisNodes: s.Conn.RedisNodes, MasterName: s.Conn.MasterName, SentinelPassword: s.Conn.SentinelPassword, SSL: s.Conn.SSL, SSLMode: s.Conn.SSLMode, SSLCA: s.Conn.SSLCA, SSLCert: s.Conn.SSLCert, SSLKey: s.Conn.SSLKey, SSHEnabled: s.Conn.SSHEnabled, SSHHost: s.Conn.SSHHost, SSHPort: s.Conn.SSHPort, SSHUser: s.Conn.SSHUser, SSHAuth: s.Conn.SSHAuth, SSHPassword: s.Conn.SSHPassword, SSHPrivateKey: s.Conn.SSHPrivateKey, SSHPassphrase: s.Conn.SSHPassphrase, }) if err != nil { return nil, err } s.redisClis[db] = cli return cli, nil } func (s *Session) ConnFor(req ConnReq) ConnReq { c := s.Conn // SQLite 的 database 是文件路径,对象树里的 "main" 不能覆盖它 if req.Database != "" && driverKey(c.Driver) != "sqlite" { c.Database = req.Database } c.SQL = req.SQL c.Table = req.Table c.Key = req.Key c.Pattern = req.Pattern c.Keyword = req.Keyword c.Format = req.Format c.Limit = req.Limit c.Offset = req.Offset c.ConfirmDanger = req.ConfirmDanger // 逐请求字段:新增 ConnReq 字段时若属于"每次请求携带"语义,必须在此补拷贝, // 否则会话路径会静默丢弃(前端全部走会话,等于功能失效) c.TimeoutMs = req.TimeoutMs c.Cursor = req.Cursor c.RequestID = req.RequestID c.SessionID = s.ID return c } func ResolveConn(req ConnReq) (ConnReq, error) { if req.SessionID == "" { return req, nil } s, err := GetSession(req.SessionID) if err != nil { return req, err } return s.ConnFor(req), nil }