argo-cd

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
}
108