From c3cbe8aebdadfc088f670cf071d54cffa3aba354 Mon Sep 17 00:00:00 2001 From: Thibault Le Ouay Date: Tue, 27 Oct 2026 10:21:10 +0100 Subject: [PATCH] private location: improvment (#3005) * private location: improvment * private location: address review feedback - tinybird batcher: wait for the per-source sender with the batch deadline and Close, instead of an uninterruptible mutex (checker and private location stay identical) - server: give the shutdown flush its own 5s bound, clamp TINYBIRD_RETRY_BUDGET to fit the WriteTimeout, kill_timeout 15s on fly - tests: no require inside goroutines, wider flush window in the ingest batch test, one shared fake Hrana server Co-Authored-By: Claude Opus 5.5 --------- Co-authored-by: Claude Opus 5.5 --- apps/checker/pkg/tinybird/batcher.go | 26 +- apps/checker/pkg/tinybird/batcher_test.go | 15 +- apps/private-location/fly.toml | 2 + apps/private-location/go.mod | 4 - apps/private-location/go.sum | 16 - .../internal/database/database_test.go | 68 +--- .../internal/database/hrana_test.go | 142 ++++++++ .../internal/database/pool_test.go | 148 ++++++++ .../internal/server/ingest_batch_test.go | 196 +++++++++++ .../internal/server/routes.go | 12 +- .../internal/server/server.go | 95 ++++- .../internal/server/tinybird_options_test.go | 58 ++++ .../internal/tinybird/batcher.go | 326 ++++++++++++++++++ .../internal/tinybird/batcher_test.go | 293 ++++++++++++++++ .../internal/tinybird/client_test.go | 2 +- .../internal/tinybird/export_test.go | 12 + 16 files changed, 1302 insertions(+), 113 deletions(-) create mode 100644 apps/private-location/internal/database/hrana_test.go create mode 100644 apps/private-location/internal/database/pool_test.go create mode 100644 apps/private-location/internal/server/ingest_batch_test.go create mode 100644 apps/private-location/internal/server/tinybird_options_test.go create mode 100644 apps/private-location/internal/tinybird/batcher.go create mode 100644 apps/private-location/internal/tinybird/batcher_test.go create mode 100644 apps/private-location/internal/tinybird/export_test.go diff --git a/apps/checker/pkg/tinybird/batcher.go b/apps/checker/pkg/tinybird/batcher.go index 79b858dc4..d9bd1d3d8 100644 --- a/apps/checker/pkg/tinybird/batcher.go +++ b/apps/checker/pkg/tinybird/batcher.go @@ -70,7 +70,8 @@ type Batcher struct { mu sync.Mutex closed bool batches map[string]*batch - senders map[string]*sync.Mutex + // senders holds one slot per data source: a flush owns it while sending. + senders map[string]chan struct{} flushes sync.WaitGroup } @@ -94,7 +95,7 @@ func NewBatcher(httpClient *http.Client, apiKey string, opts BatcherOptions) *Ba stop: stop, cancelStop: cancelStop, batches: map[string]*batch{}, - senders: map[string]*sync.Mutex{}, + senders: map[string]chan struct{}{}, } } @@ -187,22 +188,28 @@ func (b *Batcher) detachLocked(target *batch) { sender, ok := b.senders[target.dataSource] if !ok { - sender = &sync.Mutex{} + sender = make(chan struct{}, 1) b.senders[target.dataSource] = sender } b.flushes.Add(1) go func() { defer b.flushes.Done() + ctx, cancel := context.WithDeadline(b.stop, target.deadline) + defer cancel() + // 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) + // batch deadline and Close bound the wait for the sender. + select { + case sender <- struct{}{}: + defer func() { <-sender }() + case <-ctx.Done(): + } + b.flush(ctx, target) }() } -func (b *Batcher) flush(target *batch) { +func (b *Batcher) flush(ctx context.Context, target *batch) { start := time.Now() event := FlushEvent{ DataSource: target.dataSource, @@ -212,9 +219,6 @@ func (b *Batcher) flush(target *batch) { 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) diff --git a/apps/checker/pkg/tinybird/batcher_test.go b/apps/checker/pkg/tinybird/batcher_test.go index 740cc91df..9be60bbba 100644 --- a/apps/checker/pkg/tinybird/batcher_test.go +++ b/apps/checker/pkg/tinybird/batcher_test.go @@ -128,11 +128,16 @@ func TestBatcher(t *testing.T) { api := &fakeEventsAPI{} b := tinybird.NewBatcher(api.client(), "apiKey", options()) + dataSources := []string{"ping_response__v8", "tcp_response__v0"} + errs := make([]error, len(dataSources)) var wg sync.WaitGroup - for _, ds := range []string{"ping_response__v8", "tcp_response__v0"} { - wg.Go(func() { require.NoError(t, b.SendEvent(t.Context(), "event", ds)) }) + for i, ds := range dataSources { + wg.Go(func() { errs[i] = b.SendEvent(t.Context(), "event", ds) }) } wg.Wait() + for _, err := range errs { + require.NoError(t, err) + } requests := api.recorded() require.Len(t, requests, 2) @@ -238,7 +243,7 @@ func TestBatcher(t *testing.T) { api := &fakeEventsAPI{respond: func(attempt int) *http.Response { if attempt == 1 { // A slow request holds the sender past the queued batch's deadline. - time.Sleep(2 * opts.RetryBudget) + time.Sleep(5 * opts.RetryBudget) } return status(http.StatusServiceUnavailable, nil) }} @@ -254,6 +259,7 @@ func TestBatcher(t *testing.T) { start := time.Now() err := b.SendEvent(t.Context(), "event", "test") require.Len(t, api.recorded(), 1) + // It fails on its own deadline, not once the first flush frees the sender. require.Less(t, time.Since(start), 3*opts.RetryBudget) require.ErrorIs(t, err, context.DeadlineExceeded) require.ErrorContains(t, err, "flush not attempted") @@ -279,7 +285,8 @@ func TestBatcher(t *testing.T) { select { case err := <-sent: require.ErrorContains(t, err, "unexpected status code: 503") - default: + // The row's error is already queued; only the SendEvent goroutine has to be scheduled. + case <-time.After(time.Second): t.Fatal("the row should fail as soon as Close returns") } }) diff --git a/apps/private-location/fly.toml b/apps/private-location/fly.toml index 52623abc6..bd8177ca6 100644 --- a/apps/private-location/fly.toml +++ b/apps/private-location/fly.toml @@ -5,6 +5,8 @@ app = "openstatus-private-location" primary_region = "ams" +# Shutdown drains HTTP (5s) then flushes pending Tinybird batches (5s). +kill_timeout = "15s" [build] dockerfile = "./Dockerfile" diff --git a/apps/private-location/go.mod b/apps/private-location/go.mod index b64d727d0..b451d9cf2 100644 --- a/apps/private-location/go.mod +++ b/apps/private-location/go.mod @@ -10,7 +10,6 @@ require ( github.com/jmoiron/sqlx v1.4.0 github.com/joho/godotenv v1.5.1 github.com/libsql/sqlite-antlr4-parser v0.0.0-20240327125255-dbf53b6cbf06 - github.com/openstatushq/openstatus/apps/checker v0.0.0-20251012205355-e366f661c23e github.com/stretchr/testify v1.11.1 github.com/tursodatabase/go-libsql v0.0.0-20260424063416-3051e37e6e04 go.opentelemetry.io/contrib/bridges/otelslog v0.14.0 @@ -31,10 +30,7 @@ require ( github.com/go-logr/logr v1.4.3 // indirect github.com/go-logr/stdr v1.2.2 // indirect github.com/grpc-ecosystem/grpc-gateway/v2 v2.27.3 // indirect - github.com/mattn/go-colorable v0.1.13 // indirect - github.com/mattn/go-isatty v0.0.20 // indirect github.com/pmezard/go-difflib v1.0.0 // indirect - github.com/rs/zerolog v1.34.0 // indirect go.opentelemetry.io/auto/sdk v1.2.1 // indirect go.opentelemetry.io/otel/metric v1.39.0 // indirect go.opentelemetry.io/otel/trace v1.39.0 // indirect diff --git a/apps/private-location/go.sum b/apps/private-location/go.sum index 4f5d71fb2..427c614cd 100644 --- a/apps/private-location/go.sum +++ b/apps/private-location/go.sum @@ -10,7 +10,6 @@ github.com/cenkalti/backoff/v5 v5.0.3 h1:ZN+IMa753KfX5hd8vVaMixjnqRZ3y8CuJKRKj1x github.com/cenkalti/backoff/v5 v5.0.3/go.mod h1:rkhZdG3JZukswDf7f0cwqPNk4K0sa+F97BxZthm/crw= github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs= github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs= -github.com/coreos/go-systemd/v22 v22.5.0/go.mod h1:Y58oyj3AT4RCenI/lSvhwexgC+NSVTIJ3seZv2GcEnc= github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/go-chi/chi/v5 v5.2.3 h1:WQIt9uxdsAbgIYgid+BpYc+liqQZGMHRaUwp0JUcvdE= @@ -24,7 +23,6 @@ github.com/go-logr/stdr v1.2.2 h1:hSWxHoqTgW2S2qGc0LTAI563KZ5YKYRhT3MFKZMbjag= github.com/go-logr/stdr v1.2.2/go.mod h1:mMo/vtBO5dYbehREoey6XUKy/eSumjCCveDpRre4VKE= github.com/go-sql-driver/mysql v1.8.1 h1:LedoTUt/eveggdHS9qUFC1EFSa8bU2+1pZjSRpvNJ1Y= github.com/go-sql-driver/mysql v1.8.1/go.mod h1:wEBSXgmK//2ZFJyE+qWnIsVGmvmEKlqwuVSjsCm7DZg= -github.com/godbus/dbus/v5 v5.0.4/go.mod h1:xhWf0FNVPg57R7Z0UbKHbJfkEywrmjJnf7w5xrFpKfA= github.com/golang/protobuf v1.5.4 h1:i7eJL8qZTpSEXOPTxNKhASYpMn+8e5Q6AdndVa1dWek= github.com/golang/protobuf v1.5.4/go.mod h1:lnTiLA8Wa4RWRcIUkrtSVa5nRhsEGBg48fD6rSs7xps= github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8= @@ -45,25 +43,14 @@ github.com/lib/pq v1.10.9 h1:YXG7RB+JIjhP29X+OtkiDnYaXQwpS4JEWq7dtCCRUEw= github.com/lib/pq v1.10.9/go.mod h1:AlVN5x4E4T544tWzH6hKfbfQvm3HdbOxrmggDNAPY9o= github.com/libsql/sqlite-antlr4-parser v0.0.0-20240327125255-dbf53b6cbf06 h1:JLvn7D+wXjH9g4Jsjo+VqmzTUpl/LX7vfr6VOfSWTdM= github.com/libsql/sqlite-antlr4-parser v0.0.0-20240327125255-dbf53b6cbf06/go.mod h1:FUkZ5OHjlGPjnM2UyGJz9TypXQFgYqw6AFNO1UiROTM= -github.com/mattn/go-colorable v0.1.13 h1:fFA4WZxdEF4tXPZVKMLwD8oUnCTTo08duU7wxecdEvA= -github.com/mattn/go-colorable v0.1.13/go.mod h1:7S9/ev0klgBDR4GtXTXX8a3vIGJpMovkB8vQcUbaXHg= -github.com/mattn/go-isatty v0.0.16/go.mod h1:kYGgaQfpe5nmfYZH+SKPsOc2e4SrIfOl2e/yFXSvRLM= -github.com/mattn/go-isatty v0.0.19/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y= -github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY= -github.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y= github.com/mattn/go-sqlite3 v1.14.22 h1:2gZY6PC6kBnID23Tichd1K+Z0oS6nE/XwU+Vz/5o4kU= github.com/mattn/go-sqlite3 v1.14.22/go.mod h1:Uh1q+B4BYcTPb+yiD3kU8Ct7aC0hY9fxUwlHK0RXw+Y= -github.com/openstatushq/openstatus/apps/checker v0.0.0-20251012205355-e366f661c23e h1:54C0zQNHzGszQseO2QcNzM8fL7vyAYk03pRtrJIyoV0= -github.com/openstatushq/openstatus/apps/checker v0.0.0-20251012205355-e366f661c23e/go.mod h1:R84xAJYFys7XOZTDk/AyjJi4Ga9ovtLhJsfTLgTsYKg= github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4= github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0= github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= github.com/rogpeppe/go-internal v1.14.1 h1:UQB4HGPB6osV0SQTLymcB4TgvyWu6ZyliaW0tI/otEQ= github.com/rogpeppe/go-internal v1.14.1/go.mod h1:MaRKkUm5W0goXpeCfT7UZI6fk/L7L7so1lCWt35ZSgc= -github.com/rs/xid v1.6.0/go.mod h1:7XoLgs4eV+QndskICGsho+ADou8ySMSjJKDIan90Nz0= -github.com/rs/zerolog v1.34.0 h1:k43nTLIwcTVQAncfCw4KZ2VY6ukYoZaBPNOE8txlOeY= -github.com/rs/zerolog v1.34.0/go.mod h1:bJsvje4Z08ROH4Nhs5iH600c3IkWhwp44iRc54W6wYQ= github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu7U= github.com/stretchr/testify v1.11.1/go.mod h1:wZwfW3scLgRK+23gO65QZefKpKQRnfz6sD981Nm4B6U= github.com/tursodatabase/go-libsql v0.0.0-20260424063416-3051e37e6e04 h1:9nlqEMruvXDPynGbZ0RE67kKnkkg3NdnjGccvRABefc= @@ -98,9 +85,6 @@ golang.org/x/net v0.47.0 h1:Mx+4dIFzqraBXUugkia1OOvlD6LemFo1ALMHjrXDOhY= golang.org/x/net v0.47.0/go.mod h1:/jNxtkgq5yWUGYkaZGqo27cfGZ1c5Nen03aYrrKpVRU= golang.org/x/sync v0.18.0 h1:kr88TuHDroi+UVf+0hZnirlk8o8T+4MrK6mr60WkH/I= golang.org/x/sync v0.18.0/go.mod h1:9KTHXmSnoGruLpwFjVSX0lNNA75CykiMECbovNTZqGI= -golang.org/x/sys v0.0.0-20220811171246-fbc7d0a398ab/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= -golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= -golang.org/x/sys v0.12.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.39.0 h1:CvCKL8MeisomCi6qNZ+wbb0DN9E5AATixKsvNtMoMFk= golang.org/x/sys v0.39.0/go.mod h1:OgkHotnGiDImocRcuBABYBEXf8A9a87e/uXjp9XT3ks= golang.org/x/text v0.31.0 h1:aC8ghyu4JhP8VojJ2lEHBnochRno1sgL6nEi9WGFGMM= diff --git a/apps/private-location/internal/database/database_test.go b/apps/private-location/internal/database/database_test.go index b2deca705..9b04a8260 100644 --- a/apps/private-location/internal/database/database_test.go +++ b/apps/private-location/internal/database/database_test.go @@ -1,72 +1,14 @@ package database import ( - "encoding/json" - "fmt" - "net/http" - "net/http/httptest" "path/filepath" - "sync" "testing" ) func TestOpenRemoteDoesNotReuseExpiredStreams(t *testing.T) { - var mu sync.Mutex - streams := map[string]bool{} - nextStream := 0 - server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - var request struct { - Baton *string `json:"baton"` - Requests []struct { - Type string `json:"type"` - } `json:"requests"` - Batch json.RawMessage `json:"batch"` - } - if err := json.NewDecoder(r.Body).Decode(&request); err != nil { - http.Error(w, err.Error(), http.StatusBadRequest) - return - } - mu.Lock() - defer mu.Unlock() - w.Header().Set("Content-Type", "application/json") - var baton string - if request.Baton == nil { - nextStream++ - baton = fmt.Sprintf("stream-%d", nextStream) - streams[baton] = false - } else { - baton = *request.Baton - if streams[baton] { - w.WriteHeader(http.StatusBadRequest) - fmt.Fprint(w, `{"message":"The stream has expired due to inactivity","code":"STREAM_EXPIRED"}`) - return - } - } - if len(request.Batch) > 0 { - fmt.Fprintf(w, "{\"baton\":%q,\"base_url\":null}\n", baton) - fmt.Fprintln(w, `{"type":"step_begin","step":0,"cols":[{"name":"1","decltype":null}]}`) - fmt.Fprintln(w, `{"type":"row","row":[{"type":"integer","value":"1"}]}`) - fmt.Fprintln(w, `{"type":"step_end","affected_row_count":0,"last_insert_rowid":null}`) - fmt.Fprintln(w, `{"type":"replication_index","replication_index":null}`) - return - } - if len(request.Requests) == 1 && request.Requests[0].Type == "describe" { - fmt.Fprintf(w, `{"baton":%q,"base_url":null,"results":[{"type":"ok","response":{"type":"describe","result":{"params":[],"cols":[{"name":"1","decltype":null}],"is_explain":false,"is_readonly":true}}}]}`, baton) - return - } - if len(request.Requests) == 1 && request.Requests[0].Type == "close" { - fmt.Fprint(w, `{"baton":null,"base_url":null,"results":[{"type":"ok","response":{"type":"close"}}]}`) - return - } - http.Error(w, "unexpected Hrana request", http.StatusBadRequest) - })) - defer server.Close() + hrana := &fakeHrana{} + db := newFakeHrana(t, hrana) - db, err := Open(server.URL, "") - if err != nil { - t.Fatal(err) - } - defer db.Close() var value int if err := db.Get(&value, "SELECT 1"); err != nil { t.Fatal(err) @@ -76,11 +18,7 @@ func TestOpenRemoteDoesNotReuseExpiredStreams(t *testing.T) { } // Expire existing streams without a wall-clock delay. - mu.Lock() - for baton := range streams { - streams[baton] = true - } - mu.Unlock() + hrana.expireStreams() if err := db.Get(&value, "SELECT 1"); err != nil { t.Fatalf("query after stream expiry: %v", err) diff --git a/apps/private-location/internal/database/hrana_test.go b/apps/private-location/internal/database/hrana_test.go new file mode 100644 index 000000000..64f3d8c67 --- /dev/null +++ b/apps/private-location/internal/database/hrana_test.go @@ -0,0 +1,142 @@ +package database + +import ( + "encoding/json" + "fmt" + "net/http" + "net/http/httptest" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/jmoiron/sqlx" +) + +// fakeHrana is a minimal libSQL remote server: it answers SELECT 1 for any +// statement and tracks streams so tests can see what the pool holds open. +type fakeHrana struct { + // stall, when set, blocks every statement until it is closed or the client + // gives up on the request. + stall chan struct{} + delay time.Duration + + mu sync.Mutex + nextStream int + open map[string]bool + expired map[string]bool + + inFlight atomic.Int32 + maxInFlight atomic.Int32 + statements atomic.Int32 +} + +func newFakeHrana(t *testing.T, f *fakeHrana) *sqlx.DB { + t.Helper() + f.open = map[string]bool{} + f.expired = map[string]bool{} + server := httptest.NewServer(f) + t.Cleanup(server.Close) + + db, err := Open(server.URL, "") + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { db.Close() }) + + return db +} + +func (f *fakeHrana) openStreams() int { + f.mu.Lock() + defer f.mu.Unlock() + + return len(f.open) +} + +// expireStreams makes the server reject every stream opened so far, as libSQL +// does once a stream sits idle too long. +func (f *fakeHrana) expireStreams() { + f.mu.Lock() + defer f.mu.Unlock() + for i := 1; i <= f.nextStream; i++ { + f.expired[fmt.Sprintf("stream-%d", i)] = true + } +} + +func (f *fakeHrana) ServeHTTP(w http.ResponseWriter, r *http.Request) { + var request struct { + Baton *string `json:"baton"` + Requests []struct { + Type string `json:"type"` + } `json:"requests"` + Batch json.RawMessage `json:"batch"` + } + if err := json.NewDecoder(r.Body).Decode(&request); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + + f.mu.Lock() + var baton string + if request.Baton == nil { + f.nextStream++ + baton = fmt.Sprintf("stream-%d", f.nextStream) + f.open[baton] = true + } else { + baton = *request.Baton + } + expired := f.expired[baton] + f.mu.Unlock() + + w.Header().Set("Content-Type", "application/json") + + if expired { + w.WriteHeader(http.StatusBadRequest) + fmt.Fprint(w, `{"message":"The stream has expired due to inactivity","code":"STREAM_EXPIRED"}`) + return + } + + if len(request.Requests) == 1 && request.Requests[0].Type == "close" { + f.mu.Lock() + delete(f.open, baton) + f.mu.Unlock() + fmt.Fprint(w, `{"baton":null,"base_url":null,"results":[{"type":"ok","response":{"type":"close"}}]}`) + return + } + + inFlight := f.inFlight.Add(1) + defer f.inFlight.Add(-1) + for { + seen := f.maxInFlight.Load() + if inFlight <= seen || f.maxInFlight.CompareAndSwap(seen, inFlight) { + break + } + } + + if f.stall != nil { + select { + case <-f.stall: + case <-r.Context().Done(): + return + } + } + if f.delay > 0 { + time.Sleep(f.delay) + } + + if len(request.Batch) > 0 { + f.statements.Add(1) + fmt.Fprintf(w, "{\"baton\":%q,\"base_url\":null}\n", baton) + fmt.Fprintln(w, `{"type":"step_begin","step":0,"cols":[{"name":"1","decltype":null}]}`) + fmt.Fprintln(w, `{"type":"row","row":[{"type":"integer","value":"1"}]}`) + fmt.Fprintln(w, `{"type":"step_end","affected_row_count":0,"last_insert_rowid":null}`) + fmt.Fprintln(w, `{"type":"replication_index","replication_index":null}`) + return + } + if len(request.Requests) == 1 && request.Requests[0].Type == "describe" { + fmt.Fprintf(w, `{"baton":%q,"base_url":null,"results":[{"type":"ok","response":{"type":"describe","result":{"params":[],"cols":[{"name":"1","decltype":null}],"is_explain":false,"is_readonly":true}}}]}`, baton) + return + } + http.Error(w, "unexpected Hrana request", http.StatusBadRequest) +} diff --git a/apps/private-location/internal/database/pool_test.go b/apps/private-location/internal/database/pool_test.go new file mode 100644 index 000000000..63aa557ac --- /dev/null +++ b/apps/private-location/internal/database/pool_test.go @@ -0,0 +1,148 @@ +package database + +import ( + "context" + "errors" + "fmt" + "sync" + "testing" + "time" + + "github.com/jmoiron/sqlx" +) + +func waitFor(t *testing.T, msg string, cond func() bool) { + t.Helper() + deadline := time.Now().Add(2 * time.Second) + for !cond() { + if time.Now().After(deadline) { + t.Fatal(msg) + } + time.Sleep(5 * time.Millisecond) + } +} + +func TestPoolNeverExceedsMaxOpenConns(t *testing.T) { + hrana := &fakeHrana{delay: 20 * time.Millisecond} + db := newFakeHrana(t, hrana) + + const burst = 4 * maxOpenConns + errs := make(chan error, burst) + var wg sync.WaitGroup + for range burst { + wg.Go(func() { + var value int + errs <- db.GetContext(context.Background(), &value, "SELECT 1") + }) + } + wg.Wait() + close(errs) + for err := range errs { + if err != nil { + t.Fatalf("query in burst: %v", err) + } + } + + if got := hrana.maxInFlight.Load(); got > maxOpenConns { + t.Fatalf("server saw %d concurrent statements, want at most %d", got, maxOpenConns) + } + if got := db.Stats().InUse; got != 0 { + t.Fatalf("%d connections still in use after the burst", got) + } +} + +func TestPoolDoesNotLeakRemoteStreams(t *testing.T) { + hrana := &fakeHrana{} + db := newFakeHrana(t, hrana) + + for range 3 * maxOpenConns { + var value int + if err := db.GetContext(context.Background(), &value, "SELECT 1"); err != nil { + t.Fatal(err) + } + } + + // Remote connections are never kept idle, so each stream must be closed on + // the server once its query is done. + waitFor(t, fmt.Sprintf("%d remote streams left open", hrana.openStreams()), func() bool { + return hrana.openStreams() == 0 + }) + if got := db.Stats().OpenConnections; got != 0 { + t.Fatalf("%d connections still open, want 0", got) + } +} + +// stalledHrana returns a server that holds every statement until the test ends. +func stalledHrana(t *testing.T) (*fakeHrana, *sqlx.DB) { + t.Helper() + hrana := &fakeHrana{stall: make(chan struct{})} + db := newFakeHrana(t, hrana) + // Registered last so it runs first: unblock the server before it and the + // pool are closed. + t.Cleanup(func() { close(hrana.stall) }) + + return hrana, db +} + +// The ingest handlers bound every query with dbTimeout. If a timed-out query +// keeps its connection, a slow libSQL server drains the pool within a few +// requests and every later ingest fails waiting for a connection. +func TestPoolReleasesConnectionsWhenQueriesTimeOut(t *testing.T) { + t.Skip("go-libsql ignores ctx in QueryContext/ExecContext: a stalled remote query holds its connection until the server answers") + _, db := stalledHrana(t) + + returned := make(chan error, maxOpenConns) + for range maxOpenConns { + go func() { + ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond) + defer cancel() + var value int + returned <- db.GetContext(ctx, &value, "SELECT 1") + }() + } + + timeout := time.After(2 * time.Second) + for range maxOpenConns { + select { + case err := <-returned: + if !errors.Is(err, context.DeadlineExceeded) { + t.Fatalf("stalled query returned %v, want deadline exceeded", err) + } + case <-timeout: + t.Fatalf("queries ignored their 100ms deadline: %d of %d connections held by a stalled server", db.Stats().InUse, maxOpenConns) + } + } + + waitFor(t, fmt.Sprintf("%d connections still in use after their queries timed out", db.Stats().InUse), func() bool { + return db.Stats().InUse == 0 + }) +} + +func TestPoolWaitIsBoundedByContext(t *testing.T) { + _, db := stalledHrana(t) + + for range maxOpenConns { + go func() { + var value int + _ = db.GetContext(context.Background(), &value, "SELECT 1") + }() + } + waitFor(t, "pool never filled", func() bool { return db.Stats().InUse == maxOpenConns }) + + // With every connection busy, a caller must fail on its own deadline rather + // than queue forever behind them. + ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond) + defer cancel() + start := time.Now() + var value int + err := db.GetContext(ctx, &value, "SELECT 1") + if !errors.Is(err, context.DeadlineExceeded) { + t.Fatalf("query on an exhausted pool returned %v, want deadline exceeded", err) + } + if elapsed := time.Since(start); elapsed > time.Second { + t.Fatalf("query on an exhausted pool waited %s", elapsed) + } + if db.Stats().WaitCount == 0 { + t.Fatal("expected the query to wait for a connection") + } +} diff --git a/apps/private-location/internal/server/ingest_batch_test.go b/apps/private-location/internal/server/ingest_batch_test.go new file mode 100644 index 000000000..08c4c6d36 --- /dev/null +++ b/apps/private-location/internal/server/ingest_batch_test.go @@ -0,0 +1,196 @@ +package server_test + +import ( + "context" + "database/sql" + "encoding/json" + "fmt" + "io" + "net/http" + "strings" + "sync" + "testing" + "time" + + "connectrpc.com/connect" + "github.com/openstatushq/openstatus/apps/private-location/internal/server" + "github.com/openstatushq/openstatus/apps/private-location/internal/tinybird" + private_locationv1 "github.com/openstatushq/openstatus/apps/private-location/proto/private_location/v1" + "github.com/stretchr/testify/require" +) + +type eventsRequest struct { + dataSource string + lines []string +} + +type fakeEventsAPI struct { + mu sync.Mutex + requests []eventsRequest + statusCode int +} + +func (f *fakeEventsAPI) client() *http.Client { + return (&interceptorHTTPClient{f: func(req *http.Request) (*http.Response, error) { + body, _ := io.ReadAll(req.Body) + f.mu.Lock() + f.requests = append(f.requests, eventsRequest{ + dataSource: req.URL.Query().Get("name"), + lines: strings.Split(strings.TrimSuffix(string(body), "\n"), "\n"), + }) + f.mu.Unlock() + + statusCode := f.statusCode + if statusCode == 0 { + statusCode = http.StatusAccepted + } + + return &http.Response{StatusCode: statusCode, Body: io.NopCloser(strings.NewReader(""))}, nil + }}).GetHTTPClient() +} + +func (f *fakeEventsAPI) recorded() []eventsRequest { + f.mu.Lock() + defer f.mu.Unlock() + + return append([]eventsRequest(nil), f.requests...) +} + +func newBatcher(t *testing.T, api *fakeEventsAPI) *tinybird.Batcher { + t.Helper() + b := tinybird.NewBatcher(api.client(), "apiKey", tinybird.BatcherOptions{ + FlushInterval: 500 * time.Millisecond, + MaxBytes: 1 << 20, + RetryBudget: time.Second, + }) + t.Cleanup(func() { _ = b.Close(context.Background()) }) + + return b +} + +type httpIngester interface { + IngestHTTP(context.Context, *connect.Request[private_locationv1.IngestHTTPRequest]) (*connect.Response[private_locationv1.IngestHTTPResponse], error) +} + +func ingestHTTP(h httpIngester, id string) error { + req := connect.NewRequest(&private_locationv1.IngestHTTPRequest{ + Id: id, + MonitorId: "5", + Timestamp: 1234567890, + }) + req.Header().Set("openstatus-token", "my-secret-key") + _, err := h.IngestHTTP(context.Background(), req) + + return err +} + +func TestIngest_Batched(t *testing.T) { + t.Run("it should send concurrent ingests as one NDJSON request", func(t *testing.T) { + api := &fakeEventsAPI{} + h := server.NewPrivateLocationServer(testDB(), newBatcher(t, api)) + + const n = 10 + errs := make([]error, n) + var wg sync.WaitGroup + for i := range n { + wg.Go(func() { errs[i] = ingestHTTP(h, fmt.Sprintf("check-%d", i)) }) + } + wg.Wait() + for _, err := range errs { + require.NoError(t, err) + } + + requests := api.recorded() + require.Len(t, requests, 1) + require.Equal(t, tinybird.DatasourceHTTP, requests[0].dataSource) + require.Len(t, requests[0].lines, n) + + ids := make([]string, 0, n) + for _, line := range requests[0].lines { + var row server.PingData + require.NoError(t, json.Unmarshal([]byte(line), &row)) + require.Equal(t, "5", row.MonitorID) + ids = append(ids, row.ID) + } + want := make([]string, 0, n) + for i := range n { + want = append(want, fmt.Sprintf("check-%d", i)) + } + require.ElementsMatch(t, want, ids) + }) + + t.Run("it should batch each data source separately", func(t *testing.T) { + api := &fakeEventsAPI{} + h := server.NewPrivateLocationServer(testDB(), newBatcher(t, api)) + + var httpErr, tcpErr error + var wg sync.WaitGroup + wg.Go(func() { httpErr = ingestHTTP(h, "http-check") }) + wg.Go(func() { + req := connect.NewRequest(&private_locationv1.IngestTCPRequest{ + Id: "tcp-check", + MonitorId: "6", + Timestamp: 1234567890, + }) + req.Header().Set("openstatus-token", "my-secret-key") + _, tcpErr = h.IngestTCP(context.Background(), req) + }) + wg.Wait() + require.NoError(t, httpErr) + require.NoError(t, tcpErr) + + requests := api.recorded() + require.Len(t, requests, 2) + require.ElementsMatch(t, + []string{tinybird.DatasourceHTTP, tinybird.DatasourceTCP}, + []string{requests[0].dataSource, requests[1].dataSource}, + ) + for _, r := range requests { + require.Len(t, r.lines, 1) + } + }) + + t.Run("it should still ack the ingest and record last_seen_at when the flush fails", func(t *testing.T) { + api := &fakeEventsAPI{statusCode: http.StatusBadRequest} + db := testDB() + h := server.NewPrivateLocationServer(db, newBatcher(t, api)) + + require.NoError(t, ingestHTTP(h, "check-rejected")) + require.Len(t, api.recorded(), 1) + + var lastSeen sql.NullInt64 + require.NoError(t, db.Get(&lastSeen, "SELECT last_seen_at FROM private_location WHERE token = ?", "my-secret-key")) + require.True(t, lastSeen.Valid) + }) + + t.Run("it should return once the request context ends, before the batch flushes", func(t *testing.T) { + api := &fakeEventsAPI{} + b := tinybird.NewBatcher(api.client(), "apiKey", tinybird.BatcherOptions{ + FlushInterval: time.Hour, + MaxBytes: 1 << 20, + RetryBudget: time.Second, + }) + h := server.NewPrivateLocationServer(testDB(), b) + + ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond) + defer cancel() + req := connect.NewRequest(&private_locationv1.IngestHTTPRequest{ + Id: "check-cancelled", + MonitorId: "5", + Timestamp: 1234567890, + }) + req.Header().Set("openstatus-token", "my-secret-key") + + start := time.Now() + _, err := h.IngestHTTP(ctx, req) + require.NoError(t, err) + require.Less(t, time.Since(start), time.Second) + require.Empty(t, api.recorded()) + + // The abandoned row is still delivered on shutdown. + require.NoError(t, b.Close(context.Background())) + requests := api.recorded() + require.Len(t, requests, 1) + require.Len(t, requests[0].lines, 1) + }) +} diff --git a/apps/private-location/internal/server/routes.go b/apps/private-location/internal/server/routes.go index fb36cac4e..f03555728 100644 --- a/apps/private-location/internal/server/routes.go +++ b/apps/private-location/internal/server/routes.go @@ -150,16 +150,8 @@ func (s *Server) RegisterRoutes() http.Handler { r.Get("/health", s.healthHandler) - tinyBirdToken := os.Getenv("TINYBIRD_TOKEN") - - httpClient := &http.Client{ - Timeout: 45 * time.Second, - } - - tinybirdClient := tinybird.NewClient(httpClient, tinyBirdToken) - - privateLocationServer := NewPrivateLocationServer(s.db, tinybirdClient) - privateLocationServer.WorkflowsClient = workflows.NewClient(httpClient, os.Getenv("CRON_SECRET")) + privateLocationServer := NewPrivateLocationServer(s.db, s.tbClient) + privateLocationServer.WorkflowsClient = workflows.NewClient(s.httpClient, os.Getenv("CRON_SECRET")) path, handler := v1.NewPrivateLocationServiceHandler(privateLocationServer) r.Group(func(r chi.Router) { diff --git a/apps/private-location/internal/server/server.go b/apps/private-location/internal/server/server.go index c0a9bd88b..ae8992d6d 100644 --- a/apps/private-location/internal/server/server.go +++ b/apps/private-location/internal/server/server.go @@ -21,19 +21,28 @@ import ( semconv "go.opentelemetry.io/otel/semconv/v1.21.0" "github.com/openstatushq/openstatus/apps/private-location/internal/database" + "github.com/openstatushq/openstatus/apps/private-location/internal/logs" + "github.com/openstatushq/openstatus/apps/private-location/internal/tinybird" ) - type Server struct { port int db *sqlx.DB logger *slog.Logger logProvider *sdklog.LoggerProvider + httpClient *http.Client + tbClient *tinybird.Batcher } // instanceID is generated once at startup var instanceID = uuid.New().String() +const ( + writeTimeout = 30 * time.Second + // flushTimeout bounds the Tinybird flush on shutdown, after the HTTP drain. + flushTimeout = 5 * time.Second +) + // NewServer returns an HTTP server and a cleanup function to shutdown the log provider. func NewServer() (*http.Server, func(context.Context)) { portStr := os.Getenv("PORT") @@ -48,11 +57,18 @@ func NewServer() (*http.Server, func(context.Context)) { logger, logProvider := setupLogger() + httpClient := &http.Client{ + Timeout: 45 * time.Second, + } + tbClient := tinybird.NewBatcher(httpClient, os.Getenv("TINYBIRD_TOKEN"), tinybirdOptions()) + newServer := &Server{ port: port, db: database.New(), logger: logger, logProvider: logProvider, + httpClient: httpClient, + tbClient: tbClient, } // Declare Server config @@ -61,7 +77,7 @@ func NewServer() (*http.Server, func(context.Context)) { Handler: newServer.RegisterRoutes(), IdleTimeout: time.Minute, ReadTimeout: 10 * time.Second, - WriteTimeout: 30 * time.Second, + WriteTimeout: writeTimeout, } // Startup logging @@ -74,6 +90,13 @@ func NewServer() (*http.Server, func(context.Context)) { // Return cleanup function for graceful shutdown cleanup := func(ctx context.Context) { + // Ingests that finished during the HTTP drain may still sit in a batch. + // The drain can use up ctx, so the flush gets its own bound. + flushCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), flushTimeout) + defer cancel() + if err := tbClient.Close(flushCtx); err != nil { + slog.Error("failed to flush tinybird batches", "error", err) + } if logProvider != nil { logProvider.Shutdown(ctx) } @@ -137,3 +160,71 @@ func env(key, fallback string) string { } return fallback } + +func tinybirdOptions() tinybird.BatcherOptions { + opts := tinybird.DefaultBatcherOptions() + opts.FlushInterval = envDuration("TINYBIRD_FLUSH_INTERVAL", opts.FlushInterval) + opts.FlushJitter = envDuration("TINYBIRD_FLUSH_JITTER", opts.FlushJitter) + opts.MaxBytes = envInt("TINYBIRD_MAX_BATCH_BYTES", opts.MaxBytes) + opts.RetryBudget = 20 * time.Second + // A zero budget would expire before the first request. + if retryBudget := envDuration("TINYBIRD_RETRY_BUDGET", 0); retryBudget > 0 { + opts.RetryBudget = retryBudget + } + // The ingest handler waits for its row's flush, so the flush delay plus the + // retry budget must stay under the server's WriteTimeout. + if maxBudget := writeTimeout - opts.FlushInterval - opts.FlushJitter - time.Second; opts.RetryBudget > maxBudget { + slog.Warn("tinybird retry budget exceeds the write timeout, clamping it", + "retry_budget", opts.RetryBudget, + "max_retry_budget", maxBudget, + ) + opts.RetryBudget = max(maxBudget, time.Second) + } + opts.OnFlush = logTinybirdFlush + + return opts +} + +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 || logs.ShouldSample(event) { + slog.LogAttrs(ctx, slog.LevelInfo, "tinybird flush", logs.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 { + return fallback + } + + return value +} diff --git a/apps/private-location/internal/server/tinybird_options_test.go b/apps/private-location/internal/server/tinybird_options_test.go new file mode 100644 index 000000000..3e2406c24 --- /dev/null +++ b/apps/private-location/internal/server/tinybird_options_test.go @@ -0,0 +1,58 @@ +package server + +import ( + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +func TestTinybirdOptions(t *testing.T) { + t.Run("it should keep the retry budget under the server WriteTimeout by default", func(t *testing.T) { + opts := tinybirdOptions() + + require.Equal(t, 20*time.Second, opts.RetryBudget) + require.Less(t, opts.FlushInterval+opts.FlushJitter+opts.RetryBudget, writeTimeout) + require.NotNil(t, opts.OnFlush) + }) + + t.Run("it should read overrides from the environment", func(t *testing.T) { + t.Setenv("TINYBIRD_FLUSH_INTERVAL", "500ms") + t.Setenv("TINYBIRD_FLUSH_JITTER", "0s") + t.Setenv("TINYBIRD_MAX_BATCH_BYTES", "1024") + t.Setenv("TINYBIRD_RETRY_BUDGET", "5s") + + opts := tinybirdOptions() + + require.Equal(t, 500*time.Millisecond, opts.FlushInterval) + require.Zero(t, opts.FlushJitter) + require.Equal(t, 1024, opts.MaxBytes) + require.Equal(t, 5*time.Second, opts.RetryBudget) + }) + + t.Run("it should clamp a retry budget override to fit the WriteTimeout", func(t *testing.T) { + t.Setenv("TINYBIRD_FLUSH_INTERVAL", "2s") + t.Setenv("TINYBIRD_FLUSH_JITTER", "1s") + t.Setenv("TINYBIRD_RETRY_BUDGET", "1m") + + opts := tinybirdOptions() + + require.Equal(t, 26*time.Second, opts.RetryBudget) + require.Less(t, opts.FlushInterval+opts.FlushJitter+opts.RetryBudget, writeTimeout) + }) + + t.Run("it should ignore invalid or non-positive overrides", func(t *testing.T) { + defaults := tinybirdOptions() + t.Setenv("TINYBIRD_FLUSH_INTERVAL", "soon") + t.Setenv("TINYBIRD_FLUSH_JITTER", "-1s") + t.Setenv("TINYBIRD_MAX_BATCH_BYTES", "0") + t.Setenv("TINYBIRD_RETRY_BUDGET", "0s") + + opts := tinybirdOptions() + + require.Equal(t, defaults.FlushInterval, opts.FlushInterval) + require.Equal(t, defaults.FlushJitter, opts.FlushJitter) + require.Equal(t, defaults.MaxBytes, opts.MaxBytes) + require.Equal(t, defaults.RetryBudget, opts.RetryBudget) + }) +} diff --git a/apps/private-location/internal/tinybird/batcher.go b/apps/private-location/internal/tinybird/batcher.go new file mode 100644 index 000000000..d9bd1d3d8 --- /dev/null +++ b/apps/private-location/internal/tinybird/batcher.go @@ -0,0 +1,326 @@ +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 holds one slot per data source: a flush owns it while sending. + senders map[string]chan struct{} + 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]chan struct{}{}, + } +} + +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 = make(chan struct{}, 1) + b.senders[target.dataSource] = sender + } + + b.flushes.Add(1) + go func() { + defer b.flushes.Done() + ctx, cancel := context.WithDeadline(b.stop, target.deadline) + defer cancel() + + // One request per data source at a time, so retries don't pile up; the + // batch deadline and Close bound the wait for the sender. + select { + case sender <- struct{}{}: + defer func() { <-sender }() + case <-ctx.Done(): + } + b.flush(ctx, target) + }() +} + +func (b *Batcher) flush(ctx context.Context, target *batch) { + start := time.Now() + event := FlushEvent{ + DataSource: target.dataSource, + Rows: len(target.waiters), + Bytes: target.body.Len(), + RateLimitLimit: -1, + RateLimitRemaining: -1, + } + + 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) + // It fails on its own deadline, not once the first flush frees the sender. + 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") + // The row's error is already queued; only the SendEvent goroutine has to be scheduled. + case <-time.After(time.Second): + t.Fatal("the row should fail as soon as Close returns") + } + }) +} diff --git a/apps/private-location/internal/tinybird/client_test.go b/apps/private-location/internal/tinybird/client_test.go index bc0670bed..6d9155beb 100644 --- a/apps/private-location/internal/tinybird/client_test.go +++ b/apps/private-location/internal/tinybird/client_test.go @@ -5,7 +5,7 @@ import ( "net/http" "testing" - "github.com/openstatushq/openstatus/apps/checker/pkg/tinybird" + "github.com/openstatushq/openstatus/apps/private-location/internal/tinybird" "github.com/stretchr/testify/require" ) diff --git a/apps/private-location/internal/tinybird/export_test.go b/apps/private-location/internal/tinybird/export_test.go new file mode 100644 index 000000000..c1588afa1 --- /dev/null +++ b/apps/private-location/internal/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