github.com/anacrolix/torrent@v1.61.0/internal/alloclim/r.go (about)

     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. Returns true if the
    54  // reservation was released or cancelled for the first time.
    55  func (me *Reservation) Drop() (first bool) {
    56  	me.mu.Lock()
    57  	defer me.mu.Unlock()
    58  	if me.granted.IsSet() {
    59  		me.releaseOnce.Do(func() {
    60  			first = true
    61  			me.l.addValue(me.n)
    62  		})
    63  		return
    64  	}
    65  	if me.cancelled.Set() {
    66  		first = true
    67  		go me.l.doWakes()
    68  	}
    69  	return
    70  }
    71  
    72  func (me *Reservation) wake() bool {
    73  	me.mu.Lock()
    74  	defer me.mu.Unlock()
    75  	if me.cancelled.IsSet() {
    76  		return false
    77  	}
    78  	return me.granted.Set()
    79  }
    80  
    81  func (me *Reservation) Wait(ctx context.Context) error {
    82  	if me.n > me.l.Max {
    83  		return log.WithLevel(
    84  			log.Warning,
    85  			fmt.Errorf("reservation for %v exceeds limiter max %v", me.n, me.l.Max),
    86  		)
    87  	}
    88  	select {
    89  	case <-ctx.Done():
    90  	case <-me.granted.Done():
    91  	case <-me.cancelled.Done():
    92  	}
    93  	defer me.mu.Unlock()
    94  	me.mu.Lock()
    95  	switch {
    96  	case me.granted.IsSet():
    97  		return nil
    98  	case me.cancelled.IsSet():
    99  		return errors.New("reservation cancelled")
   100  	case ctx.Err() != nil:
   101  		return ctx.Err()
   102  	default:
   103  		panic("unexpected")
   104  	}
   105  }