forked from hashicorp/go-memdb
-
Notifications
You must be signed in to change notification settings - Fork 0
/
watch_test.go
186 lines (164 loc) · 3.84 KB
/
watch_test.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
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
package memdb
import (
"context"
"fmt"
"testing"
"time"
)
// testWatch makes a bunch of watch channels based on the given size and fires
// the one at the given fire index to make sure it's detected (or a timeout
// occurs if the fire index isn't hit). useCtx parameterizes whether the context
// based watch is used or timer based.
func testWatch(size, fire int, useCtx bool) error {
shouldTimeout := true
ws := NewWatchSet()
for i := 0; i < size; i++ {
watchCh := make(chan struct{})
ws.Add(watchCh)
if fire == i {
close(watchCh)
shouldTimeout = false
}
}
var timeoutCh chan time.Time
var ctx context.Context
var cancelFn context.CancelFunc
if useCtx {
ctx, cancelFn = context.WithCancel(context.Background())
defer cancelFn()
} else {
timeoutCh = make(chan time.Time)
}
doneCh := make(chan bool, 1)
go func() {
if useCtx {
doneCh <- ws.WatchCtx(ctx) != nil
} else {
doneCh <- ws.Watch(timeoutCh)
}
}()
if shouldTimeout {
select {
case <-doneCh:
return fmt.Errorf("should not trigger")
default:
}
if useCtx {
cancelFn()
} else {
close(timeoutCh)
}
select {
case didTimeout := <-doneCh:
if !didTimeout {
return fmt.Errorf("should have timed out")
}
case <-time.After(10 * time.Second):
return fmt.Errorf("should have timed out")
}
} else {
select {
case didTimeout := <-doneCh:
if didTimeout {
return fmt.Errorf("should not have timed out")
}
case <-time.After(10 * time.Second):
return fmt.Errorf("should have triggered")
}
if useCtx {
cancelFn()
} else {
close(timeoutCh)
}
}
return nil
}
func TestWatch(t *testing.T) {
testFactory := func(useCtx bool) func(t *testing.T) {
return func(t *testing.T) {
// Sweep through a bunch of chunks to hit the various cases of dividing
// the work into watchFew calls.
for size := 0; size < 3*aFew; size++ {
// Fire each possible channel slot.
for fire := 0; fire < size; fire++ {
if err := testWatch(size, fire, useCtx); err != nil {
t.Fatalf("err %d %d: %v", size, fire, err)
}
}
// Run a timeout case as well.
fire := -1
if err := testWatch(size, fire, useCtx); err != nil {
t.Fatalf("err %d %d: %v", size, fire, err)
}
}
}
}
t.Run("Timer", testFactory(false))
t.Run("Context", testFactory(true))
}
func TestWatch_AddWithLimit(t *testing.T) {
// Make sure nil doesn't crash.
{
var ws WatchSet
ch := make(chan struct{})
ws.AddWithLimit(10, ch, ch)
}
// Run a case where we trigger a channel that should be in
// there.
{
ws := NewWatchSet()
inCh := make(chan struct{})
altCh := make(chan struct{})
ws.AddWithLimit(1, inCh, altCh)
nopeCh := make(chan struct{})
ws.AddWithLimit(1, nopeCh, altCh)
close(inCh)
didTimeout := ws.Watch(time.After(1 * time.Second))
if didTimeout {
t.Fatalf("bad")
}
}
// Run a case where we trigger the alt channel that should have
// been added.
{
ws := NewWatchSet()
inCh := make(chan struct{})
altCh := make(chan struct{})
ws.AddWithLimit(1, inCh, altCh)
nopeCh := make(chan struct{})
ws.AddWithLimit(1, nopeCh, altCh)
close(altCh)
didTimeout := ws.Watch(time.After(1 * time.Second))
if didTimeout {
t.Fatalf("bad")
}
}
// Run a case where we trigger the nope channel that should not have
// been added.
{
ws := NewWatchSet()
inCh := make(chan struct{})
altCh := make(chan struct{})
ws.AddWithLimit(1, inCh, altCh)
nopeCh := make(chan struct{})
ws.AddWithLimit(1, nopeCh, altCh)
close(nopeCh)
didTimeout := ws.Watch(time.After(1 * time.Second))
if !didTimeout {
t.Fatalf("bad")
}
}
}
func BenchmarkWatch(b *testing.B) {
ws := NewWatchSet()
for i := 0; i < 1024; i++ {
watchCh := make(chan struct{})
ws.Add(watchCh)
}
timeoutCh := make(chan time.Time)
close(timeoutCh)
b.ResetTimer()
for i := 0; i < b.N; i++ {
ws.Watch(timeoutCh)
}
}