From 6f3785e971776d38aa646b7b038bbc4bd4011e25 Mon Sep 17 00:00:00 2001 From: Thibault Le Ouay Ducasse Date: Thu, 22 Oct 2026 13:44:19 +0200 Subject: [PATCH] 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 --- apps/checker/cmd/server/main.go | 53 +++- apps/checker/pkg/tinybird/batcher.go | 304 ++++++++++++++++++++++ apps/checker/pkg/tinybird/batcher_test.go | 226 ++++++++++++++++ 3 files changed, 582 insertions(+), 1 deletion(-) create mode 100644 apps/checker/pkg/tinybird/batcher.go create mode 100644 apps/checker/pkg/tinybird/batcher_test.go diff --git a/apps/checker/cmd/server/main.go b/apps/checker/cmd/server/main.go index a0374e72f..c0dd9bbfa 100644 --- a/apps/checker/cmd/server/main.go +++ b/apps/checker/cmd/server/main.go @@ -238,7 +238,13 @@ 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) + tinybirdOptions.RetryBudget = envDuration("TINYBIRD_RETRY_BUDGET", tinybirdOptions.RetryBudget) + tinybirdOptions.OnFlush = logTinybirdFlush + tinybirdClient := tinybird.NewBatcher(httpClient, tinyBirdToken, tinybirdOptions) h := &handlers.Handler{ Secret: cronSecret, @@ -301,8 +307,18 @@ 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. + flushCtx, flushCancel := context.WithTimeout(context.Background(), tinybirdFinalFlush) + 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 + // 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 +382,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..2f166dbcd --- /dev/null +++ b/apps/checker/pkg/tinybird/batcher.go @@ -0,0 +1,304 @@ +package tinybird + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "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 the retries of one batch on 429, 5xx and network errors. + 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 + + 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 + timer *time.Timer +} + +func NewBatcher(httpClient *http.Client, apiKey string, opts BatcherOptions) *Batcher { + return &Batcher{ + httpClient: httpClient, + apiKey: apiKey, + baseURL: getBaseURL(), + opts: opts, + 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() + + 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. +func (b *Batcher) Close(ctx context.Context) error { + 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(): + 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() + + 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. + 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.WithTimeout(context.Background(), b.opts.RetryBudget) + 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 { + deadline := time.Now().Add(b.opts.RetryBudget) + + for { + event.Attempts++ + resp, err := b.post(ctx, target) + var retryAfter time.Duration + if err == nil { + 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")) + 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(deadline) { + return err + } + + select { + case <-time.After(wait): + case <-ctx.Done(): + return err + } + } +} + +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<