150 lines
3.3 KiB
Go
150 lines
3.3 KiB
Go
package extraction
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"sort"
|
|
"sync"
|
|
"time"
|
|
)
|
|
|
|
type Fulfillment string
|
|
|
|
const (
|
|
Partial Fulfillment = "partial"
|
|
AllOrNothing Fulfillment = "allOrNothing"
|
|
)
|
|
|
|
type State string
|
|
|
|
const (
|
|
Available State = "AVAILABLE"
|
|
Extracted State = "EXTRACTED"
|
|
)
|
|
|
|
var ErrInsufficientProxies = errors.New("insufficient proxies")
|
|
|
|
type Candidate struct {
|
|
ID string
|
|
Protocol string
|
|
Region string
|
|
Carrier string
|
|
Upstream string
|
|
URL string
|
|
State State
|
|
ExpiresAt time.Time
|
|
LastCheckedAt time.Time
|
|
}
|
|
|
|
type Command struct {
|
|
Requested int
|
|
Fulfillment Fulfillment
|
|
Now time.Time
|
|
MinRemainingTTL time.Duration
|
|
MaxHealthCheckAge time.Duration
|
|
ReserveForGateway int
|
|
Protocols []string
|
|
Regions []string
|
|
Carriers []string
|
|
Upstreams []string
|
|
}
|
|
|
|
type Record struct {
|
|
ProxyID string
|
|
ClientID string
|
|
SourceIP string
|
|
RequestID string
|
|
Upstream string
|
|
ExtractedAt time.Time
|
|
ExpiresAt time.Time
|
|
}
|
|
|
|
type Result struct {
|
|
Requested int
|
|
Returned int
|
|
Items []Candidate
|
|
}
|
|
|
|
type Store interface {
|
|
Extract(context.Context, Command) (Result, error)
|
|
}
|
|
|
|
type MemoryStore struct {
|
|
mu sync.Mutex
|
|
candidates map[string]Candidate
|
|
}
|
|
|
|
func NewMemoryStore(candidates []Candidate) *MemoryStore {
|
|
items := make(map[string]Candidate, len(candidates))
|
|
for _, candidate := range candidates {
|
|
items[candidate.ID] = candidate
|
|
}
|
|
return &MemoryStore{candidates: items}
|
|
}
|
|
|
|
func (s *MemoryStore) Extract(_ context.Context, command Command) (Result, error) {
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
|
|
result := Result{Requested: command.Requested}
|
|
if command.Requested <= 0 {
|
|
return result, nil
|
|
}
|
|
eligible := make([]Candidate, 0, len(s.candidates))
|
|
for _, candidate := range s.candidates {
|
|
if eligibleForExtraction(candidate, command) {
|
|
eligible = append(eligible, candidate)
|
|
}
|
|
}
|
|
sort.Slice(eligible, func(i, j int) bool {
|
|
return eligible[i].ExpiresAt.After(eligible[j].ExpiresAt)
|
|
})
|
|
available := len(eligible) - command.ReserveForGateway
|
|
if available < 0 {
|
|
available = 0
|
|
}
|
|
if command.Fulfillment == AllOrNothing && available < command.Requested {
|
|
return result, ErrInsufficientProxies
|
|
}
|
|
count := command.Requested
|
|
if count > available {
|
|
count = available
|
|
}
|
|
for i := 0; i < count; i++ {
|
|
candidate := eligible[i]
|
|
candidate.State = Extracted
|
|
s.candidates[candidate.ID] = candidate
|
|
result.Items = append(result.Items, candidate)
|
|
}
|
|
result.Returned = len(result.Items)
|
|
return result, nil
|
|
}
|
|
|
|
func eligibleForExtraction(candidate Candidate, command Command) bool {
|
|
if candidate.State != Available {
|
|
return false
|
|
}
|
|
if !candidate.ExpiresAt.IsZero() && candidate.ExpiresAt.Sub(command.Now) < command.MinRemainingTTL {
|
|
return false
|
|
}
|
|
if command.MaxHealthCheckAge > 0 && command.Now.Sub(candidate.LastCheckedAt) > command.MaxHealthCheckAge {
|
|
return false
|
|
}
|
|
return matches(command.Protocols, candidate.Protocol) &&
|
|
matches(command.Regions, candidate.Region) &&
|
|
matches(command.Carriers, candidate.Carrier) &&
|
|
matches(command.Upstreams, candidate.Upstream)
|
|
}
|
|
|
|
func matches(allowed []string, value string) bool {
|
|
if len(allowed) == 0 {
|
|
return true
|
|
}
|
|
for _, candidate := range allowed {
|
|
if candidate == value {
|
|
return true
|
|
}
|
|
}
|
|
return false
|
|
}
|