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

     1  package limiter
     2  
     3  import "sync"
     4  
     5  type Key = interface{}
     6  
     7  // Manages resources with a limited number of concurrent slots for use for each key.
     8  type Instance struct {
     9  	SlotsPerKey int
    10  
    11  	mu sync.Mutex
    12  	// Limits concurrent use of a resource. Push into the channel to use a slot, and receive to free
    13  	// up a slot.
    14  	active map[Key]*activeValueType
    15  }
    16  
    17  type activeValueType struct {
    18  	ch   chan struct{}
    19  	refs int
    20  }
    21  
    22  type ActiveValueRef struct {
    23  	v *activeValueType
    24  	k Key
    25  	i *Instance
    26  }
    27  
    28  // Returns the limiting channel. Send to it to obtain a slot, and receive to release the slot.
    29  func (me ActiveValueRef) C() chan struct{} {
    30  	return me.v.ch
    31  }
    32  
    33  // Drop the reference to a key, this allows keys to be reclaimed when they're no longer in use.
    34  func (me ActiveValueRef) Drop() {
    35  	me.i.mu.Lock()
    36  	defer me.i.mu.Unlock()
    37  	me.v.refs--
    38  	if me.v.refs == 0 {
    39  		delete(me.i.active, me.k)
    40  	}
    41  }
    42  
    43  // Get a reference to the values for a key. You should make sure to call Drop exactly once on the
    44  // returned value when done.
    45  func (i *Instance) GetRef(key Key) ActiveValueRef {
    46  	i.mu.Lock()
    47  	defer i.mu.Unlock()
    48  	if i.active == nil {
    49  		i.active = make(map[Key]*activeValueType)
    50  	}
    51  	v, ok := i.active[key]
    52  	if !ok {
    53  		v = &activeValueType{
    54  			ch: make(chan struct{}, i.SlotsPerKey),
    55  		}
    56  		i.active[key] = v
    57  	}
    58  	v.refs++
    59  	return ActiveValueRef{
    60  		v: v,
    61  		k: key,
    62  		i: i,
    63  	}
    64  }