notifs: rate limit notifications
Juan Batiz-Benet committed
Feb 1, 2015 at 04:25 UTC
3f5374717d79e334d0411cb68ee53a5e19732387
2 files changed
+108
-2
thirdparty/notifier/notifier.go
+24
-2
@@ -5,6 +5,9 @@ package notifier
5
6
import (
7
"sync"
8
+
9
+ process "github.com/jbenet/go-ipfs/Godeps/_workspace/src/github.com/jbenet/goprocess"
10
+ ratelimit "github.com/jbenet/go-ipfs/Godeps/_workspace/src/github.com/jbenet/goprocess/ratelimit"
11
)
12
13
// Notifiee is a generic interface. Clients implement
@@ -31,6 +34,18 @@ type Notifiee interface{}
34
type Notifier struct {
35
mu sync.RWMutex // guards notifiees
36
nots map[Notifiee]struct{}
37
+ lim *ratelimit.RateLimiter
38
+}
39
+
40
+// RateLimited returns a rate limited Notifier. only limit goroutines
41
+// will be spawned. If limit is zero, no rate limiting happens. This
42
+// is the same as `Notifier{}`.
43
+func RateLimited(limit int) Notifier {
44
+ n := Notifier{}
45
+ if limit > 0 {
46
+ n.lim = ratelimit.NewRateLimiter(process.Background(), limit)
47
+ }
48
+ return n
49
}
50
51
// Notify signs up Notifiee e for notifications. This function
@@ -107,8 +122,15 @@ func (n *Notifier) NotifyAll(notify func(Notifiee)) {
122
n.mu.Lock()
123
if n.nots != nil { // so that zero-value is ready to be used.
124
for notifiee := range n.nots {
110
- go notify(notifiee)
111
- // TODO find a good way to rate limit this without blocking notifier.
125
+
126
+ if n.lim == nil { // no rate limit
127
+ go notify(notifiee)
128
+ } else {
129
+ notifiee := notifiee // rebind for data races
130
+ n.lim.LimitedGo(func(worker process.Process) {
131
+ notify(notifiee)
132
+ })
133
+ }
134
}
135
}
136
n.mu.Unlock()
thirdparty/notifier/notifier_test.go
+84
@@ -4,6 +4,7 @@ import (
4
"fmt"
5
"sync"
6
"testing"
7
+ "time"
8
)
9
10
// test data structures
@@ -205,3 +206,86 @@ func TestThreadsafe(t *testing.T) {
206
t.Error("counts disagree")
207
}
208
}
209
+
210
+type highwatermark struct {
211
+ mu sync.Mutex
212
+ mark int
213
+ limit int
214
+ errs chan error
215
+}
216
+
217
+func (m *highwatermark) incr() {
218
+ m.mu.Lock()
219
+ m.mark++
220
+ // fmt.Println("incr", m.mark)
221
+ if m.mark > m.limit {
222
+ m.errs <- fmt.Errorf("went over limit: %d/%d", m.mark, m.limit)
223
+ }
224
+ m.mu.Unlock()
225
+}
226
+
227
+func (m *highwatermark) decr() {
228
+ m.mu.Lock()
229
+ m.mark--
230
+ // fmt.Println("decr", m.mark)
231
+ if m.mark < 0 {
232
+ m.errs <- fmt.Errorf("went under zero: %d/%d", m.mark, m.limit)
233
+ }
234
+ m.mu.Unlock()
235
+}
236
+
237
+func TestLimited(t *testing.T) {
238
+ timeout := 10 * time.Second // huge timeout.
239
+ limit := 9
240
+
241
+ hwm := highwatermark{limit: limit, errs: make(chan error, 100)}
242
+ n := RateLimited(limit) // will stop after 3 rounds
243
+ n.Notify(1)
244
+ n.Notify(2)
245
+ n.Notify(3)
246
+
247
+ entr := make(chan struct{})
248
+ exit := make(chan struct{})
249
+ done := make(chan struct{})
250
+ go func() {
251
+ for i := 0; i < 10; i++ {
252
+ // fmt.Printf("round: %d\n", i)
253
+ n.NotifyAll(func(e Notifiee) {
254
+ hwm.incr()
255
+ entr <- struct{}{}
256
+ <-exit // wait
257
+ hwm.decr()
258
+ })
259
+ }
260
+ done <- struct{}{}
261
+ }()
262
+
263
+ for i := 0; i < 30; {
264
+ select {
265
+ case <-entr:
266
+ continue // let as many enter as possible
267
+ case <-time.After(1 * time.Millisecond):
268
+ }
269
+
270
+ // let one exit
271
+ select {
272
+ case <-entr:
273
+ continue // in case of timing issues.
274
+ case exit <- struct{}{}:
275
+ case <-time.After(timeout):
276
+ t.Error("got stuck")
277
+ }
278
+ i++
279
+ }
280
+
281
+ select {
282
+ case <-done: // two parts done
283
+ case <-time.After(timeout):
284
+ t.Error("did not finish")
285
+ }
286
+
287
+ close(hwm.errs)
288
+ for err := range hwm.errs {
289
+ t.Error(err)
290
+ }
291
+}