Implement result cache for tenant query federation (grafana#3640)

Signed-off-by: Christian Simon <[email protected]>

composite_store.go

@@ -19,7 +19,7 @@ type StoreLimits interface {
 }
 
 type CacheGenNumLoader interface {
-	GetStoreCacheGenNumber(userID string) string
+	GetStoreCacheGenNumber(tenantIDs []string) string
 }
 
 // Store for chunks.
@@ -217,7 +217,7 @@ func (c compositeStore) forStores(ctx context.Context, userID string, from, thro
 		return nil
 	}
 
-	ctx = c.injectCacheGen(ctx, userID)
+	ctx = c.injectCacheGen(ctx, []string{userID})
 
 	// first, find the schema with the highest start _before or at_ from
 	i := sort.Search(len(c.stores), func(i int) bool {
@@ -262,10 +262,10 @@ func (c compositeStore) forStores(ctx context.Context, userID string, from, thro
 	return nil
 }
 
-func (c compositeStore) injectCacheGen(ctx context.Context, userID string) context.Context {
+func (c compositeStore) injectCacheGen(ctx context.Context, tenantIDs []string) context.Context {
 	if c.cacheGenNumLoader == nil {
 		return ctx
 	}
 
-	return cache.InjectCacheGenNumber(ctx, c.cacheGenNumLoader.GetStoreCacheGenNumber(userID))
+	return cache.InjectCacheGenNumber(ctx, c.cacheGenNumLoader.GetStoreCacheGenNumber(tenantIDs))
 }

purger/tombstones.go

@@ -3,6 +3,7 @@ package purger
 import (
 	"context"
 	"sort"
+	"strconv"
 	"sync"
 	"time"
 
@@ -246,14 +247,64 @@ func (tl *TombstonesLoader) loadPendingTombstones(userID string) error {
 }
 
 // GetStoreCacheGenNumber returns store cache gen number for a user
-func (tl *TombstonesLoader) GetStoreCacheGenNumber(userID string) string {
-	return tl.getCacheGenNumbers(userID).store
-
+func (tl *TombstonesLoader) GetStoreCacheGenNumber(tenantIDs []string) string {
+	return tl.getCacheGenNumbersPerTenants(tenantIDs).store
 }
 
 // GetResultsCacheGenNumber returns results cache gen number for a user
-func (tl *TombstonesLoader) GetResultsCacheGenNumber(userID string) string {
-	return tl.getCacheGenNumbers(userID).results
+func (tl *TombstonesLoader) GetResultsCacheGenNumber(tenantIDs []string) string {
+	return tl.getCacheGenNumbersPerTenants(tenantIDs).results
+}
+
+func (tl *TombstonesLoader) getCacheGenNumbersPerTenants(tenantIDs []string) *cacheGenNumbers {
+	var result cacheGenNumbers
+
+	if len(tenantIDs) == 0 {
+		return &result
+	}
+
+	// keep the maximum value that's currently in result
+	var maxResults, maxStore int
+
+	for pos, tenantID := range tenantIDs {
+		numbers := tl.getCacheGenNumbers(tenantID)
+
+		// handle first tenant in the list
+		if pos == 0 {
+			// short cut if there is only one tenant
+			if len(tenantIDs) == 1 {
+				return numbers
+			}
+
+			// set first tenant string whatever happens next
+			result.results = numbers.results
+			result.store = numbers.store
+		}
+
+		// set results number string if it's higher than the ones before
+		if numbers.results != "" {
+			results, err := strconv.Atoi(numbers.results)
+			if err != nil {
+				level.Error(util.Logger).Log("msg", "error parsing resultsCacheGenNumber", "tenant", tenantID, "err", err)
+			} else if maxResults < results {
+				maxResults = results
+				result.results = numbers.results
+			}
+		}
+
+		// set store number string if it's higher than the ones before
+		if numbers.store != "" {
+			store, err := strconv.Atoi(numbers.store)
+			if err != nil {
+				level.Error(util.Logger).Log("msg", "error parsing storeCacheGenNumber", "tenant", tenantID, "err", err)
+			} else if maxStore < store {
+				maxStore = store
+				result.store = numbers.store
+			}
+		}
+	}
+
+	return &result
 }
 
 func (tl *TombstonesLoader) getCacheGenNumbers(userID string) *cacheGenNumbers {

purger/tombstones_test.go

@@ -8,6 +8,7 @@ import (
 
 	"github.com/prometheus/common/model"
 	"github.com/prometheus/prometheus/promql/parser"
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
 
@@ -133,6 +134,70 @@ func TestTombstonesLoader(t *testing.T) {
 	}
 }
 
+func TestTombstonesLoader_GetCacheGenNumber(t *testing.T) {
+	s := &store{
+		numbers: map[string]*cacheGenNumbers{
+			"tenant-a": {
+				results: "1000",
+				store:   "2050",
+			},
+			"tenant-b": {
+				results: "1050",
+				store:   "2000",
+			},
+			"tenant-c": {
+				results: "",
+				store:   "",
+			},
+			"tenant-d": {
+				results: "results-c",
+				store:   "store-c",
+			},
+		},
+	}
+	tombstonesLoader := NewTombstonesLoader(s, nil)
+
+	for _, tc := range []struct {
+		name                          string
+		expectedResultsCacheGenNumber string
+		expectedStoreCacheGenNumber   string
+		tenantIDs                     []string
+	}{
+		{
+			name:                          "single tenant with numeric values",
+			tenantIDs:                     []string{"tenant-a"},
+			expectedResultsCacheGenNumber: "1000",
+			expectedStoreCacheGenNumber:   "2050",
+		},
+		{
+			name:                          "single tenant with non-numeric values",
+			tenantIDs:                     []string{"tenant-d"},
+			expectedResultsCacheGenNumber: "results-c",
+			expectedStoreCacheGenNumber:   "store-c",
+		},
+		{
+			name:                          "multiple tenants with numeric values",
+			tenantIDs:                     []string{"tenant-a", "tenant-b"},
+			expectedResultsCacheGenNumber: "1050",
+			expectedStoreCacheGenNumber:   "2050",
+		},
+		{
+			name:                          "multiple tenants with numeric and non-numeric values",
+			tenantIDs:                     []string{"tenant-d", "tenant-c", "tenant-b", "tenant-a"},
+			expectedResultsCacheGenNumber: "1050",
+			expectedStoreCacheGenNumber:   "2050",
+		},
+		{
+			name: "no tenants", // not really an expected call, edge case check to avoid any panics
+		},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			assert.Equal(t, tc.expectedResultsCacheGenNumber, tombstonesLoader.GetResultsCacheGenNumber(tc.tenantIDs))
+			assert.Equal(t, tc.expectedStoreCacheGenNumber, tombstonesLoader.GetStoreCacheGenNumber(tc.tenantIDs))
+		})
+	}
+}
+
 func TestTombstonesReloadDoesntDeadlockOnFailure(t *testing.T) {
 	s := &store{}
 	tombstonesLoader := NewTombstonesLoader(s, nil)
@@ -146,10 +211,17 @@ func TestTombstonesReloadDoesntDeadlockOnFailure(t *testing.T) {
 }
 
 type store struct {
-	err error
+	numbers map[string]*cacheGenNumbers
+	err     error
 }
 
 func (f *store) getCacheGenerationNumbers(ctx context.Context, user string) (*cacheGenNumbers, error) {
+	if f.numbers != nil {
+		number, ok := f.numbers[user]
+		if ok {
+			return number, nil
+		}
+	}
 	return &cacheGenNumbers{}, f.err
 }