Add support for caching bloomfilters

table/table.go

@@ -81,7 +81,7 @@ type TableInterface interface {
 	DoesNotHave(hash uint64) bool
 }
 
-// Table represents a loaded table file with the info we have about it
+// Table represents a loaded table file with the info we have about it.
 type Table struct {
 	sync.Mutex
 
@@ -97,10 +97,11 @@ type Table struct {
 	smallest, biggest []byte // Smallest and largest keys (with timestamps).
 	id                uint64 // file id, part of filename
 
-	bf       *z.Bloom
 	Checksum []byte
 	// Stores the total size of key-values stored in this table (including the size on vlog).
 	estimatedSize uint64
+	indexStart    int
+	indexLen      int
 
 	IsInmemory bool // Set to true if the table is on level 0 and opened in memory.
 	opt        *Options
@@ -146,6 +147,13 @@ func (t *Table) DecrRef() error {
 		if err := os.Remove(filename); err != nil {
 			return err
 		}
+		// Delete all blocks from the cache.
+		for i, _ := range t.blockIndex {
+			t.opt.Cache.Del(t.blockCacheKey(i))
+		}
+		// Delete bloom filter from the cache.
+		t.opt.Cache.Del(t.bfKey())
+
 	}
 	return nil
 }
@@ -336,10 +344,12 @@ func (t *Table) readIndex() error {
 	// Read index size from the footer.
 	readPos -= 4
 	buf = t.readNoFail(readPos, 4)
-	indexLen := int(y.BytesToU32(buf))
+	t.indexLen = int(y.BytesToU32(buf))
+
 	// Read index.
-	readPos -= indexLen
-	data := t.readNoFail(readPos, indexLen)
+	readPos -= t.indexLen
+	t.indexStart = readPos
+	data := t.readNoFail(readPos, t.indexLen)
 
 	if err := y.VerifyChecksum(data, expectedChk); err != nil {
 		return y.Wrapf(err, "failed to verify checksum for table: %s", t.Filename())
@@ -358,14 +368,23 @@ func (t *Table) readIndex() error {
 	y.Check(err)
 
 	t.estimatedSize = index.EstimatedSize
-	if t.bf, err = z.JSONUnmarshal(index.BloomFilter); err != nil {
+	t.blockIndex = index.Offsets
+
+	var bf *z.Bloom
+	if bf, err = z.JSONUnmarshal(index.BloomFilter); err != nil {
 		return y.Wrapf(err, "failed to unmarshal bloom filter for the table %d in Table.readIndex",
 			t.id)
 	}
-	t.blockIndex = index.Offsets
+
+	t.opt.Cache.Set(t.bfKey(), bf, int64(len(index.BloomFilter)))
 	return nil
 }
 
+// Returns the cache key for the bloom filter.
+func (t *Table) bfKey() uint64 {
+	return math.MaxUint64 - t.id
+}
+
 func (t *Table) block(idx int) (*block, error) {
 	y.AssertTruef(idx >= 0, "idx=%d", idx)
 	if idx >= len(t.blockIndex) {
@@ -470,7 +489,31 @@ func (t *Table) ID() uint64 { return t.id }
 
 // DoesNotHave returns true if (but not "only if") the table does not have the key hash.
 // It does a bloom filter lookup.
-func (t *Table) DoesNotHave(hash uint64) bool { return !t.bf.Has(hash) }
+func (t *Table) DoesNotHave(hash uint64) bool {
+	var bf *z.Bloom
+	// Found bloomfilter in the cache.
+	if b, ok := t.opt.Cache.Get(t.bfKey()); b != nil && ok {
+		bf = b.(*z.Bloom)
+		return !bf.Has(hash)
+	}
+	// Read bloom filter from the SST.
+	data := t.readNoFail(t.indexStart, t.indexLen)
+	index := pb.TableIndex{}
+	var err error
+	// Decrypt the table index if it is encrypted.
+	if t.shouldDecrypt() {
+		data, err = t.decrypt(data)
+		y.Check(err)
+
+	}
+	y.Check(proto.Unmarshal(data, &index))
+
+	bf, err = z.JSONUnmarshal(index.BloomFilter)
+	y.Check(err)
+
+	t.opt.Cache.Set(t.bfKey(), bf, int64(len(index.BloomFilter)))
+	return !bf.Has(hash)
+}
 
 // VerifyChecksum verifies checksum for all blocks of table. This function is called by
 // OpenTable() function. This function is also called inside levelsController.VerifyChecksum().

table/table_test.go

@@ -21,6 +21,7 @@ import (
 	"crypto/sha256"
 	"fmt"
 	"hash/crc32"
+	"io/ioutil"
 	"math/rand"
 	"os"
 	"sort"
@@ -40,6 +41,8 @@ const (
 	MB = KB * 1024
 )
 
+var fileID uint64
+
 func key(prefix string, i int) string {
 	return prefix + fmt.Sprintf("%04d", i)
 }
@@ -77,13 +80,14 @@ func buildTable(t *testing.T, keyValues [][]string, opts Options) *os.File {
 	defer b.Close()
 	// TODO: Add test for file garbage collection here. No files should be left after the tests here.
 
-	filename := fmt.Sprintf("%s%s%d.sst", os.TempDir(), string(os.PathSeparator), rand.Int63())
+	dir, err := ioutil.TempDir("", "badger-test")
+	require.NoError(t, err)
+	defer os.RemoveAll(dir)
+
+	fileID++
+	filename := fmt.Sprintf("%s%s%d.sst", dir, string(os.PathSeparator), fileID)
 	f, err := y.CreateSyncedFile(filename, true)
-	if t != nil {
-		require.NoError(t, err)
-	} else {
-		y.Check(err)
-	}
+	require.NoError(t, err)
 
 	sort.Slice(keyValues, func(i, j int) bool {
 		return keyValues[i][0] < keyValues[j][0]
@@ -940,3 +944,48 @@ func TestOpenKVSize(t *testing.T) {
 	var entrySize uint64 = 15 /* DiffKey len */ + 4 /* Header Size */ + 4 /* Encoded vp */
 	require.Equal(t, entrySize, table.EstimatedSize())
 }
+
+func TestBloomFilterCache(t *testing.T) {
+	cache, err := ristretto.NewCache(&ristretto.Config{
+		NumCounters: 10,       // number of keys to track frequency of (10).
+		MaxCost:     10 << 20, // maximum cost of cache (10 MB).
+		BufferItems: 64,       // number of keys per Get buffer.
+		Metrics:     true,
+	})
+	require.NoError(t, err)
+	defer cache.Close()
+
+	opts := getTestTableOptions()
+	opts.Cache = cache
+
+	data := buildTestTable(t, "foo", 1, opts)
+
+	require.Equal(t, uint64(0), cache.Metrics.KeysAdded())
+	require.Equal(t, uint64(0), cache.Metrics.Hits())
+
+	table, err := OpenTable(data, opts)
+	require.NoError(t, err)
+	defer table.DecrRef()
+
+	// Wait for cache.set to be applied.
+	time.Sleep(time.Second)
+
+	require.Equal(t, uint64(0), cache.Metrics.Hits())
+
+	// The bloom filter lookup should hit the cache.
+	require.True(t, table.DoesNotHave(123))
+	require.Equal(t, uint64(1), cache.Metrics.Hits())
+
+	// Do it once again for fun!
+	require.True(t, table.DoesNotHave(123))
+	require.Equal(t, uint64(2), cache.Metrics.Hits())
+
+	// Drop the bloom filter from cache.
+	cache.Del(table.bfKey())
+
+	misses := cache.Metrics.Misses()
+	require.True(t, table.DoesNotHave(123))
+	require.Equal(t, uint64(2), cache.Metrics.Hits())
+	// The bf was removed from the cache. We should get a cache miss.
+	require.Equal(t, misses+1, cache.Metrics.Misses())
+}