-
-
Notifications
You must be signed in to change notification settings - Fork 134
/
wait.go
41 lines (33 loc) · 844 Bytes
/
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
package e2etest
import (
"context"
"time"
"github.com/cenkalti/backoff/v4"
"github.com/gotd/td/bin"
"github.com/gotd/td/telegram"
"github.com/gotd/td/tg"
)
type waitInvoker struct {
prev tg.Invoker
}
func retryFloodWait(ctx context.Context, cb func() error) error {
return backoff.Retry(func() error {
if err := cb(); err != nil {
if timeout, ok := telegram.AsFloodWait(err); ok {
select {
case <-ctx.Done():
return ctx.Err()
case <-time.After(timeout + 1*time.Second):
return err
}
}
return backoff.Permanent(err)
}
return nil
}, backoff.WithContext(backoff.NewExponentialBackOff(), ctx))
}
func (w waitInvoker) InvokeRaw(ctx context.Context, input bin.Encoder, output bin.Decoder) error {
return retryFloodWait(ctx, func() error {
return w.prev.InvokeRaw(ctx, input, output)
})
}