internal/alloclim/r.go

   1 package alloclim
   2
   3 import (
   4         "context"
   5         "errors"
   6         "fmt"
   7         "sync"
   8
   9         "github.com/anacrolix/chansync"
  10         "github.com/anacrolix/log"
  11 )
  12
  13 type Reservation struct {
  14         l           *Limiter
  15         n           int64
  16         releaseOnce sync.Once
  17         mu          sync.Mutex
  18         granted     chansync.SetOnce
  19         cancelled   chansync.SetOnce
  20 }
  21
  22 // Releases the alloc claim if the reservation has been granted. Does nothing if it was cancelled.
  23 // Otherwise panics.
  24 func (me *Reservation) Release() {
  25         me.mu.Lock()
  26         defer me.mu.Unlock()
  27         switch {
  28         default:
  29                 panic("not resolved")
  30         case me.cancelled.IsSet():
  31                 return
  32         case me.granted.IsSet():
  33         }
  34         me.releaseOnce.Do(func() {
  35                 me.l.addValue(me.n)
  36         })
  37 }
  38
  39 // Cancel the reservation, returns false if it was already granted. You must still release if that's
  40 // the case. See Drop.
  41 func (me *Reservation) Cancel() bool {
  42         me.mu.Lock()
  43         defer me.mu.Unlock()
  44         if me.granted.IsSet() {
  45                 return false
  46         }
  47         if me.cancelled.Set() {
  48                 go me.l.doWakes()
  49         }
  50         return true
  51 }
  52
  53 // If the reservation is granted, release it, otherwise cancel the reservation.
  54 func (me *Reservation) Drop() {
  55         me.mu.Lock()
  56         defer me.mu.Unlock()
  57         if me.granted.IsSet() {
  58                 me.releaseOnce.Do(func() {
  59                         me.l.addValue(me.n)
  60                 })
  61                 return
  62         }
  63         if me.cancelled.Set() {
  64                 go me.l.doWakes()
  65         }
  66 }
  67
  68 func (me *Reservation) wake() bool {
  69         me.mu.Lock()
  70         defer me.mu.Unlock()
  71         if me.cancelled.IsSet() {
  72                 return false
  73         }
  74         return me.granted.Set()
  75 }
  76
  77 func (me *Reservation) Wait(ctx context.Context) error {
  78         if me.n > me.l.Max {
  79                 return log.WithLevel(
  80                         log.Warning,
  81                         fmt.Errorf("reservation for %v exceeds limiter max %v", me.n, me.l.Max),
  82                 )
  83         }
  84         select {
  85         case <-ctx.Done():
  86         case <-me.granted.Done():
  87         case <-me.cancelled.Done():
  88         }
  89         defer me.mu.Unlock()
  90         me.mu.Lock()
  91         switch {
  92         case me.granted.IsSet():
  93                 return nil
  94         case me.cancelled.IsSet():
  95                 return errors.New("reservation cancelled")
  96         case ctx.Err() != nil:
  97                 return ctx.Err()
  98         default:
  99                 panic("unexpected")
 100         }
 101 }