package common import ( "fmt" "sync" "sync/atomic" "github.com/shadow1ng/fscan/common/i18n" ) /* ConcurrencyMonitor.go - 并发监控器 监控两个层级的并发: 1. 主扫描器线程数 (-t 参数控制) 2. 插件内连接线程数 (-mt 参数控制) */ // ConcurrencyMonitor 并发监控器 type ConcurrencyMonitor struct { // 主扫描器层级 activePluginTasks int64 // 当前活跃的插件任务数 totalPluginTasks int64 // 总插件任务数 // 插件内连接层级 (每个插件的连接线程数) pluginConnections sync.Map // map[string]*PluginConnectionInfo mu sync.RWMutex } // PluginConnectionInfo 单个插件的连接信息 type PluginConnectionInfo struct { PluginName string // 插件名称 Target string // 目标地址 ActiveConnections int64 // 当前活跃连接数 TotalConnections int64 // 总连接数 } var ( globalConcurrencyMonitor *ConcurrencyMonitor concurrencyMutex sync.Once ) // GetConcurrencyMonitor 获取全局并发监控器 func GetConcurrencyMonitor() *ConcurrencyMonitor { concurrencyMutex.Do(func() { globalConcurrencyMonitor = &ConcurrencyMonitor{ activePluginTasks: 0, totalPluginTasks: 0, } }) return globalConcurrencyMonitor } // ============================================================================= // 主扫描器层级监控 // ============================================================================= // StartPluginTask 开始插件任务 func (m *ConcurrencyMonitor) StartPluginTask() { atomic.AddInt64(&m.activePluginTasks, 1) atomic.AddInt64(&m.totalPluginTasks, 1) } // FinishPluginTask 完成插件任务 func (m *ConcurrencyMonitor) FinishPluginTask() { atomic.AddInt64(&m.activePluginTasks, -1) } // GetPluginTaskStats 获取插件任务统计 func (m *ConcurrencyMonitor) GetPluginTaskStats() (active int64, total int64) { return atomic.LoadInt64(&m.activePluginTasks), atomic.LoadInt64(&m.totalPluginTasks) } // ============================================================================= // 插件内连接层级监控 // ============================================================================= // StartConnection 开始连接 func (m *ConcurrencyMonitor) StartConnection(pluginName, target string) { key := fmt.Sprintf("%s@%s", pluginName, target) value, _ := m.pluginConnections.LoadOrStore(key, &PluginConnectionInfo{ PluginName: pluginName, Target: target, }) info := value.(*PluginConnectionInfo) atomic.AddInt64(&info.ActiveConnections, 1) atomic.AddInt64(&info.TotalConnections, 1) } // FinishConnection 完成连接 func (m *ConcurrencyMonitor) FinishConnection(pluginName, target string) { key := fmt.Sprintf("%s@%s", pluginName, target) if value, ok := m.pluginConnections.Load(key); ok { info := value.(*PluginConnectionInfo) atomic.AddInt64(&info.ActiveConnections, -1) } } // 已移除未使用的 GetConnectionStats 方法 // GetTotalActiveConnections 获取总活跃连接数 func (m *ConcurrencyMonitor) GetTotalActiveConnections() int64 { var total int64 m.pluginConnections.Range(func(key, value interface{}) bool { info := value.(*PluginConnectionInfo) total += atomic.LoadInt64(&info.ActiveConnections) return true }) return total } // 已移除未使用的 Reset 方法 // GetConcurrencyStatus 获取并发状态字符串 func (m *ConcurrencyMonitor) GetConcurrencyStatus() string { activePlugins, _ := m.GetPluginTaskStats() totalConnections := m.GetTotalActiveConnections() if activePlugins == 0 && totalConnections == 0 { return "" } if totalConnections == 0 { return fmt.Sprintf("%s:%d", i18n.GetText("concurrency_plugin"), activePlugins) } return fmt.Sprintf("%s:%d %s:%d", i18n.GetText("concurrency_plugin"), activePlugins, i18n.GetText("concurrency_connection"), totalConnections) } // 已移除未使用的 GetDetailedStatus 方法