Compare commits
2
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
bcdea85a2d | ||
|
|
cbf24e63e1 |
@@ -178,13 +178,14 @@ var all = map[string]Preset{
|
|||||||
//
|
//
|
||||||
// The Provider sends s.Temperature with every request whose Temperature is
|
// The Provider sends s.Temperature with every request whose Temperature is
|
||||||
// nil; a request that sets one keeps it. Neither is sent to a model the
|
// nil; a request that sets one keeps it. Neither is sent to a model the
|
||||||
// preset's OmitTemperature matches.
|
// preset's OmitTemperature matches. It is wrapped in llm.NewRetrying, so
|
||||||
|
// Generate retries transient vendor failures and has the default timeout.
|
||||||
func FromSettings(s llm.Settings, key string) (llm.Provider, error) {
|
func FromSettings(s llm.Settings, key string) (llm.Provider, error) {
|
||||||
cp, err := fromSettings(s, key)
|
cp, err := fromSettings(s, key)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return nil, redact(err, key)
|
return nil, redact(err, key)
|
||||||
}
|
}
|
||||||
return cp, nil
|
return llm.NewRetrying(cp), nil
|
||||||
}
|
}
|
||||||
|
|
||||||
func fromSettings(s llm.Settings, key string) (*configured, error) {
|
func fromSettings(s llm.Settings, key string) (*configured, error) {
|
||||||
|
|||||||
@@ -22,10 +22,8 @@ type captured struct {
|
|||||||
body map[string]json.RawMessage
|
body map[string]json.RawMessage
|
||||||
}
|
}
|
||||||
|
|
||||||
var errStop = errors.New("recorded")
|
// recordingTransport records each request and answers it 400, so no
|
||||||
|
// request leaves the process and the Provider never retries it.
|
||||||
// recordingTransport records each request and fails it, so no request
|
|
||||||
// leaves the process.
|
|
||||||
type recordingTransport struct{ reqs []captured }
|
type recordingTransport struct{ reqs []captured }
|
||||||
|
|
||||||
func (rt *recordingTransport) RoundTrip(r *http.Request) (*http.Response, error) {
|
func (rt *recordingTransport) RoundTrip(r *http.Request) (*http.Response, error) {
|
||||||
@@ -40,7 +38,12 @@ func (rt *recordingTransport) RoundTrip(r *http.Request) (*http.Response, error)
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
rt.reqs = append(rt.reqs, c)
|
rt.reqs = append(rt.reqs, c)
|
||||||
return nil, errStop
|
return &http.Response{
|
||||||
|
StatusCode: http.StatusBadRequest,
|
||||||
|
Header: http.Header{"Content-Type": {"application/json"}},
|
||||||
|
Body: io.NopCloser(strings.NewReader(`{}`)),
|
||||||
|
Request: r,
|
||||||
|
}, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
// generate builds the provider s describes and returns the request its
|
// generate builds the provider s describes and returns the request its
|
||||||
@@ -327,7 +330,7 @@ func TestNegativeTemperatureRejectedForRejectingModels(t *testing.T) {
|
|||||||
}
|
}
|
||||||
req := plainRequest()
|
req := plainRequest()
|
||||||
req.Temperature = ptr(-0.5)
|
req.Temperature = ptr(-0.5)
|
||||||
if _, err := p.Generate(context.Background(), req); err == nil || errors.Is(err, errStop) || !strings.Contains(err.Error(), "temperature") {
|
if _, err := p.Generate(context.Background(), req); err == nil || errors.Is(err, llm.ErrProvider) || !strings.Contains(err.Error(), "temperature") {
|
||||||
t.Errorf("%s: Generate = %v, want a temperature validation error", s.Model, err)
|
t.Errorf("%s: Generate = %v, want a temperature validation error", s.Model, err)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -465,7 +468,7 @@ func TestGenerateValidationErrorsNeverLeakKey(t *testing.T) {
|
|||||||
req := plainRequest()
|
req := plainRequest()
|
||||||
req.Temperature = ptr(-0.5)
|
req.Temperature = ptr(-0.5)
|
||||||
_, err = p.Generate(context.Background(), req)
|
_, err = p.Generate(context.Background(), req)
|
||||||
if err == nil || errors.Is(err, errStop) {
|
if err == nil || errors.Is(err, llm.ErrProvider) {
|
||||||
t.Fatalf("%s: Generate = %v, want a validation error", key, err)
|
t.Fatalf("%s: Generate = %v, want a validation error", key, err)
|
||||||
}
|
}
|
||||||
for e := error(err); e != nil; e = errors.Unwrap(e) {
|
for e := error(err); e != nil; e = errors.Unwrap(e) {
|
||||||
@@ -518,3 +521,86 @@ func TestConfigErrorsNeverLeakKey(t *testing.T) {
|
|||||||
})
|
})
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// scriptedTransport answers each request with the next scripted status and
|
||||||
|
// body, echoing the key in every failure body, and counts the requests.
|
||||||
|
type scriptedTransport struct {
|
||||||
|
statuses []int
|
||||||
|
key string
|
||||||
|
sent int
|
||||||
|
}
|
||||||
|
|
||||||
|
func (st *scriptedTransport) RoundTrip(r *http.Request) (*http.Response, error) {
|
||||||
|
if st.sent == len(st.statuses) {
|
||||||
|
return nil, fmt.Errorf("unscripted request %d", st.sent+1)
|
||||||
|
}
|
||||||
|
status := st.statuses[st.sent]
|
||||||
|
st.sent++
|
||||||
|
header := http.Header{"Content-Type": {"application/json"}}
|
||||||
|
body := `{"error":{"message":"bad key ` + st.key + `"}}`
|
||||||
|
switch status {
|
||||||
|
case http.StatusOK:
|
||||||
|
body = `{"model":"m","choices":[{"message":{"content":"tacos"},"finish_reason":"stop"}]}`
|
||||||
|
case http.StatusTooManyRequests:
|
||||||
|
header.Set("Retry-After", "1")
|
||||||
|
}
|
||||||
|
return &http.Response{
|
||||||
|
StatusCode: status,
|
||||||
|
Header: header,
|
||||||
|
Body: io.NopCloser(strings.NewReader(body)),
|
||||||
|
Request: r,
|
||||||
|
}, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// useTransport makes the adapters send through rt for the rest of the test.
|
||||||
|
func useTransport(t *testing.T, rt http.RoundTripper) {
|
||||||
|
t.Helper()
|
||||||
|
orig := http.DefaultTransport
|
||||||
|
http.DefaultTransport = rt
|
||||||
|
t.Cleanup(func() { http.DefaultTransport = orig })
|
||||||
|
}
|
||||||
|
|
||||||
|
// The Provider FromSettings returns retries a rate-limited Generate.
|
||||||
|
func TestFromSettingsRetriesRateLimited(t *testing.T) {
|
||||||
|
st := &scriptedTransport{statuses: []int{http.StatusTooManyRequests, http.StatusOK}, key: testKey}
|
||||||
|
useTransport(t, st)
|
||||||
|
|
||||||
|
p, err := FromSettings(llm.Settings{Provider: "openai", Model: "m"}, testKey)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("FromSettings: %v", err)
|
||||||
|
}
|
||||||
|
resp, err := p.Generate(context.Background(), plainRequest())
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("Generate: %v", err)
|
||||||
|
}
|
||||||
|
if resp.Text != "tacos" {
|
||||||
|
t.Errorf("Text = %q, want %q", resp.Text, "tacos")
|
||||||
|
}
|
||||||
|
if st.sent != 2 {
|
||||||
|
t.Errorf("sent %d requests, want 2", st.sent)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// When every retry fails, the error FromSettings's Provider returns keeps
|
||||||
|
// its sentinel and never holds the key the vendor echoed.
|
||||||
|
func TestFromSettingsRetryErrorNeverLeaksKey(t *testing.T) {
|
||||||
|
st := &scriptedTransport{statuses: []int{http.StatusTooManyRequests, http.StatusTooManyRequests, http.StatusServiceUnavailable}, key: testKey}
|
||||||
|
useTransport(t, st)
|
||||||
|
|
||||||
|
p, err := FromSettings(llm.Settings{Provider: "openai", Model: "m"}, testKey)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("FromSettings: %v", err)
|
||||||
|
}
|
||||||
|
_, err = p.Generate(context.Background(), plainRequest())
|
||||||
|
if !errors.Is(err, llm.ErrProvider) {
|
||||||
|
t.Fatalf("Generate = %v, want %v", err, llm.ErrProvider)
|
||||||
|
}
|
||||||
|
if st.sent != llm.MaxRetries+1 {
|
||||||
|
t.Errorf("sent %d requests, want %d", st.sent, llm.MaxRetries+1)
|
||||||
|
}
|
||||||
|
for e := error(err); e != nil; e = errors.Unwrap(e) {
|
||||||
|
if s := e.Error() + fmt.Sprintf("%#v", e); strings.Contains(s, testKey) {
|
||||||
|
t.Fatalf("error leaks the key: %s", s)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
@@ -0,0 +1,107 @@
|
|||||||
|
package llm
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"errors"
|
||||||
|
"math/rand/v2"
|
||||||
|
"time"
|
||||||
|
)
|
||||||
|
|
||||||
|
const (
|
||||||
|
// DefaultTimeout bounds each Generate attempt whose Request sets no
|
||||||
|
// Timeout.
|
||||||
|
DefaultTimeout = 60 * time.Second
|
||||||
|
// MaxRetries is how many times a failed Generate is retried, so a call
|
||||||
|
// makes at most MaxRetries+1 attempts.
|
||||||
|
MaxRetries = 2
|
||||||
|
// MaxRetryAfter caps the wait a rate-limited vendor may ask for.
|
||||||
|
MaxRetryAfter = 20 * time.Second
|
||||||
|
// backoffBase is the first retry's backoff before jitter; each later
|
||||||
|
// retry doubles it.
|
||||||
|
backoffBase = time.Second
|
||||||
|
)
|
||||||
|
|
||||||
|
// Retrying is a Provider that retries transient Generate failures of the
|
||||||
|
// Provider it wraps. It retries ErrRateLimited, ErrTimeout and a 5xx
|
||||||
|
// ErrProvider at most MaxRetries times, and nothing else. A rate-limited
|
||||||
|
// retry waits for the vendor's RetryAfter, capped at MaxRetryAfter; every
|
||||||
|
// other retry, and a rate limit that names no wait, waits a jittered
|
||||||
|
// exponential backoff. A retry whose wait would not end before the
|
||||||
|
// caller's context deadline is not made: the last error is returned at
|
||||||
|
// once. ListModels and Verify pass straight through.
|
||||||
|
type Retrying struct {
|
||||||
|
Provider
|
||||||
|
// now is the clock; tests replace it.
|
||||||
|
now func() time.Time
|
||||||
|
// sleep waits d or until ctx is done; tests replace it.
|
||||||
|
sleep func(ctx context.Context, d time.Duration) error
|
||||||
|
// jitter returns a random duration in [0, n) for n > 0; tests replace
|
||||||
|
// it.
|
||||||
|
jitter func(n time.Duration) time.Duration
|
||||||
|
}
|
||||||
|
|
||||||
|
// NewRetrying returns inner wrapped with retries and the default Generate
|
||||||
|
// timeout.
|
||||||
|
func NewRetrying(inner Provider) *Retrying {
|
||||||
|
return &Retrying{
|
||||||
|
Provider: inner,
|
||||||
|
now: time.Now,
|
||||||
|
sleep: sleep,
|
||||||
|
jitter: rand.N[time.Duration],
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Generate sends req to the wrapped Provider, retrying transient failures.
|
||||||
|
// A req.Timeout of 0 becomes DefaultTimeout; either bounds each attempt,
|
||||||
|
// not the whole call, which ctx bounds.
|
||||||
|
func (r *Retrying) Generate(ctx context.Context, req Request) (Response, error) {
|
||||||
|
if req.Timeout == 0 {
|
||||||
|
req.Timeout = DefaultTimeout
|
||||||
|
}
|
||||||
|
for retry := 0; ; retry++ {
|
||||||
|
resp, err := r.Provider.Generate(ctx, req)
|
||||||
|
if err == nil || retry == MaxRetries || ctx.Err() != nil {
|
||||||
|
return resp, err
|
||||||
|
}
|
||||||
|
wait, ok := r.wait(err, retry)
|
||||||
|
if !ok {
|
||||||
|
return resp, err
|
||||||
|
}
|
||||||
|
if deadline, has := ctx.Deadline(); has && !r.now().Add(wait).Before(deadline) {
|
||||||
|
return resp, err
|
||||||
|
}
|
||||||
|
if serr := r.sleep(ctx, wait); serr != nil {
|
||||||
|
return resp, err
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// wait returns how long to wait before retrying after err, the retry-th
|
||||||
|
// retry counting from 0, and false when err is not retried.
|
||||||
|
func (r *Retrying) wait(err error, retry int) (time.Duration, bool) {
|
||||||
|
var rl *RateLimitedError
|
||||||
|
if errors.As(err, &rl) && rl.RetryAfter > 0 {
|
||||||
|
return min(rl.RetryAfter, MaxRetryAfter), true
|
||||||
|
}
|
||||||
|
var pe *ProviderError
|
||||||
|
switch {
|
||||||
|
case errors.Is(err, ErrRateLimited), errors.Is(err, ErrTimeout):
|
||||||
|
case errors.As(err, &pe) && pe.Status() >= 500 && pe.Status() <= 599:
|
||||||
|
default:
|
||||||
|
return 0, false
|
||||||
|
}
|
||||||
|
half := (backoffBase << retry) / 2
|
||||||
|
return half + r.jitter(half), true
|
||||||
|
}
|
||||||
|
|
||||||
|
// sleep waits d or until ctx is done, returning ctx's error in that case.
|
||||||
|
func sleep(ctx context.Context, d time.Duration) error {
|
||||||
|
t := time.NewTimer(d)
|
||||||
|
defer t.Stop()
|
||||||
|
select {
|
||||||
|
case <-t.C:
|
||||||
|
return nil
|
||||||
|
case <-ctx.Done():
|
||||||
|
return ctx.Err()
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,271 @@
|
|||||||
|
package llm
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"errors"
|
||||||
|
"fmt"
|
||||||
|
"slices"
|
||||||
|
"testing"
|
||||||
|
"time"
|
||||||
|
)
|
||||||
|
|
||||||
|
// fakeClock is a clock whose sleeps return at once, advance it and are
|
||||||
|
// recorded.
|
||||||
|
type fakeClock struct {
|
||||||
|
t time.Time
|
||||||
|
sleeps []time.Duration
|
||||||
|
}
|
||||||
|
|
||||||
|
func (c *fakeClock) now() time.Time { return c.t }
|
||||||
|
|
||||||
|
func (c *fakeClock) sleep(ctx context.Context, d time.Duration) error {
|
||||||
|
if err := ctx.Err(); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
c.sleeps = append(c.sleeps, d)
|
||||||
|
c.t = c.t.Add(d)
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// newTestRetrying returns a Retrying over a fresh Fake with a fake clock
|
||||||
|
// at the real current time, so real context deadlines compare sensibly,
|
||||||
|
// and jitter that always returns its maximum.
|
||||||
|
func newTestRetrying() (*Retrying, *Fake, *fakeClock) {
|
||||||
|
f := &Fake{}
|
||||||
|
c := &fakeClock{t: time.Now()}
|
||||||
|
r := NewRetrying(f)
|
||||||
|
r.now = c.now
|
||||||
|
r.sleep = c.sleep
|
||||||
|
r.jitter = func(n time.Duration) time.Duration { return n - 1 }
|
||||||
|
return r, f, c
|
||||||
|
}
|
||||||
|
|
||||||
|
var retryReq = Request{Messages: []Message{{Role: RoleUser, Content: "plan dinners"}}}
|
||||||
|
|
||||||
|
func ok(text string) GenerateResult { return GenerateResult{Response: Response{Text: text}} }
|
||||||
|
|
||||||
|
func fail(err error) GenerateResult { return GenerateResult{Err: err} }
|
||||||
|
|
||||||
|
func assertConsumed(t *testing.T, f *Fake, attempts int) {
|
||||||
|
t.Helper()
|
||||||
|
if g, _, _ := f.Remaining(); g != 0 {
|
||||||
|
t.Errorf("%d scripted Generate results unused", g)
|
||||||
|
}
|
||||||
|
if n := len(f.Requests()); n != attempts {
|
||||||
|
t.Errorf("made %d attempts, want %d", n, attempts)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestRetryingRateLimitedThenSuccess(t *testing.T) {
|
||||||
|
r, f, c := newTestRetrying()
|
||||||
|
f.QueueGenerate(fail(NewRateLimitedError(3*time.Second, "")), ok("dinner"))
|
||||||
|
|
||||||
|
resp, err := r.Generate(context.Background(), retryReq)
|
||||||
|
if err != nil || resp.Text != "dinner" {
|
||||||
|
t.Fatalf("Generate = %+v, %v", resp, err)
|
||||||
|
}
|
||||||
|
if want := []time.Duration{3 * time.Second}; !slices.Equal(c.sleeps, want) {
|
||||||
|
t.Errorf("sleeps = %v, want %v", c.sleeps, want)
|
||||||
|
}
|
||||||
|
assertConsumed(t, f, 2)
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestRetryingRateLimitedWaitIsCapped(t *testing.T) {
|
||||||
|
r, f, c := newTestRetrying()
|
||||||
|
f.QueueGenerate(fail(NewRateLimitedError(time.Hour, "")), ok("dinner"))
|
||||||
|
|
||||||
|
if _, err := r.Generate(context.Background(), retryReq); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if want := []time.Duration{MaxRetryAfter}; !slices.Equal(c.sleeps, want) {
|
||||||
|
t.Errorf("sleeps = %v, want %v", c.sleeps, want)
|
||||||
|
}
|
||||||
|
assertConsumed(t, f, 2)
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestRetryingRateLimitedWithoutRetryAfterBacksOff(t *testing.T) {
|
||||||
|
r, f, c := newTestRetrying()
|
||||||
|
f.QueueGenerate(fail(NewRateLimitedError(0, "")), ok("dinner"))
|
||||||
|
|
||||||
|
if _, err := r.Generate(context.Background(), retryReq); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if want := []time.Duration{time.Second - 1}; !slices.Equal(c.sleeps, want) {
|
||||||
|
t.Errorf("sleeps = %v, want %v", c.sleeps, want)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestRetrying5xxTwiceThenSuccess(t *testing.T) {
|
||||||
|
r, f, c := newTestRetrying()
|
||||||
|
f.QueueGenerate(
|
||||||
|
fail(NewProviderError(503, []byte("busy"), "")),
|
||||||
|
// Adapters wrap the error they return.
|
||||||
|
fail(fmt.Errorf("openai: generate: %w", NewProviderError(503, []byte("busy"), ""))),
|
||||||
|
ok("dinner"),
|
||||||
|
)
|
||||||
|
|
||||||
|
resp, err := r.Generate(context.Background(), retryReq)
|
||||||
|
if err != nil || resp.Text != "dinner" {
|
||||||
|
t.Fatalf("Generate = %+v, %v", resp, err)
|
||||||
|
}
|
||||||
|
// Max jitter: [0.5 s, 1 s) then [1 s, 2 s).
|
||||||
|
if want := []time.Duration{time.Second - 1, 2*time.Second - 1}; !slices.Equal(c.sleeps, want) {
|
||||||
|
t.Errorf("sleeps = %v, want %v", c.sleeps, want)
|
||||||
|
}
|
||||||
|
assertConsumed(t, f, 3)
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestRetryingBackoffJitterLowerBound(t *testing.T) {
|
||||||
|
r, f, c := newTestRetrying()
|
||||||
|
r.jitter = func(time.Duration) time.Duration { return 0 }
|
||||||
|
f.QueueGenerate(fail(ErrTimeout), fail(ErrTimeout), ok("dinner"))
|
||||||
|
|
||||||
|
if _, err := r.Generate(context.Background(), retryReq); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if want := []time.Duration{500 * time.Millisecond, time.Second}; !slices.Equal(c.sleeps, want) {
|
||||||
|
t.Errorf("sleeps = %v, want %v", c.sleeps, want)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestRetrying5xxThreeTimesFails(t *testing.T) {
|
||||||
|
r, f, c := newTestRetrying()
|
||||||
|
last := NewProviderError(503, []byte("third"), "")
|
||||||
|
f.QueueGenerate(
|
||||||
|
fail(NewProviderError(503, []byte("first"), "")),
|
||||||
|
fail(NewProviderError(503, []byte("second"), "")),
|
||||||
|
fail(last),
|
||||||
|
)
|
||||||
|
|
||||||
|
_, err := r.Generate(context.Background(), retryReq)
|
||||||
|
if err != last {
|
||||||
|
t.Fatalf("err = %v, want the last attempt's error", err)
|
||||||
|
}
|
||||||
|
if len(c.sleeps) != MaxRetries {
|
||||||
|
t.Errorf("slept %d times, want %d", len(c.sleeps), MaxRetries)
|
||||||
|
}
|
||||||
|
assertConsumed(t, f, 3)
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestRetryingDoesNotRetryPermanentErrors(t *testing.T) {
|
||||||
|
for name, err := range map[string]error{
|
||||||
|
"401 auth": ErrAuth,
|
||||||
|
"400": NewProviderError(400, []byte("bad request"), ""),
|
||||||
|
"404": NewProviderError(404, []byte("no such model"), ""),
|
||||||
|
"bad output": ErrBadOutput,
|
||||||
|
"other": errors.New("llm: request has no messages"),
|
||||||
|
} {
|
||||||
|
t.Run(name, func(t *testing.T) {
|
||||||
|
r, f, c := newTestRetrying()
|
||||||
|
f.QueueGenerate(fail(err))
|
||||||
|
|
||||||
|
if _, got := r.Generate(context.Background(), retryReq); got != err {
|
||||||
|
t.Fatalf("err = %v, want %v", got, err)
|
||||||
|
}
|
||||||
|
if len(c.sleeps) != 0 {
|
||||||
|
t.Errorf("slept %v", c.sleeps)
|
||||||
|
}
|
||||||
|
assertConsumed(t, f, 1)
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestRetryingDeadlineShorterThanRetryAfter(t *testing.T) {
|
||||||
|
r, f, c := newTestRetrying()
|
||||||
|
rl := NewRateLimitedError(10*time.Second, "")
|
||||||
|
f.QueueGenerate(fail(rl))
|
||||||
|
ctx, cancel := context.WithDeadline(context.Background(), c.t.Add(5*time.Second))
|
||||||
|
defer cancel()
|
||||||
|
|
||||||
|
if _, err := r.Generate(ctx, retryReq); err != rl {
|
||||||
|
t.Fatalf("err = %v, want the rate-limited error", err)
|
||||||
|
}
|
||||||
|
if len(c.sleeps) != 0 {
|
||||||
|
t.Errorf("slept %v", c.sleeps)
|
||||||
|
}
|
||||||
|
assertConsumed(t, f, 1)
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestRetryingDeadlineShorterThanBackoff(t *testing.T) {
|
||||||
|
r, f, c := newTestRetrying()
|
||||||
|
f.QueueGenerate(fail(ErrTimeout))
|
||||||
|
// The first backoff with maximum jitter is just under 1 s.
|
||||||
|
ctx, cancel := context.WithDeadline(context.Background(), c.t.Add(time.Second-1))
|
||||||
|
defer cancel()
|
||||||
|
|
||||||
|
if _, err := r.Generate(ctx, retryReq); !errors.Is(err, ErrTimeout) {
|
||||||
|
t.Fatalf("err = %v, want ErrTimeout", err)
|
||||||
|
}
|
||||||
|
if len(c.sleeps) != 0 {
|
||||||
|
t.Errorf("slept %v", c.sleeps)
|
||||||
|
}
|
||||||
|
assertConsumed(t, f, 1)
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestRetryingRetriesWithinDeadline(t *testing.T) {
|
||||||
|
r, f, c := newTestRetrying()
|
||||||
|
f.QueueGenerate(fail(NewRateLimitedError(3*time.Second, "")), ok("dinner"))
|
||||||
|
ctx, cancel := context.WithDeadline(context.Background(), c.t.Add(time.Minute))
|
||||||
|
defer cancel()
|
||||||
|
|
||||||
|
if _, err := r.Generate(ctx, retryReq); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
assertConsumed(t, f, 2)
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestRetryingStopsWhenContextDone(t *testing.T) {
|
||||||
|
r, f, _ := newTestRetrying()
|
||||||
|
ctx, cancel := context.WithCancel(context.Background())
|
||||||
|
cancel()
|
||||||
|
f.QueueGenerate(fail(ErrTimeout))
|
||||||
|
|
||||||
|
if _, err := r.Generate(ctx, retryReq); !errors.Is(err, ErrTimeout) {
|
||||||
|
t.Fatalf("err = %v, want ErrTimeout", err)
|
||||||
|
}
|
||||||
|
assertConsumed(t, f, 1)
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestRetryingDefaultTimeout(t *testing.T) {
|
||||||
|
r, f, _ := newTestRetrying()
|
||||||
|
f.QueueGenerate(ok("a"), ok("b"))
|
||||||
|
custom := retryReq
|
||||||
|
custom.Timeout = 5 * time.Second
|
||||||
|
|
||||||
|
for _, req := range []Request{retryReq, custom} {
|
||||||
|
if _, err := r.Generate(context.Background(), req); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
got := f.Requests()
|
||||||
|
if got[0].Timeout != DefaultTimeout || got[1].Timeout != 5*time.Second {
|
||||||
|
t.Fatalf("timeouts = %v, %v; want %v, 5s", got[0].Timeout, got[1].Timeout, DefaultTimeout)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestRetryingPassesThroughListModelsAndVerify(t *testing.T) {
|
||||||
|
r, f, _ := newTestRetrying()
|
||||||
|
f.QueueListModels(ListModelsResult{Err: ErrTimeout})
|
||||||
|
f.QueueVerify(ErrTimeout)
|
||||||
|
|
||||||
|
if _, err := r.ListModels(context.Background()); !errors.Is(err, ErrTimeout) {
|
||||||
|
t.Fatalf("ListModels err = %v", err)
|
||||||
|
}
|
||||||
|
if err := r.Verify(context.Background()); !errors.Is(err, ErrTimeout) {
|
||||||
|
t.Fatalf("Verify err = %v", err)
|
||||||
|
}
|
||||||
|
if _, l, v := f.Remaining(); l != 0 || v != 0 {
|
||||||
|
t.Fatalf("Remaining = %d, %d", l, v)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestSleepStopsWhenContextDone(t *testing.T) {
|
||||||
|
ctx, cancel := context.WithCancel(context.Background())
|
||||||
|
cancel()
|
||||||
|
if err := sleep(ctx, time.Hour); !errors.Is(err, context.Canceled) {
|
||||||
|
t.Fatalf("sleep err = %v", err)
|
||||||
|
}
|
||||||
|
if err := sleep(context.Background(), time.Millisecond); err != nil {
|
||||||
|
t.Fatalf("sleep err = %v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
Reference in New Issue
Block a user