-
Notifications
You must be signed in to change notification settings - Fork 8
/
wait.go
92 lines (77 loc) · 1.86 KB
/
wait.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
package wait
import (
"context"
"time"
"github.com/pkg/errors"
)
type WaitFn func() (res interface{}, done bool, err error)
type Handler struct {
fn WaitFn
throttle time.Duration
timeout time.Duration
}
// New creates a new Wait instance
func New(f WaitFn) *Handler {
return &Handler{
fn: f,
throttle: 5 * time.Second,
timeout: 30 * time.Minute,
}
}
// SetThrottle sets the duration between func triggering
func (w *Handler) SetThrottle(d time.Duration) error {
if d == 0 {
return errors.New("Throttle can't be 0")
}
w.throttle = d
return nil
}
// SetTimeout sets the duration for wait timeout
func (w *Handler) SetTimeout(d time.Duration) *Handler {
w.timeout = d
return w
}
// Wait starts the wait until there's an error or wait is done
func (w *Handler) Wait() (res interface{}, err error) {
var done bool
ctx, cancel := context.WithTimeout(context.Background(), w.timeout)
defer cancel()
for {
res, done, err = w.fn()
if err != nil || done {
return
}
// context timeout was chosen in order to support throttle = 0
tick, cancelTick := context.WithTimeout(context.Background(), w.throttle)
defer cancelTick()
select {
case <-tick.Done():
// continue
case <-ctx.Done():
return res, errors.New("Wait() has timed out")
}
}
}
// WaitWithContext starts the wait until there's an error or wait is done
func (w Handler) WaitWithContext(ctx context.Context) (res interface{}, err error) {
var done bool
ctx, cancel := context.WithTimeout(ctx, w.timeout)
defer cancel()
ticker := time.NewTicker(w.throttle)
defer ticker.Stop()
for {
res, done, err = w.fn()
if err != nil {
return res, errors.Wrap(err, "defined wait function returned an error")
}
if done {
return res, nil
}
select {
case <-ticker.C:
// continue
case <-ctx.Done():
return res, errors.New("Wait() has timed out")
}
}
}