throttle.go

   1 // Copyright (C) The Lightning Authors. All rights reserved.
   2 //
   3 // SPDX-License-Identifier: AGPL-3.0
   4
   5 package lightning
   6
   7 import (
   8         "sync"
   9         "sync/atomic"
  10 )
  11
  12 type throttle struct {
  13         Max       int
  14         wg        sync.WaitGroup
  15         ch        chan bool
  16         err       atomic.Value
  17         setupOnce sync.Once
  18         errorOnce sync.Once
  19 }
  20
  21 func (t *throttle) Acquire() {
  22         t.setupOnce.Do(func() { t.ch = make(chan bool, t.Max) })
  23         t.wg.Add(1)
  24         t.ch <- true
  25 }
  26
  27 func (t *throttle) Release() {
  28         t.wg.Done()
  29         <-t.ch
  30 }
  31
  32 func (t *throttle) Report(err error) {
  33         if err != nil {
  34                 t.errorOnce.Do(func() { t.err.Store(err) })
  35         }
  36 }
  37
  38 func (t *throttle) Err() error {
  39         err, _ := t.err.Load().(error)
  40         return err
  41 }
  42
  43 func (t *throttle) Wait() error {
  44         t.wg.Wait()
  45         return t.Err()
  46 }
  47
  48 func (t *throttle) Go(f func() error) error {
  49         t.Acquire()
  50         if t.Err() != nil {
  51                 return t.Err()
  52         }
  53         go func() {
  54                 t.Report(f())
  55                 t.Release()
  56         }()
  57         return nil
  58 }