102 lines
1.9 KiB
Go
102 lines
1.9 KiB
Go
package alloclim
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"fmt"
|
|
"sync"
|
|
|
|
"github.com/anacrolix/chansync"
|
|
"github.com/anacrolix/log"
|
|
)
|
|
|
|
type Reservation struct {
|
|
l *Limiter
|
|
n int64
|
|
releaseOnce sync.Once
|
|
mu sync.Mutex
|
|
granted chansync.SetOnce
|
|
cancelled chansync.SetOnce
|
|
}
|
|
|
|
// Releases the alloc claim if the reservation has been granted. Does nothing if it was cancelled.
|
|
// Otherwise panics.
|
|
func (me *Reservation) Release() {
|
|
me.mu.Lock()
|
|
defer me.mu.Unlock()
|
|
switch {
|
|
default:
|
|
panic("not resolved")
|
|
case me.cancelled.IsSet():
|
|
return
|
|
case me.granted.IsSet():
|
|
}
|
|
me.releaseOnce.Do(func() {
|
|
me.l.addValue(me.n)
|
|
})
|
|
}
|
|
|
|
// Cancel the reservation, returns false if it was already granted. You must still release if that's
|
|
// the case. See Drop.
|
|
func (me *Reservation) Cancel() bool {
|
|
me.mu.Lock()
|
|
defer me.mu.Unlock()
|
|
if me.granted.IsSet() {
|
|
return false
|
|
}
|
|
if me.cancelled.Set() {
|
|
go me.l.doWakes()
|
|
}
|
|
return true
|
|
}
|
|
|
|
// If the reservation is granted, release it, otherwise cancel the reservation.
|
|
func (me *Reservation) Drop() {
|
|
me.mu.Lock()
|
|
defer me.mu.Unlock()
|
|
if me.granted.IsSet() {
|
|
me.releaseOnce.Do(func() {
|
|
me.l.addValue(me.n)
|
|
})
|
|
return
|
|
}
|
|
if me.cancelled.Set() {
|
|
go me.l.doWakes()
|
|
}
|
|
}
|
|
|
|
func (me *Reservation) wake() bool {
|
|
me.mu.Lock()
|
|
defer me.mu.Unlock()
|
|
if me.cancelled.IsSet() {
|
|
return false
|
|
}
|
|
return me.granted.Set()
|
|
}
|
|
|
|
func (me *Reservation) Wait(ctx context.Context) error {
|
|
if me.n > me.l.Max {
|
|
return log.WithLevel(
|
|
log.Warning,
|
|
fmt.Errorf("reservation for %v exceeds limiter max %v", me.n, me.l.Max),
|
|
)
|
|
}
|
|
select {
|
|
case <-ctx.Done():
|
|
case <-me.granted.Done():
|
|
case <-me.cancelled.Done():
|
|
}
|
|
defer me.mu.Unlock()
|
|
me.mu.Lock()
|
|
switch {
|
|
case me.granted.IsSet():
|
|
return nil
|
|
case me.cancelled.IsSet():
|
|
return errors.New("reservation cancelled")
|
|
case ctx.Err() != nil:
|
|
return ctx.Err()
|
|
default:
|
|
panic("unexpected")
|
|
}
|
|
}
|