sql: add Bloom filters for ATXs and malicious identities

Author: ivan4th

go.mod

@@ -5,6 +5,7 @@ go 1.22.4
 require (
 	cloud.google.com/go/storage v1.43.0
 	github.com/ALTree/bigfloat v0.2.0
+	github.com/bits-and-blooms/bloom/v3 v3.7.0
 	github.com/chaos-mesh/chaos-mesh/api v0.0.0-20240913055630-bfe8736306b4
 	github.com/cosmos/btcutil v1.0.5
 	github.com/go-llsqlite/crawshaw v0.5.5
@@ -35,6 +36,7 @@ require (
 	github.com/prometheus/client_model v0.6.1
 	github.com/prometheus/common v0.59.1
 	github.com/quic-go/quic-go v0.46.0
+	github.com/rqlite/sql v0.0.0-20240312185922-ffac88a740bd
 	github.com/rs/cors v1.11.1
 	github.com/santhosh-tekuri/jsonschema/v5 v5.3.1
 	github.com/seehuhn/mt19937 v1.0.0
@@ -80,6 +82,7 @@ require (
 	github.com/anacrolix/sync v0.3.0 // indirect
 	github.com/benbjohnson/clock v1.3.5 // indirect
 	github.com/beorn7/perks v1.0.1 // indirect
+	github.com/bits-and-blooms/bitset v1.10.0 // indirect
 	github.com/c0mm4nd/go-ripemd v0.0.0-20200326052756-bd1759ad7d10 // indirect
 	github.com/cespare/xxhash/v2 v2.3.0 // indirect
 	github.com/containerd/cgroups v1.1.0 // indirect

go.sum

@@ -48,6 +48,10 @@ github.com/benbjohnson/clock v1.3.5/go.mod h1:J11/hYXuz8f4ySSvYwY0FKfm+ezbsZBKZx
 github.com/beorn7/perks v0.0.0-20180321164747-3a771d992973/go.mod h1:Dwedo/Wpr24TaqPxmxbtue+5NUziq4I4S80YR8gNf3Q=
 github.com/beorn7/perks v1.0.1 h1:VlbKKnNfV8bJzeqoa4cOKqO6bYr3WgKZxO8Z16+hsOM=
 github.com/beorn7/perks v1.0.1/go.mod h1:G2ZrVWU2WbWT9wwq4/hrbKbnv/1ERSJQ0ibhJ6rlkpw=
+github.com/bits-and-blooms/bitset v1.10.0 h1:ePXTeiPEazB5+opbv5fr8umg2R/1NlzgDsyepwsSr88=
+github.com/bits-and-blooms/bitset v1.10.0/go.mod h1:7hO7Gc7Pp1vODcmWvKMRA9BNmbv6a/7QIWpPxHddWR8=
+github.com/bits-and-blooms/bloom/v3 v3.7.0 h1:VfknkqV4xI+PsaDIsoHueyxVDZrfvMn56jeWUzvzdls=
+github.com/bits-and-blooms/bloom/v3 v3.7.0/go.mod h1:VKlUSvp0lFIYqxJjzdnSsZEw4iHb1kOL2tfHTgyJBHg=
 github.com/bradfitz/go-smtpd v0.0.0-20170404230938-deb6d6237625/go.mod h1:HYsPBTaaSFSlLx/70C2HPIMNZpVV8+vt/A+FMnYP11g=
 github.com/bradfitz/iter v0.0.0-20140124041915-454541ec3da2/go.mod h1:PyRFw1Lt2wKX4ZVSQ2mk+PeDa1rxyObEDlApuIsUKuo=
 github.com/bradfitz/iter v0.0.0-20190303215204-33e6a9893b0c h1:FUUopH4brHNO2kJoNN3pV+OBEYmgraLT/KHZrMM69r0=
@@ -152,6 +156,8 @@ github.com/go-openapi/swag v0.22.4/go.mod h1:UzaqsxGiab7freDnrUUra0MwWfN/q7tE4j+
 github.com/go-stack/stack v1.8.0/go.mod h1:v0f6uXyyMGvRgIKkXu+yp6POWl0qKG85gN/melR3HDY=
 github.com/go-task/slim-sprig/v3 v3.0.0 h1:sUs3vkvUymDpBKi3qH1YSqBQk9+9D/8M2mN1vB6EwHI=
 github.com/go-task/slim-sprig/v3 v3.0.0/go.mod h1:W848ghGpv3Qj3dhTPRyJypKRiqCdHZiAzKg9hl15HA8=
+github.com/go-test/deep v1.0.8 h1:TDsG77qcSprGbC6vTN8OuXp5g+J+b5Pcguhf7Zt61VM=
+github.com/go-test/deep v1.0.8/go.mod h1:5C2ZWiW0ErCdrYzpqxLbTX7MG14M9iiw8DgHncVwcsE=
 github.com/go-yaml/yaml v2.1.0+incompatible/go.mod h1:w2MrLa16VYP0jy6N7M5kHaCkaLENm+P+Tv+MfurjSw0=
 github.com/godbus/dbus/v5 v5.0.3/go.mod h1:xhWf0FNVPg57R7Z0UbKHbJfkEywrmjJnf7w5xrFpKfA=
 github.com/godbus/dbus/v5 v5.0.4/go.mod h1:xhWf0FNVPg57R7Z0UbKHbJfkEywrmjJnf7w5xrFpKfA=
@@ -551,6 +557,8 @@ github.com/robfig/cron/v3 v3.0.1/go.mod h1:eQICP3HwyT7UooqI/z+Ov+PtYAWygg1TEWWzG
 github.com/rogpeppe/go-internal v1.3.0/go.mod h1:M8bDsm7K2OlrFYOpmOWEs/qY81heoFRclV5y23lUDJ4=
 github.com/rogpeppe/go-internal v1.12.0 h1:exVL4IDcn6na9z1rAb56Vxr+CgyK3nn3O+epU5NdKM8=
 github.com/rogpeppe/go-internal v1.12.0/go.mod h1:E+RYuTGaKKdloAfM02xzb0FW3Paa99yedzYV+kq4uf4=
+github.com/rqlite/sql v0.0.0-20240312185922-ffac88a740bd h1:wW6BtayFoKaaDeIvXRE3SZVPOscSKlYD+X3bB749+zk=
+github.com/rqlite/sql v0.0.0-20240312185922-ffac88a740bd/go.mod h1:ib9zVtNgRKiGuoMyUqqL5aNpk+r+++YlyiVIkclVqPg=
 github.com/rs/cors v1.11.1 h1:eU3gRzXLRK57F5rKMGMZURNdIG4EoAmX8k94r9wXWHA=
 github.com/rs/cors v1.11.1/go.mod h1:XyqrcTp5zjWr1wsJ8PIRZssZ8b/WMcMf71DJnit4EMU=
 github.com/russross/blackfriday v1.5.2/go.mod h1:JO/DiYxRf+HjHt06OyowR9PTA263kcR/rfWxYHBV53g=
@@ -656,6 +664,8 @@ github.com/syndtr/goleveldb v1.0.1-0.20210819022825-2ae1ddf74ef7 h1:epCh84lMvA70
 github.com/syndtr/goleveldb v1.0.1-0.20210819022825-2ae1ddf74ef7/go.mod h1:q4W45IWZaF22tdD+VEXcAWRA037jwmWEB5VWYORlTpc=
 github.com/tarm/serial v0.0.0-20180830185346-98f6abe2eb07/go.mod h1:kDXzergiv9cbyO7IOYJZWg1U88JhDg3PB6klq9Hg2pA=
 github.com/tinylib/msgp v1.0.2/go.mod h1:+d+yLhGm8mzTaHzB+wgMYrodPfmZrzkirds8fDWklFE=
+github.com/twmb/murmur3 v1.1.6 h1:mqrRot1BRxm+Yct+vavLMou2/iJt0tNVTTC0QoIjaZg=
+github.com/twmb/murmur3 v1.1.6/go.mod h1:Qq/R7NUyOfr65zD+6Q5IHKsJLwP7exErjN6lyyq3OSQ=
 github.com/urfave/cli v1.22.2/go.mod h1:Gos4lmkARVdJ6EkW0WaNv/tZAAMe9V7XWyB60NtXRu0=
 github.com/urfave/cli v1.22.10/go.mod h1:Gos4lmkARVdJ6EkW0WaNv/tZAAMe9V7XWyB60NtXRu0=
 github.com/viant/assertly v0.4.8/go.mod h1:aGifi++jvCrUaklKEKT0BU95igDNaqkvz+49uaYMPRU=

node/node.go

@@ -74,6 +74,7 @@ import (
 	"github.com/spacemeshos/go-spacemesh/sql"
 	"github.com/spacemeshos/go-spacemesh/sql/activesets"
 	"github.com/spacemeshos/go-spacemesh/sql/atxs"
+	"github.com/spacemeshos/go-spacemesh/sql/identities"
 	"github.com/spacemeshos/go-spacemesh/sql/layers"
 	"github.com/spacemeshos/go-spacemesh/sql/localsql"
 	localmigrations "github.com/spacemeshos/go-spacemesh/sql/localsql/migrations"
@@ -1979,6 +1980,15 @@ func (app *App) setupDBs(ctx context.Context, lg log.Log) error {
 	}
 	{
 		warmupLog := app.log.Zap().Named("warmup")
+		app.log.Info("loading Bloom filters")
+		_, err := atxs.LoadBloomFilter(app.db, warmupLog)
+		if err != nil {
+			return fmt.Errorf("loading ATX Bloom filter: %w", err)
+		}
+		_, err = identities.LoadBloomFilter(app.db, warmupLog)
+		if err != nil {
+			return fmt.Errorf("loading malicious identity Bloom filter: %w", err)
+		}
 		app.log.Info("starting cache warmup")
 		applied, err := layers.GetLastApplied(app.db)
 		if err != nil {

sql/atxs/atxs.go

@@ -2,10 +2,12 @@ package atxs
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"time"
 
 	sqlite "github.com/go-llsqlite/crawshaw"
+	"go.uber.org/zap"
 
 	"github.com/spacemeshos/go-spacemesh/common/types"
 	"github.com/spacemeshos/go-spacemesh/sql"
@@ -15,6 +17,11 @@ import (
 const (
 	CacheKindEpochATXs sql.QueryCacheKind = "epoch-atxs"
 	CacheKindATXBlob   sql.QueryCacheKind = "atx-blob"
+	// Bloom filter size is < 115 MiB while below 100M ATXs.
+	// TODO: adjust Bloom filter settings after ATX merge & checkpointing.
+	BloomFilterFalsePositiveRate = 0.01
+	BloomFilterMinSize           = 100_000_000
+	BloomFilterExtraCoef         = 1.2
 )
 
 // Query to retrieve ATXs.
@@ -120,8 +127,28 @@ func GetByEpochAndNodeID(
 	return id, nil
 }
 
+// LoadBloomFilter intializes and loads the bloom filter for ATXs.
+func LoadBloomFilter(db sql.StateDatabase, logger *zap.Logger) (*sql.DBBloomFilter, error) {
+	bf := sql.NewDBBloomFilter(
+		"atxs", "select id from atxs", "id",
+		BloomFilterMinSize, BloomFilterExtraCoef, BloomFilterFalsePositiveRate)
+	if err := bf.Load(db, logger); err != nil {
+		return nil, fmt.Errorf("load bloom filter: %w", err)
+	}
+	db.AddSet(bf)
+	return bf, nil
+}
+
 // Has checks if an ATX exists by a given ATX ID.
+// It tries to do so using Bloom filter first, and falls back to a direct query if the filter is not available.
 func Has(db sql.Executor, id types.ATXID) (bool, error) {
+	has, err := sql.Contains(db, "atxs", id[:])
+	if err == nil {
+		return has, nil
+	} else if !errors.Is(err, sql.ErrNoSet) {
+		return false, fmt.Errorf("check if have id %s: %w", id, err)
+	}
+
 	rows, err := db.Exec("select 1 from atxs where id = ?1;",
 		func(stmt *sql.Statement) {
 			stmt.BindBytes(1, id.Bytes())
@@ -486,7 +513,12 @@ func Add(db sql.Executor, atx *types.ActivationTx, blob types.AtxBlob) error {
 		return fmt.Errorf("insert ATX ID %v: %w", atx.ID(), err)
 	}
 
-	return AddBlob(db, atx.ID(), blob.Blob, blob.Version)
+	if err := AddBlob(db, atx.ID(), blob.Blob, blob.Version); err != nil {
+		return err
+	}
+
+	sql.AddToSet(db, "atxs", atx.ID().Bytes())
+	return err
 }
 
 func AddBlob(db sql.Executor, id types.ATXID, blob []byte, version types.AtxVersion) error {