-
-
Notifications
You must be signed in to change notification settings - Fork 158
/
atomic_maybe_work.go
87 lines (75 loc) · 2.28 KB
/
atomic_maybe_work.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
package kgo
import "sync/atomic"
// a helper type for some places
type atomicBool uint32
func (b *atomicBool) set(v bool) {
if v {
atomic.StoreUint32((*uint32)(b), 1)
} else {
atomic.StoreUint32((*uint32)(b), 0)
}
}
func (b *atomicBool) get() bool { return atomic.LoadUint32((*uint32)(b)) == 1 }
func (b *atomicBool) swap(v bool) bool {
var swap uint32
if v {
swap = 1
}
return atomic.SwapUint32((*uint32)(b), swap) == 1
}
const (
stateUnstarted = iota
stateWorking
stateContinueWorking
)
type workLoop struct{ state uint32 }
// maybeBegin returns whether a work loop should begin.
func (l *workLoop) maybeBegin() bool {
var state uint32
var done bool
for !done {
switch state = atomic.LoadUint32(&l.state); state {
case stateUnstarted:
done = atomic.CompareAndSwapUint32(&l.state, state, stateWorking)
state = stateWorking
case stateWorking:
done = atomic.CompareAndSwapUint32(&l.state, state, stateContinueWorking)
state = stateContinueWorking
case stateContinueWorking:
done = true
}
}
return state == stateWorking
}
// maybeFinish demotes loop's internal state and returns whether work should
// keep going. This function should be called before looping to continue
// work.
//
// If again is true, this will avoid demoting from working to not
// working. Again would be true if the loop knows it should continue working;
// calling this function is necessary even in this case to update loop's
// internal state.
//
// This function is a no-op if the loop is already finished, but generally,
// since the loop itself calls MaybeFinish after it has been started, this
// should never be called if the loop is unstarted.
func (l *workLoop) maybeFinish(again bool) bool {
switch state := atomic.LoadUint32(&l.state); state {
// Working:
// If again, we know we should continue; keep our state.
// If not again, we try to downgrade state and stop.
// If we cannot, then something slipped in to say keep going.
case stateWorking:
if !again {
again = !atomic.CompareAndSwapUint32(&l.state, state, stateUnstarted)
}
// Continue: demote ourself and run again no matter what.
case stateContinueWorking:
atomic.StoreUint32(&l.state, stateWorking)
again = true
}
return again
}
func (l *workLoop) hardFinish() {
atomic.StoreUint32(&l.state, stateUnstarted)
}