122 lines
2.3 KiB
Go
122 lines
2.3 KiB
Go
package collector
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"sync"
|
|
)
|
|
|
|
// Registry 采集器注册中心
|
|
type Registry struct {
|
|
collectors map[string]Collector
|
|
mu sync.RWMutex
|
|
}
|
|
|
|
// NewRegistry 创建采集器注册中心
|
|
func NewRegistry() *Registry {
|
|
return &Registry{
|
|
collectors: make(map[string]Collector),
|
|
}
|
|
}
|
|
|
|
// Register 注册一个采集器
|
|
func (r *Registry) Register(c Collector) error {
|
|
r.mu.Lock()
|
|
defer r.mu.Unlock()
|
|
|
|
name := c.Name()
|
|
if _, exists := r.collectors[name]; exists {
|
|
return fmt.Errorf("collector already registered: %s", name)
|
|
}
|
|
|
|
r.collectors[name] = c
|
|
return nil
|
|
}
|
|
|
|
// Get 获取指定名称的采集器
|
|
func (r *Registry) Get(name string) (Collector, bool) {
|
|
r.mu.RLock()
|
|
defer r.mu.RUnlock()
|
|
|
|
c, ok := r.collectors[name]
|
|
return c, ok
|
|
}
|
|
|
|
// List 列出所有已注册的采集器
|
|
func (r *Registry) List() []Collector {
|
|
r.mu.RLock()
|
|
defer r.mu.RUnlock()
|
|
|
|
result := make([]Collector, 0, len(r.collectors))
|
|
for _, c := range r.collectors {
|
|
result = append(result, c)
|
|
}
|
|
return result
|
|
}
|
|
|
|
// ListByCategory 按分类列出采集器
|
|
func (r *Registry) ListByCategory(cat Category) []Collector {
|
|
r.mu.RLock()
|
|
defer r.mu.RUnlock()
|
|
|
|
var result []Collector
|
|
for _, c := range r.collectors {
|
|
if c.Category() == cat {
|
|
result = append(result, c)
|
|
}
|
|
}
|
|
return result
|
|
}
|
|
|
|
// Available 返回当前系统上可用的采集器
|
|
func (r *Registry) Available(ctx context.Context) []Collector {
|
|
r.mu.RLock()
|
|
defer r.mu.RUnlock()
|
|
|
|
var result []Collector
|
|
for _, c := range r.collectors {
|
|
if c.IsAvailable(ctx) {
|
|
result = append(result, c)
|
|
}
|
|
}
|
|
return result
|
|
}
|
|
|
|
// ScanAll 使用所有可用的采集器扫描
|
|
func (r *Registry) ScanAll(ctx context.Context, opts ScanOptions) ([]*ScanResult, error) {
|
|
available := r.Available(ctx)
|
|
|
|
var (
|
|
results []*ScanResult
|
|
mu sync.Mutex
|
|
wg sync.WaitGroup
|
|
errs []error
|
|
)
|
|
|
|
for _, c := range available {
|
|
wg.Add(1)
|
|
go func(collector Collector) {
|
|
defer wg.Done()
|
|
|
|
result, err := collector.Scan(ctx, opts)
|
|
|
|
mu.Lock()
|
|
defer mu.Unlock()
|
|
|
|
if err != nil {
|
|
errs = append(errs, fmt.Errorf("%s: %w", collector.Name(), err))
|
|
return
|
|
}
|
|
results = append(results, result)
|
|
}(c)
|
|
}
|
|
|
|
wg.Wait()
|
|
|
|
if len(errs) > 0 && len(results) == 0 {
|
|
return nil, fmt.Errorf("all collectors failed: %v", errs)
|
|
}
|
|
|
|
return results, nil
|
|
}
|