From 51e2d798fcd9c44bba865671e6f6356a4eb080ba Mon Sep 17 00:00:00 2001 From: Filip Petkovski Date: Fri, 8 Jul 2022 06:35:34 +0200 Subject: [PATCH] Simplify TSDBStats interface Signed-off-by: Filip Petkovski --- pkg/api/status/v1.go | 1 - pkg/receive/handler.go | 10 ++---- pkg/receive/multitsdb.go | 40 +++++++++++------------- pkg/receive/multitsdb_test.go | 59 +++++++++++++++++++++++++++++++++++ 4 files changed, 80 insertions(+), 30 deletions(-) diff --git a/pkg/api/status/v1.go b/pkg/api/status/v1.go index 4fb00888ec2..0f10a5b5b8d 100644 --- a/pkg/api/status/v1.go +++ b/pkg/api/status/v1.go @@ -21,7 +21,6 @@ import ( "net/http" "github.com/prometheus/client_golang/prometheus" - "github.com/prometheus/prometheus/tsdb" "github.com/go-kit/log" diff --git a/pkg/receive/handler.go b/pkg/receive/handler.go index 7fd49b229e7..8bb14e1808b 100644 --- a/pkg/receive/handler.go +++ b/pkg/receive/handler.go @@ -18,7 +18,6 @@ import ( "time" "github.com/thanos-io/thanos/pkg/api" - statusapi "github.com/thanos-io/thanos/pkg/api/status" "github.com/thanos-io/thanos/pkg/logging" @@ -288,19 +287,14 @@ func (h *Handler) getStats(r *http.Request, statsByLabelName string) ([]statusap } if getAllTenantStats { - return h.options.TSDBStats.AllTenantStats(statsByLabelName), nil + return h.options.TSDBStats.TenantStats(statsByLabelName), nil } if tenantID == "" { tenantID = h.options.DefaultTenantID } - stats := h.options.TSDBStats.SingleTenantStats(tenantID, statsByLabelName) - if stats == nil { - return nil, nil - } - - return []statusapi.TenantStats{*stats}, nil + return h.options.TSDBStats.TenantStats(statsByLabelName, tenantID), nil } // Close stops the Handler. diff --git a/pkg/receive/multitsdb.go b/pkg/receive/multitsdb.go index 284a57f24cd..dcee55bdafc 100644 --- a/pkg/receive/multitsdb.go +++ b/pkg/receive/multitsdb.go @@ -9,11 +9,10 @@ import ( "os" "path" "path/filepath" + "sort" "sync" "time" - "github.com/thanos-io/thanos/pkg/api/status" - "github.com/go-kit/log" "github.com/go-kit/log/level" "github.com/pkg/errors" @@ -21,6 +20,7 @@ import ( "github.com/prometheus/prometheus/model/labels" "github.com/prometheus/prometheus/storage" "github.com/prometheus/prometheus/tsdb" + "github.com/thanos-io/thanos/pkg/api/status" "go.uber.org/atomic" "golang.org/x/sync/errgroup" @@ -35,8 +35,9 @@ import ( ) type TSDBStats interface { - SingleTenantStats(tenantID, statsByLabelName string) *status.TenantStats - AllTenantStats(statsByLabelName string) []status.TenantStats + // TenantStats returns TSDB head stats for the given tenants. + // If no tenantIDs are provided, stats for all tenants are returned. + TenantStats(statsByLabelName string, tenantIDs ...string) []status.TenantStats } type MultiTSDB struct { @@ -389,32 +390,26 @@ func (t *MultiTSDB) TSDBExemplars() map[string]*exemplars.TSDB { return res } -func (t *MultiTSDB) SingleTenantStats(tenantID, statsByLabelName string) *status.TenantStats { +func (t *MultiTSDB) TenantStats(statsByLabelName string, tenantIDs ...string) []status.TenantStats { t.mtx.RLock() defer t.mtx.RUnlock() - - tenant, ok := t.tenants[tenantID] - if !ok { - return nil - } - - stats := tenant.readyS.get().db.Head().Stats(statsByLabelName) - return &status.TenantStats{ - Tenant: tenantID, - Stats: stats, + if len(tenantIDs) == 0 { + for tenantID := range t.tenants { + tenantIDs = append(tenantIDs, tenantID) + } } -} - -func (t *MultiTSDB) AllTenantStats(statsByLabelName string) []status.TenantStats { - t.mtx.RLock() - defer t.mtx.RUnlock() var ( mu sync.Mutex wg sync.WaitGroup result = make([]status.TenantStats, 0, len(t.tenants)) ) - for tenantID, tenantInstance := range t.tenants { + for _, tenantID := range tenantIDs { + tenantInstance, ok := t.tenants[tenantID] + if !ok { + continue + } + wg.Add(1) go func(tenantID string, tenantInstance *tenant) { defer wg.Done() @@ -430,6 +425,9 @@ func (t *MultiTSDB) AllTenantStats(statsByLabelName string) []status.TenantStats } wg.Wait() + sort.Slice(result, func(i, j int) bool { + return result[i].Tenant < result[j].Tenant + }) return result } diff --git a/pkg/receive/multitsdb_test.go b/pkg/receive/multitsdb_test.go index efa55a8b4fc..917353d0740 100644 --- a/pkg/receive/multitsdb_test.go +++ b/pkg/receive/multitsdb_test.go @@ -477,6 +477,65 @@ func TestMultiTSDBPrune(t *testing.T) { } } +func TestMultiTSDBStats(t *testing.T) { + tests := []struct { + name string + tenants []string + expectedStats int + }{ + { + name: "single tenant", + tenants: []string{"foo"}, + expectedStats: 1, + }, + { + name: "missing tenant", + tenants: []string{"missing-foo"}, + expectedStats: 0, + }, + { + name: "multiple tenants with missing tenant", + tenants: []string{"foo", "missing-foo"}, + expectedStats: 1, + }, + { + name: "all tenants", + tenants: []string{"foo", "bar", "baz"}, + expectedStats: 3, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + dir, err := ioutil.TempDir("", "tsdb-stats") + testutil.Ok(t, err) + defer func() { testutil.Ok(t, os.RemoveAll(dir)) }() + + m := NewMultiTSDB(dir, log.NewNopLogger(), prometheus.NewRegistry(), + &tsdb.Options{ + MinBlockDuration: (2 * time.Hour).Milliseconds(), + MaxBlockDuration: (2 * time.Hour).Milliseconds(), + RetentionDuration: (6 * time.Hour).Milliseconds(), + }, + labels.FromStrings("replica", "test"), + "tenant_id", + nil, + false, + metadata.NoneFunc, + ) + defer func() { testutil.Ok(t, m.Close()) }() + + testutil.Ok(t, appendSample(m, "foo", time.Now())) + testutil.Ok(t, appendSample(m, "bar", time.Now())) + testutil.Ok(t, appendSample(m, "baz", time.Now())) + testutil.Equals(t, 3, len(m.TSDBStores())) + + stats := m.TenantStats(labels.MetricName, test.tenants...) + testutil.Equals(t, test.expectedStats, len(stats)) + }) + } +} + func appendSample(m *MultiTSDB, tenant string, timestamp time.Time) error { ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) defer cancel()