Files
mautrix-telegram/pkg/gotd/telegram/internal/e2etest/retry.go
T
2025-06-27 20:03:37 -07:00

62 lines
1.4 KiB
Go

package e2etest
import (
"context"
"time"
"github.com/cenkalti/backoff/v4"
"go.mau.fi/mautrix-telegram/pkg/gotd/bin"
"go.mau.fi/mautrix-telegram/pkg/gotd/tg"
"go.mau.fi/mautrix-telegram/pkg/gotd/tgerr"
)
type retryInvoker struct {
prev tg.Invoker
}
func retryResult[T any](ctx context.Context, cb func() (T, error)) (T, error) {
var zero T
return backoff.RetryWithData[T](func() (T, error) {
res, err := cb()
if err != nil {
if tgerr.IsCode(err, -500) {
return zero, err
}
if tgerr.Is(err, "CONNECTION_NOT_INITED") {
return zero, err
}
if ok, err := tgerr.FloodWait(ctx, err); ok {
return zero, err
}
return zero, backoff.Permanent(err)
}
return res, nil
}, backoff.WithContext(backoff.NewConstantBackOff(time.Millisecond*500), ctx))
}
func retry(ctx context.Context, cb func() error) error {
return backoff.Retry(func() error {
if err := cb(); err != nil {
if tgerr.IsCode(err, -500) {
return err
}
if tgerr.Is(err, "CONNECTION_NOT_INITED") {
return err
}
if ok, err := tgerr.FloodWait(ctx, err); ok {
return err
}
return backoff.Permanent(err)
}
return nil
}, backoff.WithContext(backoff.NewExponentialBackOff(), ctx))
}
func (w retryInvoker) Invoke(ctx context.Context, input bin.Encoder, output bin.Decoder) error {
return retry(ctx, func() error {
return w.prev.Invoke(ctx, input, output)
})
}