104 lines
2.4 KiB
Go
104 lines
2.4 KiB
Go
package provider
|
|
|
|
import (
|
|
"errors"
|
|
"math"
|
|
"sync"
|
|
|
|
"proxy-pool/internal/domain/upstream"
|
|
)
|
|
|
|
var ErrInvalidStatsRecorder = errors.New("invalid Provider stats recorder")
|
|
|
|
type Stats struct {
|
|
UpstreamID string
|
|
ConsecutiveEmptyFetch int64
|
|
FetchErrorCount int64
|
|
}
|
|
|
|
type StatsReader interface {
|
|
ReadProviderStats([]string) []Stats
|
|
}
|
|
|
|
type StatsRetainer interface {
|
|
RetainProviderStats([]string)
|
|
}
|
|
|
|
type StatsRecorder struct {
|
|
mu sync.Mutex
|
|
maximum int
|
|
byID map[string]Stats
|
|
}
|
|
|
|
func NewStatsRecorder(maximum int) (*StatsRecorder, error) {
|
|
if maximum <= 0 {
|
|
return nil, ErrInvalidStatsRecorder
|
|
}
|
|
return &StatsRecorder{maximum: maximum, byID: make(map[string]Stats)}, nil
|
|
}
|
|
|
|
func (recorder *StatsRecorder) Record(result Result) {
|
|
if recorder == nil || result.UpstreamID == "" {
|
|
return
|
|
}
|
|
recorder.mu.Lock()
|
|
defer recorder.mu.Unlock()
|
|
stats, exists := recorder.byID[result.UpstreamID]
|
|
if !exists {
|
|
if len(recorder.byID) >= recorder.maximum {
|
|
return
|
|
}
|
|
stats.UpstreamID = result.UpstreamID
|
|
}
|
|
switch result.Class {
|
|
case upstream.FetchEmpty:
|
|
if stats.ConsecutiveEmptyFetch < math.MaxInt64 {
|
|
stats.ConsecutiveEmptyFetch++
|
|
}
|
|
case upstream.FetchValid, upstream.FetchDuplicateOnly:
|
|
stats.ConsecutiveEmptyFetch = 0
|
|
case upstream.FetchError:
|
|
if stats.FetchErrorCount < math.MaxInt64 {
|
|
stats.FetchErrorCount++
|
|
}
|
|
default:
|
|
return
|
|
}
|
|
recorder.byID[result.UpstreamID] = stats
|
|
}
|
|
|
|
func (recorder *StatsRecorder) ReadProviderStats(upstreamIDs []string) []Stats {
|
|
result := make([]Stats, len(upstreamIDs))
|
|
if recorder == nil {
|
|
return result
|
|
}
|
|
recorder.mu.Lock()
|
|
defer recorder.mu.Unlock()
|
|
for index, upstreamID := range upstreamIDs {
|
|
result[index] = recorder.byID[upstreamID]
|
|
result[index].UpstreamID = upstreamID
|
|
}
|
|
return result
|
|
}
|
|
|
|
// RetainProviderStats removes observations for upstreams no longer present in
|
|
// the complete configuration. Disabled but configured upstreams must be kept.
|
|
func (recorder *StatsRecorder) RetainProviderStats(upstreamIDs []string) {
|
|
if recorder == nil {
|
|
return
|
|
}
|
|
retained := make(map[string]struct{}, len(upstreamIDs))
|
|
for _, upstreamID := range upstreamIDs {
|
|
if upstreamID != "" {
|
|
retained[upstreamID] = struct{}{}
|
|
}
|
|
}
|
|
recorder.mu.Lock()
|
|
defer recorder.mu.Unlock()
|
|
for upstreamID := range recorder.byID {
|
|
if _, keep := retained[upstreamID]; !keep {
|
|
delete(recorder.byID, upstreamID)
|
|
}
|
|
}
|
|
}
|