fscan/core/adaptive_pool.go

package core

import (
	"fmt"
	"sync"
	"sync/atomic"
	"time"

	"github.com/panjf2000/ants/v2"
	"github.com/shadow1ng/fscan/common"
	"github.com/shadow1ng/fscan/common/i18n"
)

// AdaptivePool 自适应线程池
// 封装 ants.PoolWithFunc，支持根据资源耗尽率动态调整线程数
type AdaptivePool struct {
	pool  *ants.PoolWithFunc
	state *common.State

	initialSize int
	minSize     int
	maxSize     int
	currentSize int32 // 原子操作

	// 监控参数
	checkInterval      time.Duration
	lastCheckNano      atomic.Int64 // UnixNano
	lastExhaustedCount int64
	lastPacketCount    int64

	// 阈值
	exhaustedThreshold float64 // 资源耗尽率阈值（触发降级）
	recoveryThreshold  float64 // 恢复阈值（允许升级）

	mu sync.Mutex
}

// NewAdaptivePool 创建自适应线程池
func NewAdaptivePool(size int, fn func(interface{}), state *common.State) (*AdaptivePool, error) {
	// 移除 WithPreAlloc(true)，在大规模扫描时预分配可能导致内存问题
	pool, err := ants.NewPoolWithFunc(size, fn)
	if err != nil {
		return nil, err
	}

	minSize := size / 4
	if minSize < 10 {
		minSize = 10
	}

	return &AdaptivePool{
		pool:               pool,
		state:              state,
		initialSize:        size,
		minSize:            minSize,
		maxSize:            size,
		currentSize:        int32(size),
		checkInterval:      time.Second,
		exhaustedThreshold: 0.10, // 10% 资源耗尽率触发降级
		recoveryThreshold:  0.02, // 2% 以下允许恢复
	}, nil
}

// Invoke 提交任务，并在适当时机检查是否需要调整线程数
func (ap *AdaptivePool) Invoke(task interface{}) error {
	ap.maybeAdjust()
	return ap.pool.Invoke(task)
}

// maybeAdjust 检查并可能调整线程池大小
// 使用原子 CAS 进行时间检查，99%+ 的调用零锁开销
func (ap *AdaptivePool) maybeAdjust() {
	lastCheck := ap.lastCheckNano.Load()
	now := time.Now().UnixNano()
	if now-lastCheck < int64(ap.checkInterval) {
		return
	}
	if !ap.lastCheckNano.CompareAndSwap(lastCheck, now) {
		return // 其他 goroutine 已在检查
	}

	// 获取当前计数
	currentExhausted := ap.state.GetResourceExhaustedCount()
	currentPackets := ap.state.GetPacketCount()

	ap.mu.Lock()
	// 计算增量（本周期内的耗尽率）
	deltaExhausted := currentExhausted - ap.lastExhaustedCount
	deltaPackets := currentPackets - ap.lastPacketCount

	ap.lastExhaustedCount = currentExhausted
	ap.lastPacketCount = currentPackets
	ap.mu.Unlock()

	// 需要足够的样本才能判断
	if deltaPackets < 100 {
		return
	}

	rate := float64(deltaExhausted) / float64(deltaPackets)
	currentSize := int(atomic.LoadInt32(&ap.currentSize))

	if rate > ap.exhaustedThreshold && currentSize > ap.minSize {
		// 降级：减少 20% 线程
		newSize := int(float64(currentSize) * 0.8)
		if newSize < ap.minSize {
			newSize = ap.minSize
		}
		ap.tune(newSize)
		common.LogInfo(i18n.Tr("adaptive_pool_resource_exhausted", fmt.Sprintf("%.1f", rate*100), currentSize, newSize))
	} else if rate < ap.recoveryThreshold && currentSize < ap.maxSize {
		// 恢复：增加 10% 线程（保守恢复）
		newSize := int(float64(currentSize) * 1.1)
		if newSize > ap.maxSize {
			newSize = ap.maxSize
		}
		if newSize > currentSize {
			ap.tune(newSize)
		}
	}
}

// tune 调整线程池大小
func (ap *AdaptivePool) tune(newSize int) {
	ap.pool.Tune(newSize)
	atomic.StoreInt32(&ap.currentSize, int32(newSize))
}

// Running 返回当前运行中的 goroutine 数量
func (ap *AdaptivePool) Running() int {
	return ap.pool.Running()
}

// Cap 返回当前池容量
func (ap *AdaptivePool) Cap() int {
	return int(atomic.LoadInt32(&ap.currentSize))
}

// Release 释放线程池
func (ap *AdaptivePool) Release() {
	ap.pool.Release()
}

// Wait 等待所有任务完成
func (ap *AdaptivePool) Wait() {
	// ants 没有原生 Wait，通过 Running() == 0 轮询
	for ap.pool.Running() > 0 {
		time.Sleep(10 * time.Millisecond)
	}
}