Simplify TSDBStats interface

Signed-off-by: Filip Petkovski <filip.petkovsky@gmail.com>

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"

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.

pkg/receive/multitsdb.go

@@ -9,18 +9,18 @@ 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"
 	"github.com/prometheus/client_golang/prometheus"
 	"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
 }
 

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()