mayfly-go/server/pkg/pool/cache_pool.go

package pool

import (
	"context"
	"math/rand"
	"mayfly-go/pkg/logx"
	"mayfly-go/pkg/utils/stringx"
	"sync"
	"time"
)

type cacheEntry[T Conn] struct {
	conn       T
	lastActive time.Time
}

type CachePool[T Conn] struct {
	factory func() (T, error)
	mu      sync.RWMutex
	cache   map[string]*cacheEntry[T] // 使用字符串键的缓存
	config  PoolConfig[T]
	closeCh chan struct{}
	closed  bool
}

func NewCachePool[T Conn](factory func() (T, error), opts ...Option[T]) *CachePool[T] {
	config := PoolConfig[T]{
		MaxConns:            1,
		IdleTimeout:         60 * time.Minute,
		WaitTimeout:         10 * time.Second,
		HealthCheckInterval: 10 * time.Minute,
	}
	for _, opt := range opts {
		opt(&config)
	}

	p := &CachePool[T]{
		factory: factory,
		cache:   make(map[string]*cacheEntry[T]),
		config:  config,
		closeCh: make(chan struct{}),
	}

	go p.backgroundMaintenance()
	return p
}

// Get 获取连接（自动创建或复用缓存连接）
func (p *CachePool[T]) Get(ctx context.Context, opts ...GetOption) (T, error) {
	var zero T

	options := defaultGetOptions // 默认更新 lastActive
	for _, apply := range opts {
		apply(&options)
	}

	// 先尝试加读锁，仅用于查找可用连接
	p.mu.RLock()
	if len(p.cache) >= p.config.MaxConns {
		keys := make([]string, 0, len(p.cache))
		for k := range p.cache {
			keys = append(keys, k)
		}

		randomKey := keys[rand.Intn(len(keys))]
		entry := p.cache[randomKey]
		conn := entry.conn

		if options.updateLastActive {
			// 更新最后活跃时间
			entry.lastActive = time.Now()
		}
		p.mu.RUnlock()
		return conn, nil
	}
	p.mu.RUnlock()

	if !options.newConn {
		return zero, ErrNoAvailableConn
	}

	// 没有找到可用连接，升级为写锁进行创建
	p.mu.Lock()
	defer p.mu.Unlock()

	if p.closed {
		return zero, ErrPoolClosed
	}

	// 再次检查是否已创建（防止并发）
	if len(p.cache) >= p.config.MaxConns {
		for _, entry := range p.cache {
			if options.updateLastActive {
				// 更新最后活跃时间
				entry.lastActive = time.Now()
			}
			return entry.conn, nil
		}
	}

	// 创建新连接
	conn, err := p.factory()
	if err != nil {
		return zero, err
	}

	p.cache[generateCacheKey()] = &cacheEntry[T]{
		conn:       conn,
		lastActive: time.Now(),
	}

	return conn, nil
}

// Put 将连接放回缓存
func (p *CachePool[T]) Put(conn T) error {
	p.mu.Lock()
	defer p.mu.Unlock()

	if p.closed {
		return conn.Close()
	}

	p.cache[generateCacheKey()] = &cacheEntry[T]{
		conn:       conn,
		lastActive: time.Now(),
	}

	// 如果超出最大连接数，清理最久未使用的
	if len(p.cache) > p.config.MaxConns {
		p.removeOldest()
	}

	return nil
}

// 移除最久未使用的连接
func (p *CachePool[T]) removeOldest() {
	var oldestKey string
	var oldestTime time.Time

	for key, entry := range p.cache {
		if oldestKey == "" || entry.lastActive.Before(oldestTime) {
			oldestKey = key
			oldestTime = entry.lastActive
		}
	}

	if oldestKey != "" {
		p.cache[oldestKey].conn.Close()
		delete(p.cache, oldestKey)
	}
}

// Close 关闭连接池
func (p *CachePool[T]) Close() {
	p.mu.Lock()
	defer p.mu.Unlock()

	if p.closed {
		return
	}

	p.closed = true
	close(p.closeCh)

	for key, entry := range p.cache {
		// 强制关闭连接
		p.closeConn(key, entry, true)
	}

	// 触发关闭回调
	if p.config.OnPoolClose != nil {
		p.config.OnPoolClose()
	}

	p.cache = make(map[string]*cacheEntry[T])
}

// Resize 动态调整大小
func (p *CachePool[T]) Resize(newSize int) {
	p.mu.Lock()
	defer p.mu.Unlock()

	if p.closed || newSize == p.config.MaxConns {
		return
	}

	p.config.MaxConns = newSize

	// 如果新大小小于当前缓存数量，清理多余的连接
	for len(p.cache) > newSize {
		p.removeOldest()
	}
}

// Stats 获取统计信息
func (p *CachePool[T]) Stats() PoolStats {
	p.mu.RLock()
	defer p.mu.RUnlock()

	return PoolStats{
		TotalConns: int32(len(p.cache)),
	}
}

// 后台维护协程
func (p *CachePool[T]) backgroundMaintenance() {
	ticker := time.NewTicker(p.config.HealthCheckInterval)
	defer ticker.Stop()

	for {
		select {
		case <-ticker.C:
			p.cleanupIdle()
		case <-p.closeCh:
			return
		}
	}
}

// 清理闲置超时的连接
func (p *CachePool[T]) cleanupIdle() {
	p.mu.Lock()
	defer p.mu.Unlock()

	cutoff := time.Now().Add(-p.config.IdleTimeout)
	for key, entry := range p.cache {
		if entry.lastActive.Before(cutoff) || !p.ping(entry.conn) {
			logx.Infof("cache pool - cleaning up idle connection, key: %s", key)
			// 如果连接超时或不可用，则关闭连接
			p.closeConn(key, entry, false)
		}
	}
}

func (p *CachePool[T]) ping(conn T) bool {
	done := make(chan struct{})
	var result bool
	go func() {
		result = conn.Ping() == nil
		close(done)
	}()
	select {
	case <-done:
		return result
	case <-time.After(2 * time.Second): // 设置超时
		logx.Debug("ping timeout")
		return false // 超时认为不可用
	}
}

func (p *CachePool[T]) closeConn(key string, entry *cacheEntry[T], force bool) bool {
	if !force {
		// 如果不是强制关闭且有连接关闭回调，则调用回调
		// 如果回调返回错误，则不关闭连接
		if onConnClose := p.config.OnConnClose; onConnClose != nil {
			if err := onConnClose(entry.conn); err != nil {
				logx.Infof("cache pool - connection close callback returned error, skip closing connection:: %v", err)
				return false
			}
		}
	}

	if err := entry.conn.Close(); err != nil {
		logx.Errorf("cache pool - closing connection error: %v", err)
	}
	delete(p.cache, key)
	return true
}

// 生成缓存键
func generateCacheKey() string {
	return stringx.RandUUID()
}