From 739c89a6184e628ceb38ee9ad0beb6be4356d2da Mon Sep 17 00:00:00 2001 From: Thibault Le Ouay Date: Thu, 22 Oct 2026 14:06:03 +0200 Subject: [PATCH] checker: batch Tinybird events per data source (#2974) * checker: batch Tinybird events per data source Rows are grouped per data source and sent as one NDJSON request to the Events API after 2s (plus up to 1s of jitter) or 5 MB. SendEvent still returns once its row is recorded, so a Pub/Sub ack keeps meaning the result is in Tinybird. Tinybird allows 100 appends/s per workspace, whatever the row count. Under Pub/Sub every checker wakes at cronTimestamp, so one POST per check sent thousands of appends in the first seconds of each minute. The 429s nacked the message and re-ran the probe, which fed the storm. A batch is retried on 429 (honoring Retry-After), 5xx and network errors for up to 30s before its rows get the error. One request per data source is in flight at a time. Each flush logs a sampled 'tinybird flush' wide event, and pending batches are flushed on shutdown. Co-Authored-By: Claude Opus 5.5 * checker: bound queued Tinybird flushes and the shutdown flush Address review on #2974: the retry budget now starts when a batch is detached, Close aborts in-flight retries, the final flush gets an absolute shutdown deadline, response bodies are drained, request-building errors are not retried, and TINYBIRD_RETRY_BUDGET<=0 keeps the default. Co-Authored-By: Claude Opus 5.5 --------- Co-authored-by: Claude Opus 5.5 --- apps/checker/cmd/server/main.go | 59 +++- apps/checker/pkg/tinybird/batcher.go | 322 ++++++++++++++++++++++ apps/checker/pkg/tinybird/batcher_test.go | 286 +++++++++++++++++++ apps/checker/pkg/tinybird/export_test.go | 12 + 4 files changed, 678 insertions(+), 1 deletion(-) create mode 100644 apps/checker/pkg/tinybird/batcher.go create mode 100644 apps/checker/pkg/tinybird/batcher_test.go create mode 100644 apps/checker/pkg/tinybird/export_test.go diff --git a/apps/checker/cmd/server/main.go b/apps/checker/cmd/server/main.go index a0374e72f..0a1cf617f 100644 --- a/apps/checker/cmd/server/main.go +++ b/apps/checker/cmd/server/main.go @@ -238,7 +238,16 @@ func main() { defer httpClient.CloseIdleConnections() - tinybirdClient := tinybird.NewClient(httpClient, tinyBirdToken) + tinybirdOptions := tinybird.DefaultBatcherOptions() + 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) + // 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) h := &handlers.Handler{ Secret: cronSecret, @@ -294,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() @@ -301,8 +311,20 @@ func main() { log.Ctx(ctx).Error().Err(err).Msg("failed to shutdown http server") } <-consumerDone + + // 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") + } } +// 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. const consumerDrain = 50 * time.Second @@ -366,6 +388,41 @@ func logConsumerEvent(ctx context.Context, e queue.Event) { } } +func logTinybirdFlush(ctx context.Context, e tinybird.FlushEvent) { + event := map[string]any{ + "timestamp": time.Now().Format(time.RFC3339), + "datasource": e.DataSource, + "rows": e.Rows, + "bytes": e.Bytes, + "attempts": e.Attempts, + "ambiguous_attempts": e.AmbiguousAttempts, + "status_code": e.StatusCode, + "duration_ms": int(e.Duration.Milliseconds()), + "oldest_wait_ms": int(e.OldestWait.Milliseconds()), + "ratelimit_limit": e.RateLimitLimit, + "ratelimit_remaining": e.RateLimitRemaining, + } + if e.Err != nil { + event["error"] = map[string]any{ + "type": "TinybirdFlushError", + "message": e.Err.Error(), + } + } + + if e.Attempts > 1 || shouldSample(event) { + slog.LogAttrs(ctx, slog.LevelInfo, "tinybird flush", MapToAttrs(event)...) + } +} + +func envDuration(key string, fallback time.Duration) time.Duration { + value, err := time.ParseDuration(env(key, "")) + if err != nil || value < 0 { + return fallback + } + + return value +} + func envInt(key string, fallback int) int { value, err := strconv.Atoi(env(key, "")) if err != nil || value <= 0 { diff --git a/apps/checker/pkg/tinybird/batcher.go b/apps/checker/pkg/tinybird/batcher.go new file mode 100644 index 000000000..79b858dc4 --- /dev/null +++ b/apps/checker/pkg/tinybird/batcher.go @@ -0,0 +1,322 @@ +package tinybird + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "io" + "math/rand/v2" + "net/http" + "net/url" + "strconv" + "sync" + "time" +) + +var ErrBatcherClosed = errors.New("tinybird batcher closed") + +type BatcherOptions struct { + // A batch is flushed FlushInterval plus a random share of FlushJitter after + // its first row, or as soon as it reaches MaxBytes. + FlushInterval time.Duration + FlushJitter time.Duration + MaxBytes int + // 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) +} + +func DefaultBatcherOptions() BatcherOptions { + return BatcherOptions{ + FlushInterval: 2 * time.Second, + FlushJitter: time.Second, + MaxBytes: 5 << 20, + RetryBudget: 30 * time.Second, + } +} + +type FlushEvent struct { + DataSource string + Rows int + Bytes int + Attempts int + // AmbiguousAttempts counts failed attempts Tinybird may still have ingested + // (network errors and 5xx), so a later success can duplicate the rows. + AmbiguousAttempts int + StatusCode int + Duration time.Duration + OldestWait time.Duration + RateLimitLimit int + RateLimitRemaining int + Err error +} + +// Batcher implements Client by grouping rows per data source into one NDJSON +// request to the Events API. SendEvent returns once its row's batch is flushed. +type Batcher struct { + httpClient *http.Client + apiKey string + 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 + senders map[string]*sync.Mutex + flushes sync.WaitGroup +} + +type batch struct { + dataSource string + 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{}, + } +} + +func (b *Batcher) SendEvent(ctx context.Context, event any, dataSourceName string) error { + line, err := json.Marshal(event) + if err != nil { + return fmt.Errorf("unable to encode payload: %w", err) + } + + done := make(chan error, 1) + + b.mu.Lock() + if b.closed { + b.mu.Unlock() + return ErrBatcherClosed + } + current, ok := b.batches[dataSourceName] + if !ok { + current = &batch{dataSource: dataSourceName, startedAt: time.Now()} + current.timer = time.AfterFunc(b.flushDelay(), func() { b.flushOnTimer(current) }) + b.batches[dataSourceName] = current + } + current.body.Write(line) + current.body.WriteByte('\n') + current.waiters = append(current.waiters, done) + if current.body.Len() >= b.opts.MaxBytes { + b.detachLocked(current) + } + 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 + case <-ctx.Done(): + return ctx.Err() + } +} + +// 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 { + b.detachLocked(pending) + } + b.mu.Unlock() + + done := make(chan struct{}) + go func() { + b.flushes.Wait() + close(done) + }() + + select { + case <-done: + return nil + case <-ctx.Done(): + b.cancelStop() + <-done + return ctx.Err() + } +} + +func (b *Batcher) flushDelay() time.Duration { + if b.opts.FlushJitter <= 0 { + return b.opts.FlushInterval + } + + return b.opts.FlushInterval + rand.N(b.opts.FlushJitter) +} + +func (b *Batcher) flushOnTimer(target *batch) { + b.mu.Lock() + defer b.mu.Unlock() + if b.batches[target.dataSource] == target { + b.detachLocked(target) + } +} + +// detachLocked hands the batch to a flush goroutine; new rows start a new 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 { + sender = &sync.Mutex{} + b.senders[target.dataSource] = sender + } + + b.flushes.Add(1) + go func() { + defer b.flushes.Done() + // 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) + }() +} + +func (b *Batcher) flush(target *batch) { + start := time.Now() + event := FlushEvent{ + DataSource: target.dataSource, + Rows: len(target.waiters), + Bytes: target.body.Len(), + RateLimitLimit: -1, + RateLimitRemaining: -1, + } + + ctx, cancel := context.WithDeadline(b.stop, target.deadline) + defer cancel() + + event.Err = b.sendWithRetry(ctx, target, &event) + event.Duration = time.Since(start) + event.OldestWait = time.Since(target.startedAt) + + for _, waiter := range target.waiters { + waiter <- event.Err + } + + if b.opts.OnFlush != nil { + b.opts.OnFlush(context.Background(), event) + } +} + +func (b *Batcher) sendWithRetry(ctx context.Context, target *batch, event *FlushEvent) error { + 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.httpClient.Do(req) + var retryAfter time.Duration + 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 { + return nil + } + err = fmt.Errorf("unexpected status code: %d", resp.StatusCode) + } + + retryable := resp == nil || resp.StatusCode == http.StatusTooManyRequests || resp.StatusCode >= 500 + if !retryable { + return err + } + if resp == nil || resp.StatusCode >= 500 { + event.AmbiguousAttempts++ + } + + wait := retryAfter + if wait <= 0 { + wait = backoff(event.Attempts) + } + if time.Now().Add(wait).After(target.deadline) { + return err + } + + select { + case <-time.After(wait): + case <-ctx.Done(): + return err + } + } +} + +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 +} -- 2.54.0