From 6b84d685d56f680650ffd710b3b21e4be5ee4f51 Mon Sep 17 00:00:00 2001
From: Ben Johnson <benbjohnson@yahoo.com>
Date: Tue, 15 Feb 2022 08:25:45 -0700
Subject: [PATCH] Periodically invalidate rank cache during bulk add

This commit changes `RankCache.BulkAdd()` so that entries are
limited to an upper bound of 2x `maxEntries`. When this bound
is exceeded then the cache is automatically recalculated.
---
 cache.go      |  8 ++++++++
 cache_test.go | 12 ++++++++++++
 2 files changed, 20 insertions(+)

diff --git a/cache.go b/cache.go
index 4dd5d5e4c..f558a28cf 100644
--- a/cache.go
+++ b/cache.go
@@ -194,6 +194,14 @@ func (c *rankCache) BulkAdd(id uint64, n uint64) {
 	}
 
 	c.entries[id] = n
+
+	// FB-1206: Periodically invalidate the cache when we are bulk loading
+	// as this can take up an upbounded amount of memory. This is especially
+	// true when restoring shards as all rows will be added.
+	if len(c.entries) > int(2*c.maxEntries) {
+		c.stats.Count(MetricRecalculateCache, 1, 1.0)
+		c.recalculate()
+	}
 }
 
 // Get returns a count for a given id.
diff --git a/cache_test.go b/cache_test.go
index 1ff8d0fff..0da077f1c 100644
--- a/cache_test.go
+++ b/cache_test.go
@@ -70,3 +70,15 @@ func TestCache_Rank_Dirty(t *testing.T) {
 		t.Fatalf("wrote %v but got %v", expect, got)
 	}
 }
+
+func TestCache_Rank_BulkAdd(t *testing.T) {
+	const cacheSize = 10
+	cache := pilosa.NewRankCache(uint32(cacheSize))
+
+	for i := uint64(0); i < 1000; i++ {
+		cache.BulkAdd(i, i)
+		if n := cache.Len(); n > cacheSize*2 {
+			t.Fatalf("entry count exceed 2x cache size: %d", n)
+		}
+	}
+}