110 lines
2.5 KiB
Go
110 lines
2.5 KiB
Go
package dispatch
|
|
|
|
import (
|
|
"errors"
|
|
"sync/atomic"
|
|
"time"
|
|
|
|
proxyDomain "github.com/proxy-pool/proxy-pool/internal/domain/proxy"
|
|
"github.com/proxy-pool/proxy-pool/internal/gateway/snapshot"
|
|
)
|
|
|
|
var ErrNoCandidate = errors.New("no local proxy candidate is available")
|
|
|
|
type Request struct {
|
|
Now time.Time
|
|
Scheme proxyDomain.Scheme
|
|
Upstreams []string
|
|
RequiredTags map[string]string
|
|
Exclude map[string]struct{}
|
|
SafetyMargin time.Duration
|
|
}
|
|
|
|
type Lease struct {
|
|
Proxy proxyDomain.Proxy
|
|
Epoch uint64
|
|
Version uint64
|
|
reserved *proxyDomain.Reservation
|
|
}
|
|
|
|
func (l *Lease) Commit() error { return l.reserved.Commit() }
|
|
func (l *Lease) Cancel() error { return l.reserved.Cancel() }
|
|
func (l *Lease) Release() error { return l.reserved.Release() }
|
|
|
|
type Dispatcher struct {
|
|
store *snapshot.Store
|
|
cursor atomic.Uint64
|
|
}
|
|
|
|
func New(store *snapshot.Store) *Dispatcher {
|
|
return &Dispatcher{store: store}
|
|
}
|
|
|
|
func (d *Dispatcher) Acquire(request Request) (*Lease, error) {
|
|
if d == nil || d.store == nil {
|
|
return nil, ErrNoCandidate
|
|
}
|
|
view := d.store.Current()
|
|
if view == nil || len(view.Entries) == 0 {
|
|
return nil, ErrNoCandidate
|
|
}
|
|
if request.Now.IsZero() {
|
|
request.Now = time.Now().UTC()
|
|
}
|
|
|
|
start := int((d.cursor.Add(1) - 1) % uint64(len(view.Entries)))
|
|
for offset := 0; offset < len(view.Entries); offset++ {
|
|
entry := view.Entries[(start+offset)%len(view.Entries)]
|
|
if !eligible(entry.Proxy, request) {
|
|
continue
|
|
}
|
|
reservation, ok := entry.Runtime.Reserve()
|
|
if !ok {
|
|
continue
|
|
}
|
|
return &Lease{
|
|
Proxy: entry.Proxy,
|
|
Epoch: view.Epoch,
|
|
Version: view.Version,
|
|
reserved: reservation,
|
|
}, nil
|
|
}
|
|
return nil, ErrNoCandidate
|
|
}
|
|
|
|
func eligible(candidate proxyDomain.Proxy, request Request) bool {
|
|
if candidate.State != proxyDomain.StateAvailable {
|
|
return false
|
|
}
|
|
if request.Scheme != "" && candidate.Scheme != request.Scheme {
|
|
return false
|
|
}
|
|
if _, excluded := request.Exclude[candidate.ID]; excluded {
|
|
return false
|
|
}
|
|
if candidate.ExpiresAt != nil && !candidate.ExpiresAt.After(request.Now.Add(request.SafetyMargin)) {
|
|
return false
|
|
}
|
|
if !contains(request.Upstreams, candidate.SourceUpstream) {
|
|
return false
|
|
}
|
|
for key, value := range request.RequiredTags {
|
|
if candidate.Tags[key] != value {
|
|
return false
|
|
}
|
|
}
|
|
return true
|
|
}
|
|
|
|
func contains(allowed []string, value string) bool {
|
|
if len(allowed) == 0 {
|
|
return true
|
|
}
|
|
for _, candidate := range allowed {
|
|
if candidate == value {
|
|
return true
|
|
}
|
|
}
|
|
return false
|
|
}
|