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<