sync counter for processing things
Juan Batiz-Benet committed
Sep 18, 2014 at 19:42 UTC
7a785ded741de648ea50455afe32be0fe7da2470
1 file changed
+118
util/todocounter/counter.go
new
+118
@@ -0,0 +1,118 @@
1
+package todocounter
2
+
3
+import (
4
+ "sync"
5
+)
6
+
7
+// Counter records things remaining to process. It is needed for complicated
8
+// cases where multiple goroutines are spawned to process items, and they may
9
+// generate more items to process. For example, say a query over a set of nodes
10
+// may yield either a result value, or more nodes to query. Signaling is subtly
11
+// complicated, because the queue may be empty while items are being processed,
12
+// that will end up adding more items to the queue.
13
+//
14
+// Use Counter like this:
15
+//
16
+// todos := make(chan int, 10)
17
+// ctr := todoctr.NewCounter()
18
+//
19
+// process := func(item int) {
20
+// fmt.Println("processing %d\n...", item)
21
+//
22
+// // this task may randomly generate more tasks
23
+// if rand.Intn(5) == 0 {
24
+// todos<- item + 1
25
+// ctr.Increment(1) // increment counter for new task.
26
+// }
27
+//
28
+// ctr.Decrement(1) // decrement one to signal the task being done.
29
+// }
30
+//
31
+// // add some tasks.
32
+// todos<- 1
33
+// todos<- 2
34
+// todos<- 3
35
+// todos<- 4
36
+// ctr.Increment(4)
37
+//
38
+// for {
39
+// select {
40
+// case item := <- todos:
41
+// go process(item)
42
+// case <-ctr.Done():
43
+// fmt.Println("done processing everything.")
44
+// close(todos)
45
+// }
46
+// }
47
+type Counter interface {
48
+ // Incrememnt adds a number of todos to track.
49
+ // If the counter is **below** zero, it panics.
50
+ Increment(i uint32)
51
+
52
+ // Decrement removes a number of todos to track.
53
+ // If the count drops to zero, signals done and destroys the counter.
54
+ // If the count drops **below** zero, panics. It means you have tried to remove
55
+ // more things than you added, i.e. sync issues.
56
+ Decrement(i uint32)
57
+
58
+ // Done returns a channel to wait upon. Use it in selects:
59
+ //
60
+ // select {
61
+ // case <-ctr.Done():
62
+ // // done processing all items
63
+ // }
64
+ //
65
+ Done() <-chan struct{}
66
+}
67
+
68
+type todoCounter struct {
69
+ count int32
70
+ done chan struct{}
71
+ sync.RWMutex
72
+}
73
+
74
+// NewSyncCounter constructs a new counter
75
+func NewSyncCounter() Counter {
76
+ return &todoCounter{
77
+ done: make(chan struct{}),
78
+ }
79
+}
80
+
81
+func (c *todoCounter) Increment(i uint32) {
82
+ c.Lock()
83
+ defer c.Unlock()
84
+
85
+ if c.count < 0 {
86
+ panic("counter already signaled done. use a new counter.")
87
+ }
88
+
89
+ // increment count
90
+ c.count += int32(i)
91
+}
92
+
93
+// Decrement removes a number of todos to track.
94
+// If the count drops to zero, signals done and destroys the counter.
95
+// If the count drops **below** zero, panics. It means you have tried to remove
96
+// more things than you added, i.e. sync issues.
97
+func (c *todoCounter) Decrement(i uint32) {
98
+ c.Lock()
99
+ defer c.Unlock()
100
+
101
+ if c.count < 0 {
102
+ panic("counter already signaled done. probably have sync issues.")
103
+ }
104
+
105
+ if int32(i) > c.count {
106
+ panic("decrement amount creater than counter. sync issues.")
107
+ }
108
+
109
+ c.count -= int32(i)
110
+ if c.count == 0 { // done! signal it.
111
+ c.count-- // set it to -1 to prevent reuse
112
+ close(c.done) // a closed channel will always return nil
113
+ }
114
+}
115
+
116
+func (c *todoCounter) Done() <-chan struct{} {
117
+ return c.done
118
+}