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