opsxdb/apps/server/internal/db/session.go

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
}