|
@@ -4,33 +4,33 @@
|
4
|
4
|
package utils
|
5
|
5
|
|
6
|
6
|
import (
|
|
7
|
+ "context"
|
7
|
8
|
"log"
|
8
|
9
|
"runtime/debug"
|
9
|
10
|
"time"
|
10
|
11
|
)
|
11
|
12
|
|
12
|
|
-// Semaphore is a counting semaphore. Note that a capacity of n requires O(n) storage.
|
|
13
|
+type e struct{}
|
|
14
|
+
|
|
15
|
+// Semaphore is a counting semaphore.
|
13
|
16
|
// A semaphore of capacity 1 can be used as a trylock.
|
14
|
|
-type Semaphore (chan bool)
|
|
17
|
+type Semaphore (chan e)
|
15
|
18
|
|
16
|
19
|
// Initialize initializes a semaphore to a given capacity.
|
17
|
20
|
func (semaphore *Semaphore) Initialize(capacity int) {
|
18
|
|
- *semaphore = make(chan bool, capacity)
|
19
|
|
- for i := 0; i < capacity; i++ {
|
20
|
|
- (*semaphore) <- true
|
21
|
|
- }
|
|
21
|
+ *semaphore = make(chan e, capacity)
|
22
|
22
|
}
|
23
|
23
|
|
24
|
24
|
// Acquire acquires a semaphore, blocking if necessary.
|
25
|
25
|
func (semaphore *Semaphore) Acquire() {
|
26
|
|
- <-(*semaphore)
|
|
26
|
+ (*semaphore) <- e{}
|
27
|
27
|
}
|
28
|
28
|
|
29
|
29
|
// TryAcquire tries to acquire a semaphore, returning whether the acquire was
|
30
|
30
|
// successful. It never blocks.
|
31
|
31
|
func (semaphore *Semaphore) TryAcquire() (acquired bool) {
|
32
|
32
|
select {
|
33
|
|
- case <-(*semaphore):
|
|
33
|
+ case (*semaphore) <- e{}:
|
34
|
34
|
return true
|
35
|
35
|
default:
|
36
|
36
|
return false
|
|
@@ -47,7 +47,7 @@ func (semaphore *Semaphore) AcquireWithTimeout(timeout time.Duration) (acquired
|
47
|
47
|
|
48
|
48
|
timer := time.NewTimer(timeout)
|
49
|
49
|
select {
|
50
|
|
- case <-(*semaphore):
|
|
50
|
+ case (*semaphore) <- e{}:
|
51
|
51
|
acquired = true
|
52
|
52
|
case <-timer.C:
|
53
|
53
|
acquired = false
|
|
@@ -56,11 +56,24 @@ func (semaphore *Semaphore) AcquireWithTimeout(timeout time.Duration) (acquired
|
56
|
56
|
return
|
57
|
57
|
}
|
58
|
58
|
|
|
59
|
+// AcquireWithContext tries to acquire a semaphore, blocking at most until
|
|
60
|
+// the context expires. It returns whether the acquire was successful.
|
|
61
|
+// Note that if the context is already expired, the acquire may succeed anyway.
|
|
62
|
+func (semaphore *Semaphore) AcquireWithContext(ctx context.Context) (acquired bool) {
|
|
63
|
+ select {
|
|
64
|
+ case (*semaphore) <- e{}:
|
|
65
|
+ acquired = true
|
|
66
|
+ case <-ctx.Done():
|
|
67
|
+ acquired = false
|
|
68
|
+ }
|
|
69
|
+ return
|
|
70
|
+}
|
|
71
|
+
|
59
|
72
|
// Release releases a semaphore. It never blocks. (This is not a license
|
60
|
73
|
// to program spurious releases.)
|
61
|
74
|
func (semaphore *Semaphore) Release() {
|
62
|
75
|
select {
|
63
|
|
- case (*semaphore) <- true:
|
|
76
|
+ case <-(*semaphore):
|
64
|
77
|
// good
|
65
|
78
|
default:
|
66
|
79
|
// spurious release
|