124 lines
2.5 KiB
Go
124 lines
2.5 KiB
Go
package agenttask
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"fmt"
|
|
"sync"
|
|
)
|
|
|
|
var (
|
|
ErrAlreadyAdmitted = errors.New("agenttask work already has an active dispatch ticket")
|
|
ErrProviderLimitMismatch = errors.New("agenttask provider capacity changed within one scheduler revision")
|
|
)
|
|
|
|
type DispatchCandidate struct {
|
|
ProjectID ProjectID
|
|
WorkspaceID WorkspaceID
|
|
WorkUnitID WorkUnitID
|
|
AttemptID AttemptID
|
|
Target ExecutionTarget
|
|
}
|
|
|
|
func (c DispatchCandidate) key() string {
|
|
return fmt.Sprintf("%s\x00%s\x00%s", c.ProjectID, c.WorkUnitID, c.AttemptID)
|
|
}
|
|
|
|
type schedulerPool struct {
|
|
limit int
|
|
inUse int
|
|
notify chan struct{}
|
|
}
|
|
|
|
type Scheduler struct {
|
|
mu sync.Mutex
|
|
pools map[string]*schedulerPool
|
|
tickets map[string]struct{}
|
|
}
|
|
|
|
func NewScheduler() *Scheduler {
|
|
return &Scheduler{
|
|
pools: make(map[string]*schedulerPool),
|
|
tickets: make(map[string]struct{}),
|
|
}
|
|
}
|
|
|
|
func (s *Scheduler) Acquire(
|
|
ctx context.Context,
|
|
candidate DispatchCandidate,
|
|
) (*DispatchLease, error) {
|
|
if candidate.Target.Capacity <= 0 {
|
|
return nil, fmt.Errorf("%w: capacity must be positive", ErrProviderLimitMismatch)
|
|
}
|
|
ticketKey := candidate.key()
|
|
poolKey := candidate.Target.PoolKey()
|
|
for {
|
|
s.mu.Lock()
|
|
if _, exists := s.tickets[ticketKey]; exists {
|
|
s.mu.Unlock()
|
|
return nil, ErrAlreadyAdmitted
|
|
}
|
|
pool := s.pools[poolKey]
|
|
if pool == nil {
|
|
pool = &schedulerPool{
|
|
limit: candidate.Target.Capacity,
|
|
notify: make(chan struct{}),
|
|
}
|
|
s.pools[poolKey] = pool
|
|
} else if pool.limit != candidate.Target.Capacity {
|
|
s.mu.Unlock()
|
|
return nil, ErrProviderLimitMismatch
|
|
}
|
|
if pool.inUse < pool.limit {
|
|
pool.inUse++
|
|
s.tickets[ticketKey] = struct{}{}
|
|
s.mu.Unlock()
|
|
return &DispatchLease{
|
|
scheduler: s,
|
|
poolKey: poolKey,
|
|
ticketKey: ticketKey,
|
|
}, nil
|
|
}
|
|
notify := pool.notify
|
|
s.mu.Unlock()
|
|
select {
|
|
case <-ctx.Done():
|
|
return nil, ctx.Err()
|
|
case <-notify:
|
|
}
|
|
}
|
|
}
|
|
|
|
type DispatchLease struct {
|
|
once sync.Once
|
|
scheduler *Scheduler
|
|
poolKey string
|
|
ticketKey string
|
|
}
|
|
|
|
func (l *DispatchLease) Release() {
|
|
if l == nil || l.scheduler == nil {
|
|
return
|
|
}
|
|
l.once.Do(func() {
|
|
s := l.scheduler
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
pool := s.pools[l.poolKey]
|
|
if pool != nil && pool.inUse > 0 {
|
|
pool.inUse--
|
|
close(pool.notify)
|
|
pool.notify = make(chan struct{})
|
|
}
|
|
delete(s.tickets, l.ticketKey)
|
|
})
|
|
}
|
|
|
|
func (s *Scheduler) Active(poolKey string) int {
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
if pool := s.pools[poolKey]; pool != nil {
|
|
return pool.inUse
|
|
}
|
|
return 0
|
|
}
|