Add configurable bloom filters

modules/ingester/ingester_test.go

@@ -259,7 +259,8 @@ func defaultIngester(t *testing.T, tmpDir string) (*Ingester, []*tempopb.Trace,
 			},
 			Block: &encoding.BlockConfig{
 				IndexDownsampleBytes: 2,
-				BloomFP:              .01,
+				BloomFP:              0.01,
+				BloomFilterShardSize: 100_000,
 				Encoding:             backend.EncLZ4_1M,
 				IndexPageSizeBytes:   1000,
 			},

modules/ingester/instance_test.go

@@ -520,7 +520,8 @@ func defaultInstance(t require.TestingT, tmpDir string) *instance {
 			},
 			Block: &encoding.BlockConfig{
 				IndexDownsampleBytes: 2,
-				BloomFP:              .01,
+				BloomFP:              0.01,
+				BloomFilterShardSize: 100_000,
 				Encoding:             backend.EncLZ4_1M,
 				IndexPageSizeBytes:   1000,
 			},

modules/querier/querier_test.go

@@ -57,7 +57,8 @@ func TestReturnAllHits(t *testing.T) {
 		Block: &encoding.BlockConfig{
 			Encoding:             backend.EncNone,
 			IndexDownsampleBytes: 10,
-			BloomFP:              .05,
+			BloomFP:              0.01,
+			BloomFilterShardSize: 100_000,
 			IndexPageSizeBytes:   1000,
 		},
 		WAL: &wal.Config{

modules/storage/config.go

@@ -5,6 +5,7 @@ import (
 	"time"
 
 	cortex_cache "github.com/cortexproject/cortex/pkg/chunk/cache"
+
 	"github.com/grafana/tempo/pkg/util"
 	"github.com/grafana/tempo/tempodb"
 	"github.com/grafana/tempo/tempodb/backend"
@@ -37,7 +38,8 @@ func (cfg *Config) RegisterFlagsAndApplyDefaults(prefix string, f *flag.FlagSet)
 	cfg.Trace.WAL.Encoding = backend.EncNone
 
 	cfg.Trace.Block = &encoding.BlockConfig{}
-	f.Float64Var(&cfg.Trace.Block.BloomFP, util.PrefixConfig(prefix, "trace.block.bloom-filter-false-positive"), .05, "Bloom False Positive.")
+	f.Float64Var(&cfg.Trace.Block.BloomFP, util.PrefixConfig(prefix, "trace.block.bloom-filter-false-positive"), .05, "Bloom Filter False Positive.")
+	f.IntVar(&cfg.Trace.Block.BloomFilterShardSize, util.PrefixConfig(prefix, "trace.block.bloom-filter-shard-size"), 250*1024, "Bloom Filter Shard Size.")
 	f.IntVar(&cfg.Trace.Block.IndexDownsampleBytes, util.PrefixConfig(prefix, "trace.block.index-downsample-bytes"), 1024*1024, "Number of bytes (before compression) per index record.")
 	f.IntVar(&cfg.Trace.Block.IndexPageSizeBytes, util.PrefixConfig(prefix, "trace.block.index-page-size-bytes"), 250*1024, "Number of bytes per index page.")
 	cfg.Trace.Block.Encoding = backend.EncZstd

tempodb/backend/block_meta.go

@@ -28,6 +28,7 @@ type BlockMeta struct {
 	IndexPageSize   uint32    `json:"indexPageSize"`   // Size of each index page in bytes
 	TotalRecords    uint32    `json:"totalRecords"`    // Total Records stored in the index file
 	DataEncoding    string    `json:"dataEncoding"`    // DataEncoding is a string provided externally, but tracked by tempodb that indicates the way the bytes are encoded
+	BloomShardCount uint16    `json:"bloomShards"`     // Number of bloom filter shards
 }
 
 func NewBlockMeta(tenantID string, blockID uuid.UUID, version string, encoding Encoding, dataEncoding string) *BlockMeta {

tempodb/backend/block_meta_test.go

@@ -2,6 +2,7 @@ package backend
 
 import (
 	"bytes"
+	"encoding/json"
 	"math/rand"
 	"testing"
 
@@ -41,3 +42,29 @@ func TestBlockMeta(t *testing.T) {
 	assert.Equal(t, 1, bytes.Compare(b.MaxID, b.MinID))
 	assert.Equal(t, 2, b.TotalObjects)
 }
+
+func TestBlockMetaParsing(t *testing.T) {
+	inputJSON := `
+{
+    "format": "v0",
+    "blockID": "00000000-0000-0000-0000-000000000000",
+    "minID": "AAAAAAAAAAAAOO0z0LnnHg==",
+    "maxID": "AAAAAAAAAAD/o61w2bYIDg==",
+    "tenantID": "single-tenant",
+    "startTime": "2021-01-01T00:00:00.0000000Z",
+    "endTime": "2021-01-02T00:00:00.0000000Z",
+    "totalObjects": 10,
+    "size": 12345,
+    "compactionLevel": 0,
+    "encoding": "zstd",
+    "indexPageSize": 250000,
+    "totalRecords": 124356,
+    "dataEncoding": "",
+    "bloomShards": 244
+}
+`
+
+	blockMeta := BlockMeta{}
+	err := json.Unmarshal([]byte(inputJSON), &blockMeta)
+	assert.NoError(t, err, "expected to be able to unmarshal from JSON")
+}

tempodb/backend/local/local_test.go

@@ -10,9 +10,10 @@ import (
 	"testing"
 
 	"github.com/google/uuid"
+	"github.com/stretchr/testify/assert"
+
 	"github.com/grafana/tempo/tempodb/backend"
 	"github.com/grafana/tempo/tempodb/encoding/common"
-	"github.com/stretchr/testify/assert"
 )
 
 const objectName = "test"
@@ -107,7 +108,7 @@ func TestCompaction(t *testing.T) {
 		BlockID: blockID,
 	}
 
-	shardNum := common.GetShardNum()
+	shardNum := common.ValidateShardCount(int(fakeMeta.BloomShardCount))
 	fakeBloom := make([][]byte, shardNum)
 	fakeIndex := make([]byte, 20)
 	fakeTraces := make([]byte, 200)

tempodb/compactor_bookmark_test.go

@@ -35,6 +35,7 @@ func TestCurrentClear(t *testing.T) {
 		Block: &encoding.BlockConfig{
 			IndexDownsampleBytes: 17,
 			BloomFP:              .01,
+			BloomFilterShardSize: 100_000,
 			Encoding:             backend.EncGZIP,
 			IndexPageSizeBytes:   1000,
 		},

tempodb/compactor_test.go

@@ -64,6 +64,7 @@ func TestCompaction(t *testing.T) {
 		Block: &encoding.BlockConfig{
 			IndexDownsampleBytes: 11,
 			BloomFP:              .01,
+			BloomFilterShardSize: 100_000,
 			Encoding:             backend.EncLZ4_4M,
 			IndexPageSizeBytes:   1000,
 		},
@@ -192,6 +193,7 @@ func TestSameIDCompaction(t *testing.T) {
 		Block: &encoding.BlockConfig{
 			IndexDownsampleBytes: 11,
 			BloomFP:              .01,
+			BloomFilterShardSize: 100_000,
 			Encoding:             backend.EncSnappy,
 			IndexPageSizeBytes:   1000,
 		},
@@ -279,6 +281,7 @@ func TestCompactionUpdatesBlocklist(t *testing.T) {
 		Block: &encoding.BlockConfig{
 			IndexDownsampleBytes: 11,
 			BloomFP:              .01,
+			BloomFilterShardSize: 100_000,
 			Encoding:             backend.EncNone,
 			IndexPageSizeBytes:   1000,
 		},
@@ -345,6 +348,7 @@ func TestCompactionMetrics(t *testing.T) {
 		Block: &encoding.BlockConfig{
 			IndexDownsampleBytes: 11,
 			BloomFP:              .01,
+			BloomFilterShardSize: 100_000,
 			Encoding:             backend.EncNone,
 			IndexPageSizeBytes:   1000,
 		},
@@ -415,6 +419,7 @@ func TestCompactionIteratesThroughTenants(t *testing.T) {
 		Block: &encoding.BlockConfig{
 			IndexDownsampleBytes: 11,
 			BloomFP:              .01,
+			BloomFilterShardSize: 100_000,
 			Encoding:             backend.EncLZ4_64k,
 			IndexPageSizeBytes:   1000,
 		},

tempodb/encoding/backend_block.go

@@ -48,7 +48,7 @@ func (b *BackendBlock) Find(ctx context.Context, id common.ID) ([]byte, error) {
 
 	span.SetTag("block", b.meta.BlockID.String())
 
-	shardKey := common.ShardKeyForTraceID(id)
+	shardKey := common.ShardKeyForTraceID(id, int(b.meta.BloomShardCount))
 	blockID := b.meta.BlockID
 	tenantID := b.meta.TenantID
 

tempodb/encoding/block.go

@@ -26,7 +26,7 @@ func bloomName(shard int) string {
 
 // writeBlockMeta writes the bloom filter, meta and index to the passed in backend.Writer
 func writeBlockMeta(ctx context.Context, w backend.Writer, meta *backend.BlockMeta, indexBytes []byte, b *common.ShardedBloomFilter) error {
-	blooms, err := b.WriteTo()
+	blooms, err := b.Write()
 	if err != nil {
 		return err
 	}
@@ -81,7 +81,7 @@ func CopyBlock(ctx context.Context, meta *backend.BlockMeta, src backend.Reader,
 	}
 
 	// Bloom
-	for i := 0; i < common.GetShardNum(); i++ {
+	for i := 0; i < common.ValidateShardCount(int(meta.BloomShardCount)); i++ {
 		err = copy(bloomName(i))
 		if err != nil {
 			return err

tempodb/encoding/common/bloom.go

@@ -3,40 +3,53 @@ package common
 import (
 	"bytes"
 
-	"github.com/grafana/tempo/pkg/util"
 	"github.com/willf/bloom"
+
+	"github.com/grafana/tempo/pkg/util"
 )
 
-const shardNum = 10
+const legacyShardCount = 10
 
 type ShardedBloomFilter struct {
 	blooms []*bloom.BloomFilter
 }
 
-func NewWithEstimates(n uint, fp float64) *ShardedBloomFilter {
-	b := &ShardedBloomFilter{
-		blooms: make([]*bloom.BloomFilter, shardNum),
+// NewBloom creates a ShardedBloomFilter
+func NewBloom(fp float64, shardSize, estimatedObjects uint) *ShardedBloomFilter {
+	// estimate the number of shards needed. an approximate value is enough
+	var shardCount uint
+	var kPerBloom uint
+	for {
+		shardCount++
+		var m, k uint
+		// m: number of bits in the filter
+		// k: number of hash functions
+		if m, k = bloom.EstimateParameters(estimatedObjects/shardCount, fp); m < shardSize {
+			kPerBloom = k
+			break
+		}
 	}
 
-	itemsPerBloom := n / shardNum
-	if itemsPerBloom == 0 {
-		itemsPerBloom = 1
+	b := &ShardedBloomFilter{
+		blooms: make([]*bloom.BloomFilter, shardCount),
 	}
-	for i := 0; i < shardNum; i++ {
-		b.blooms[i] = bloom.NewWithEstimates(itemsPerBloom, fp)
+
+	for i := 0; i < int(shardCount); i++ {
+		// New(m uint, k uint) creates a new Bloom filter with _m_ bits and _k_ hashing functions
+		b.blooms[i] = bloom.New(shardSize*8, kPerBloom)
 	}
 
 	return b
 }
 
 func (b *ShardedBloomFilter) Add(traceID []byte) {
-	shardKey := ShardKeyForTraceID(traceID)
+	shardKey := ShardKeyForTraceID(traceID, len(b.blooms))
 	b.blooms[shardKey].Add(traceID)
 }
 
-// WriteTo is a wrapper around bloom.WriteTo
-func (b *ShardedBloomFilter) WriteTo() ([][]byte, error) {
-	bloomBytes := make([][]byte, shardNum)
+// Write is a wrapper around bloom.WriteTo
+func (b *ShardedBloomFilter) Write() ([][]byte, error) {
+	bloomBytes := make([][]byte, len(b.blooms))
 	for i, f := range b.blooms {
 		bloomBuffer := &bytes.Buffer{}
 		_, err := f.WriteTo(bloomBuffer)
@@ -48,16 +61,24 @@ func (b *ShardedBloomFilter) WriteTo() ([][]byte, error) {
 	return bloomBytes, nil
 }
 
-func ShardKeyForTraceID(traceID []byte) int {
-	return int(util.TokenForTraceID(traceID)) % shardNum
+func (b *ShardedBloomFilter) GetShardCount() int {
+	return len(b.blooms)
 }
 
 // Test implements bloom.Test -> required only for testing
 func (b *ShardedBloomFilter) Test(traceID []byte) bool {
-	shardKey := ShardKeyForTraceID(traceID)
+	shardKey := ShardKeyForTraceID(traceID, len(b.blooms))
 	return b.blooms[shardKey].Test(traceID)
 }
 
-func GetShardNum() int {
-	return shardNum
+func ShardKeyForTraceID(traceID []byte, shardCount int) int {
+	return int(util.TokenForTraceID(traceID)) % ValidateShardCount(shardCount)
+}
+
+// For backward compatibility
+func ValidateShardCount(shardCount int) int {
+	if shardCount == 0 {
+		return legacyShardCount
+	}
+	return shardCount
 }