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() {
  23                 if t.Max < 1 {
  24                         panic("throttle.Max < 1")
  25                 }
  26                 t.ch = make(chan bool, t.Max)
  27         })
  28         t.wg.Add(1)
  29         t.ch <- true
  30 }
  31
  32 func (t *throttle) Release() {
  33         t.wg.Done()
  34         <-t.ch
  35 }
  36
  37 func (t *throttle) Report(err error) {
  38         if err != nil {
  39                 t.errorOnce.Do(func() { t.err.Store(err) })
  40         }
  41 }
  42
  43 func (t *throttle) Err() error {
  44         err, _ := t.err.Load().(error)
  45         return err
  46 }
  47
  48 func (t *throttle) Wait() error {
  49         t.wg.Wait()
  50         return t.Err()
  51 }
  52
  53 func (t *throttle) Go(f func() error) error {
  54         t.Acquire()
  55         if t.Err() != nil {
  56                 t.Release()
  57                 return t.Err()
  58         }
  59         go func() {
  60                 t.Report(f())
  61                 t.Release()
  62         }()
  63         return nil
  64 }