From 5fdf78424309a41a4a06c480516a6d0c0ac27094 Mon Sep 17 00:00:00 2001
From: Oleg Zaytsev <mail@olegzaytsev.com>
Date: Wed, 27 Sep 2023 17:23:03 +0200
Subject: [PATCH] Refactor PostingsForMatcherCache promise

Extract promise payload as a struct, to make size calculation easier.

Signed-off-by: Oleg Zaytsev <mail@olegzaytsev.com>
---
 tsdb/postings_for_matchers_cache.go      | 44 +++++++++++++-----------
 tsdb/postings_for_matchers_cache_test.go |  4 +--
 2 files changed, 25 insertions(+), 23 deletions(-)

diff --git a/tsdb/postings_for_matchers_cache.go b/tsdb/postings_for_matchers_cache.go
index e0c32923b0..aea96ca3dc 100644
--- a/tsdb/postings_for_matchers_cache.go
+++ b/tsdb/postings_for_matchers_cache.go
@@ -81,41 +81,43 @@ func (c *PostingsForMatchersCache) PostingsForMatchers(ctx context.Context, ix I
 	return c.postingsForMatchersPromise(ctx, ix, ms)()
 }
 
-func (c *PostingsForMatchersCache) postingsForMatchersPromise(ctx context.Context, ix IndexPostingsReader, ms []*labels.Matcher) func() (index.Postings, error) {
-	var (
-		wg       sync.WaitGroup
-		cloner   *index.PostingsCloner
-		outerErr error
-	)
-	wg.Add(1)
+type postingsForMatcherPromise struct {
+	sync.WaitGroup
 
-	promise := func() (index.Postings, error) {
-		wg.Wait()
-		if outerErr != nil {
-			return nil, outerErr
-		}
-		return cloner.Clone(), nil
+	cloner *index.PostingsCloner
+	err    error
+}
+
+func (p *postingsForMatcherPromise) result() (index.Postings, error) {
+	p.Wait()
+	if p.err != nil {
+		return nil, p.err
 	}
+	return p.cloner.Clone(), nil
+}
+
+func (c *PostingsForMatchersCache) postingsForMatchersPromise(ctx context.Context, ix IndexPostingsReader, ms []*labels.Matcher) func() (index.Postings, error) {
+	promise := new(postingsForMatcherPromise)
+	promise.Add(1)
 
 	key := matchersKey(ms)
 	oldPromise, loaded := c.calls.LoadOrStore(key, promise)
 	if loaded {
-		return oldPromise.(func() (index.Postings, error))
+		promise = oldPromise.(*postingsForMatcherPromise)
+		return promise.result
 	}
-	defer wg.Done()
+	defer promise.Done()
 
 	if postings, err := c.postingsForMatchers(ctx, ix, ms...); err != nil {
-		outerErr = err
+		promise.err = err
 	} else {
-		cloner = index.NewPostingsCloner(postings)
+		promise.cloner = index.NewPostingsCloner(postings)
 	}
 
-	// Estimate the size of the cache entry, in bytes. We use max() because
-	// size.Of() returns -1 if the value is nil.
-	sizeBytes := int64(len(key)) + max(0, int64(size.Of(outerErr))) + max(0, int64(size.Of(cloner)))
+	sizeBytes := int64(len(key) + size.Of(promise))
 
 	c.created(key, c.timeNow(), sizeBytes)
-	return promise
+	return promise.result
 }
 
 type postingsForMatchersCachedCall struct {
diff --git a/tsdb/postings_for_matchers_cache_test.go b/tsdb/postings_for_matchers_cache_test.go
index b9a13effdc..e6ee36e96f 100644
--- a/tsdb/postings_for_matchers_cache_test.go
+++ b/tsdb/postings_for_matchers_cache_test.go
@@ -267,7 +267,7 @@ func TestPostingsForMatchersCache(t *testing.T) {
 	t.Run("cached value is evicted because cache exceeds max bytes", func(t *testing.T) {
 		const (
 			maxItems         = 100 // Never hit it.
-			maxBytes         = 1000
+			maxBytes         = 1100
 			numMatchers      = 5
 			postingsListSize = 30 // 8 bytes per posting ref, so 30 x 8 = 240 bytes.
 		)
@@ -313,7 +313,7 @@ func TestPostingsForMatchersCache(t *testing.T) {
 
 		// At this point we expect that the postings have been computed only once for the 3 matchers.
 		for i := 0; i < 3; i++ {
-			assert.Equal(t, 1, callsPerMatchers[matchersKey(matchersLists[i])])
+			assert.Equalf(t, 1, callsPerMatchers[matchersKey(matchersLists[i])], "matcher %d", i)
 		}
 
 		// Call PostingsForMatchers() for a 4th matcher. We expect this will evict the oldest cached entry.