diff --git a/apps/checker/cmd/server/main.go b/apps/checker/cmd/server/main.go index c0dd9bbfa..0a1cf617f 100644 --- a/apps/checker/cmd/server/main.go +++ b/apps/checker/cmd/server/main.go @@ -242,7 +242,10 @@ func main() { tinybirdOptions.FlushInterval = envDuration("TINYBIRD_FLUSH_INTERVAL", tinybirdOptions.FlushInterval) tinybirdOptions.FlushJitter = envDuration("TINYBIRD_FLUSH_JITTER", tinybirdOptions.FlushJitter) tinybirdOptions.MaxBytes = envInt("TINYBIRD_MAX_BATCH_BYTES", tinybirdOptions.MaxBytes) - tinybirdOptions.RetryBudget = envDuration("TINYBIRD_RETRY_BUDGET", tinybirdOptions.RetryBudget) + // A zero budget would expire before the first request. + if retryBudget := envDuration("TINYBIRD_RETRY_BUDGET", 0); retryBudget > 0 { + tinybirdOptions.RetryBudget = retryBudget + } tinybirdOptions.OnFlush = logTinybirdFlush tinybirdClient := tinybird.NewBatcher(httpClient, tinyBirdToken, tinybirdOptions) @@ -300,6 +303,7 @@ func main() { } <-ctx.Done() + shutdownStart := time.Now() // The consumer drains on its own from ctx; stop the listener alongside it. shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), consumerDrain) defer shutdownCancel() @@ -308,16 +312,18 @@ func main() { } <-consumerDone - // Checks that finished during the drain may still sit in a batch. - flushCtx, flushCancel := context.WithTimeout(context.Background(), tinybirdFinalFlush) + // Checks that finished during the drain may still sit in a batch. The flush + // gets whatever the drain left of the shutdown budget. + flushCtx, flushCancel := context.WithDeadline(context.Background(), shutdownStart.Add(shutdownBudget)) defer flushCancel() if err := tinybirdClient.Close(flushCtx); err != nil { log.Error().Err(err).Msg("failed to flush tinybird batches") } } -// tinybirdFinalFlush fits in what the 60s fly kill_timeout leaves after consumerDrain. -const tinybirdFinalFlush = 8 * time.Second +// shutdownBudget keeps the whole shutdown, final Tinybird flush included, under +// the 60s fly kill_timeout. +const shutdownBudget = 58 * time.Second // consumerDrain stays under the 60s fly kill_timeout, so unfinished checks are // nacked for redelivery before the machine is killed. diff --git a/apps/checker/pkg/tinybird/batcher.go b/apps/checker/pkg/tinybird/batcher.go index 2f166dbcd..79b858dc4 100644 --- a/apps/checker/pkg/tinybird/batcher.go +++ b/apps/checker/pkg/tinybird/batcher.go @@ -6,6 +6,7 @@ import ( "encoding/json" "errors" "fmt" + "io" "math/rand/v2" "net/http" "net/url" @@ -22,7 +23,9 @@ type BatcherOptions struct { FlushInterval time.Duration FlushJitter time.Duration MaxBytes int - // RetryBudget bounds the retries of one batch on 429, 5xx and network errors. + // RetryBudget bounds one batch from the moment it is detached, including the + // wait behind an earlier flush of the same data source, so a Tinybird outage + // cannot queue batches indefinitely. RetryBudget time.Duration OnFlush func(context.Context, FlushEvent) } @@ -60,6 +63,10 @@ type Batcher struct { baseURL string opts BatcherOptions + // stop aborts in-flight flushes once Close gives up waiting. + stop context.Context + cancelStop context.CancelFunc + mu sync.Mutex closed bool batches map[string]*batch @@ -72,15 +79,20 @@ type batch struct { body bytes.Buffer waiters []chan error startedAt time.Time + deadline time.Time timer *time.Timer } func NewBatcher(httpClient *http.Client, apiKey string, opts BatcherOptions) *Batcher { + stop, cancelStop := context.WithCancel(context.Background()) + return &Batcher{ httpClient: httpClient, apiKey: apiKey, baseURL: getBaseURL(), opts: opts, + stop: stop, + cancelStop: cancelStop, batches: map[string]*batch{}, senders: map[string]*sync.Mutex{}, } @@ -113,6 +125,8 @@ func (b *Batcher) SendEvent(ctx context.Context, event any, dataSourceName strin } b.mu.Unlock() + // A cancelled caller's row stays in the batch and is still sent, so the + // redelivered check can record it twice: delivery is at least once. select { case err := <-done: return err @@ -121,8 +135,11 @@ func (b *Batcher) SendEvent(ctx context.Context, event any, dataSourceName strin } } -// Close flushes every pending batch and waits for in-flight flushes, bounded by ctx. +// Close flushes every pending batch and waits for in-flight flushes, bounded by +// ctx. When ctx ends first, in-flight flushes are aborted and their rows fail. func (b *Batcher) Close(ctx context.Context) error { + defer b.cancelStop() + b.mu.Lock() b.closed = true for _, pending := range b.batches { @@ -140,6 +157,8 @@ func (b *Batcher) Close(ctx context.Context) error { case <-done: return nil case <-ctx.Done(): + b.cancelStop() + <-done return ctx.Err() } } @@ -164,6 +183,7 @@ func (b *Batcher) flushOnTimer(target *batch) { func (b *Batcher) detachLocked(target *batch) { delete(b.batches, target.dataSource) target.timer.Stop() + target.deadline = time.Now().Add(b.opts.RetryBudget) sender, ok := b.senders[target.dataSource] if !ok { @@ -174,7 +194,8 @@ func (b *Batcher) detachLocked(target *batch) { b.flushes.Add(1) go func() { defer b.flushes.Done() - // One request per data source at a time, so retries don't pile up. + // One request per data source at a time, so retries don't pile up; the + // batch deadline bounds the wait for the lock. sender.Lock() defer sender.Unlock() b.flush(target) @@ -191,7 +212,7 @@ func (b *Batcher) flush(target *batch) { RateLimitRemaining: -1, } - ctx, cancel := context.WithTimeout(context.Background(), b.opts.RetryBudget) + ctx, cancel := context.WithDeadline(b.stop, target.deadline) defer cancel() event.Err = b.sendWithRetry(ctx, target, &event) @@ -208,17 +229,37 @@ func (b *Batcher) flush(target *batch) { } func (b *Batcher) sendWithRetry(ctx context.Context, target *batch, event *FlushEvent) error { - deadline := time.Now().Add(b.opts.RetryBudget) + if err := ctx.Err(); err != nil { + return fmt.Errorf("flush not attempted: %w", err) + } + + requestURL, err := url.Parse(b.baseURL) + if err != nil { + return fmt.Errorf("unable to parse url: %w", err) + } + q := requestURL.Query() + q.Add("name", target.dataSource) + requestURL.RawQuery = q.Encode() for { + req, err := http.NewRequestWithContext(ctx, http.MethodPost, requestURL.String(), bytes.NewReader(target.body.Bytes())) + if err != nil { + return fmt.Errorf("unable to create request: %w", err) + } + req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", b.apiKey)) + event.Attempts++ - resp, err := b.post(ctx, target) + resp, err := b.httpClient.Do(req) var retryAfter time.Duration - if err == nil { + if err != nil { + err = fmt.Errorf("unable to send request: %w", err) + } else { event.StatusCode = resp.StatusCode event.RateLimitLimit = headerInt(resp.Header, "X-RateLimit-Limit") event.RateLimitRemaining = headerInt(resp.Header, "X-RateLimit-Remaining") retryAfter = parseRetryAfter(resp.Header.Get("Retry-After")) + // Drain so the connection is reused by the next attempt. + _, _ = io.Copy(io.Discard, io.LimitReader(resp.Body, 64<<10)) resp.Body.Close() if resp.StatusCode == http.StatusAccepted { @@ -239,7 +280,7 @@ func (b *Batcher) sendWithRetry(ctx context.Context, target *batch, event *Flush if wait <= 0 { wait = backoff(event.Attempts) } - if time.Now().Add(wait).After(deadline) { + if time.Now().Add(wait).After(target.deadline) { return err } @@ -251,29 +292,6 @@ func (b *Batcher) sendWithRetry(ctx context.Context, target *batch, event *Flush } } -func (b *Batcher) post(ctx context.Context, target *batch) (*http.Response, error) { - requestURL, err := url.Parse(b.baseURL) - if err != nil { - return nil, fmt.Errorf("unable to parse url: %w", err) - } - q := requestURL.Query() - q.Add("name", target.dataSource) - requestURL.RawQuery = q.Encode() - - req, err := http.NewRequestWithContext(ctx, http.MethodPost, requestURL.String(), bytes.NewReader(target.body.Bytes())) - if err != nil { - return nil, fmt.Errorf("unable to create request: %w", err) - } - req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", b.apiKey)) - - resp, err := b.httpClient.Do(req) - if err != nil { - return nil, fmt.Errorf("unable to send request: %w", err) - } - - return resp, nil -} - func backoff(attempt int) time.Duration { base := min(250*time.Millisecond< 0 }, time.Second, time.Millisecond) + + // Queued behind the first batch, which keeps its sender busy past this batch's budget. + start := time.Now() + err := b.SendEvent(t.Context(), "event", "test") + require.Len(t, api.recorded(), 1) + require.Less(t, time.Since(start), 3*opts.RetryBudget) + require.ErrorIs(t, err, context.DeadlineExceeded) + require.ErrorContains(t, err, "flush not attempted") + require.ErrorContains(t, <-first, "unexpected status code: 503") + }) + + t.Run("it should abort in-flight retries once Close gives up", func(t *testing.T) { + t.Parallel() + api := &fakeEventsAPI{respond: func(int) *http.Response { return status(http.StatusServiceUnavailable, nil) }} + opts := options() + opts.FlushInterval = time.Hour + opts.MaxBytes = 1 + opts.RetryBudget = time.Minute + b := tinybird.NewBatcher(api.client(), "apiKey", opts) + + sent := make(chan error, 1) + go func() { sent <- b.SendEvent(t.Context(), "event", "test") }() + require.Eventually(t, func() bool { return len(api.recorded()) > 0 }, time.Second, time.Millisecond) + + ctx, cancel := context.WithTimeout(t.Context(), 50*time.Millisecond) + defer cancel() + require.ErrorIs(t, b.Close(ctx), context.DeadlineExceeded) + select { + case err := <-sent: + require.ErrorContains(t, err, "unexpected status code: 503") + default: + t.Fatal("the row should fail as soon as Close returns") + } + }) } diff --git a/apps/checker/pkg/tinybird/export_test.go b/apps/checker/pkg/tinybird/export_test.go new file mode 100644 index 000000000..c1588afa1 --- /dev/null +++ b/apps/checker/pkg/tinybird/export_test.go @@ -0,0 +1,12 @@ +package tinybird + +// PendingRows reports the rows buffered for a data source and not yet detached. +func (b *Batcher) PendingRows(dataSource string) int { + b.mu.Lock() + defer b.mu.Unlock() + if pending, ok := b.batches[dataSource]; ok { + return len(pending.waiters) + } + + return 0 +}