86 lines
1.9 KiB
Go
86 lines
1.9 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()
|
|
}
|
|
|
|
selection := view.Select(snapshot.Query{
|
|
Now: request.Now,
|
|
Scheme: request.Scheme,
|
|
Upstreams: request.Upstreams,
|
|
RequiredTags: request.RequiredTags,
|
|
Exclude: request.Exclude,
|
|
SafetyMargin: request.SafetyMargin,
|
|
})
|
|
if selection.Len() == 0 {
|
|
return nil, ErrNoCandidate
|
|
}
|
|
|
|
start := int((d.cursor.Add(1) - 1) % uint64(selection.Len()))
|
|
for offset := 0; offset < selection.Len(); offset++ {
|
|
entry, ok := selection.EntryAt((start + offset) % selection.Len())
|
|
if !ok {
|
|
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
|
|
}
|