github.com/argoproj/argo-cd/v2@v2.10.9/hack/gen-resources/util/sizedwaitgroup.go (about)

     1  // The MIT License (MIT)
     2  
     3  // Copyright (c) 2018 Rémy Mathieu
     4  
     5  // Permission is hereby granted, free of charge, to any person obtaining a copy
     6  // of this software and associated documentation files (the "Software"), to deal
     7  // in the Software without restriction, including without limitation the rights
     8  // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
     9  // copies of the Software, and to permit persons to whom the Software is
    10  // furnished to do so, subject to the following conditions:
    11  
    12  // The above copyright notice and this permission notice shall be included in all
    13  // copies or substantial portions of the Software.
    14  
    15  // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
    16  // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
    17  // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
    18  // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
    19  // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
    20  // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
    21  // SOFTWARE.
    22  // https://github.com/remeh/sizedwaitgroup
    23  
    24  // Based upon sync.WaitGroup, SizedWaitGroup allows to start multiple
    25  // routines and to wait for their end using the simple API.
    26  
    27  // Package util SizedWaitGroup adds the feature of limiting the maximum number of
    28  // concurrently started routines. It could for example be used to start
    29  // multiples routines querying a database but without sending too much
    30  // queries in order to not overload the given database.
    31  //
    32  // Rémy Mathieu © 2016
    33  package util
    34  
    35  import (
    36  	"context"
    37  	"math"
    38  	"sync"
    39  )
    40  
    41  // SizedWaitGroup has the same role and close to the
    42  // same API as the Golang sync.WaitGroup but adds a limit of
    43  // the amount of goroutines started concurrently.
    44  type SizedWaitGroup struct {
    45  	Size int
    46  
    47  	current chan struct{}
    48  	wg      sync.WaitGroup
    49  }
    50  
    51  // New creates a SizedWaitGroup.
    52  // The limit parameter is the maximum amount of
    53  // goroutines which can be started concurrently.
    54  func New(limit int) SizedWaitGroup {
    55  	size := math.MaxInt32 // 2^31 - 1
    56  	if limit > 0 {
    57  		size = limit
    58  	}
    59  	return SizedWaitGroup{
    60  		Size: size,
    61  
    62  		current: make(chan struct{}, size),
    63  		wg:      sync.WaitGroup{},
    64  	}
    65  }
    66  
    67  // Add increments the internal WaitGroup counter.
    68  // It can be blocking if the limit of spawned goroutines
    69  // has been reached. It will stop blocking when Done is
    70  // been called.
    71  //
    72  // See sync.WaitGroup documentation for more information.
    73  func (s *SizedWaitGroup) Add() {
    74  	_ = s.AddWithContext(context.Background())
    75  }
    76  
    77  // AddWithContext increments the internal WaitGroup counter.
    78  // It can be blocking if the limit of spawned goroutines
    79  // has been reached. It will stop blocking when Done is
    80  // been called, or when the context is canceled. Returns nil on
    81  // success or an error if the context is canceled before the lock
    82  // is acquired.
    83  //
    84  // See sync.WaitGroup documentation for more information.
    85  func (s *SizedWaitGroup) AddWithContext(ctx context.Context) error {
    86  	select {
    87  	case <-ctx.Done():
    88  		return ctx.Err()
    89  	case s.current <- struct{}{}:
    90  		break
    91  	}
    92  	s.wg.Add(1)
    93  	return nil
    94  }
    95  
    96  // Done decrements the SizedWaitGroup counter.
    97  // See sync.WaitGroup documentation for more information.
    98  func (s *SizedWaitGroup) Done() {
    99  	<-s.current
   100  	s.wg.Done()
   101  }
   102  
   103  // Wait blocks until the SizedWaitGroup counter is zero.
   104  // See sync.WaitGroup documentation for more information.
   105  func (s *SizedWaitGroup) Wait() {
   106  	s.wg.Wait()
   107  }