-
Notifications
You must be signed in to change notification settings - Fork 1
/
moderator.go
133 lines (115 loc) · 2.74 KB
/
moderator.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
/*
© 2020–present Harald Rudell <harald.rudell@gmail.com> (https://haraldrudell.github.io/haraldrudell/)
ISC License
*/
package parl
import (
"errors"
"fmt"
"sync"
"sync/atomic"
)
const (
defaultParallelism = 20
uint64MinusOne = ^uint64(0)
)
// Moderator invokes functions at a limited level of parallelism
type Moderator struct {
parallelism uint64
shutdownLock sync.Once
isShutdown AtomicBool
dataLock sync.Mutex
available uint64 // behind dataLock
waitCount uint64 // atomic
waitLock sync.Mutex
}
var ErrModeratorShutdown = errors.New("Moderator shut down")
// NewModerator creates a new Moderator used to limit parallelism
func NewModerator(parallelism uint64) (mo *Moderator) {
if parallelism == 0 {
parallelism = defaultParallelism
}
return &Moderator{parallelism: parallelism, available: parallelism}
}
// Do calls fn limited by the moderator’s parallelism.
// If the moderator is shut down, ErrModeratorShutdown is returned
func (mo *Moderator) Do(fn func() error) (err error) {
if fn == nil {
panic(New("Moderator.Do with nil function"))
}
if mo.isShutdown.IsTrue() {
return ErrModeratorShutdown
}
defer mo.returnTicket() // we will always get a ticket, and it should be returned
if err = mo.getTicket(); err != nil {
return // shutdown
}
return fn()
}
func (mo *Moderator) getTicket() (err error) {
if mo.getAvailableTicket() {
return // a ticket was available
}
// wait for ticket
mo.waitLock.Lock()
// we now have a ticket!
if mo.isShutdown.IsTrue() {
return ErrModeratorShutdown
}
return
}
func (mo *Moderator) getAvailableTicket() (hasTicket bool) {
mo.dataLock.Lock()
defer mo.dataLock.Unlock()
if mo.available > 0 {
mo.available--
if mo.available == 0 {
mo.waitLock.Lock() // the next thread will wait
}
return true
}
mo.waitCount++
return
}
func (mo *Moderator) returnTicket() {
mo.dataLock.Lock()
defer mo.dataLock.Unlock()
// hand ticket to the queue
if mo.waitCount > 0 {
mo.waitCount--
mo.waitLock.Unlock()
return
}
if mo.available == 0 {
mo.waitLock.Unlock()
}
// note an additional ticket available
mo.available++
}
func (mo *Moderator) Status() (parallelism uint64, available uint64, waiting uint64, isShutdown bool) {
parallelism = mo.parallelism
mo.dataLock.Lock()
available = mo.available
mo.dataLock.Unlock()
waiting = atomic.LoadUint64(&mo.waitCount)
isShutdown = mo.isShutdown.IsTrue()
return
}
func (mo *Moderator) String() (s string) {
p, a, w, sd := mo.Status()
if a > 0 {
s = fmt.Sprintf("available: %d(%d)", a, p)
} else {
s = fmt.Sprintf("waiting: %d(%d)", w, p)
}
if sd {
s += " shutdown"
}
return
}
func (mo *Moderator) shutdown() {
mo.isShutdown.Set()
}
func (mo *Moderator) Shutdown() {
mo.shutdownLock.Do(mo.shutdown)
}