throttle.go

   1 package lightning
   2
   3 import (
   4         "sync"
   5         "sync/atomic"
   6 )
   7
   8 type throttle struct {
   9         Max       int
  10         wg        sync.WaitGroup
  11         ch        chan bool
  12         err       atomic.Value
  13         setupOnce sync.Once
  14         errorOnce sync.Once
  15 }
  16
  17 func (t *throttle) Acquire() {
  18         t.setupOnce.Do(func() { t.ch = make(chan bool, t.Max) })
  19         t.wg.Add(1)
  20         t.ch <- true
  21 }
  22
  23 func (t *throttle) Release() {
  24         t.wg.Done()
  25         <-t.ch
  26 }
  27
  28 func (t *throttle) Report(err error) {
  29         if err != nil {
  30                 t.errorOnce.Do(func() { t.err.Store(err) })
  31         }
  32 }
  33
  34 func (t *throttle) Err() error {
  35         err, _ := t.err.Load().(error)
  36         return err
  37 }
  38
  39 func (t *throttle) Wait() error {
  40         t.wg.Wait()
  41         return t.Err()
  42 }