internal/limiter/limiter.go

   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 }