TBS: Replace badger with pebble

go.mod

@@ -5,7 +5,7 @@ go 1.23.0
 require (
 	github.com/KimMachineGun/automemlimit v0.7.0
 	github.com/cespare/xxhash/v2 v2.3.0
-	github.com/dgraph-io/badger/v2 v2.2007.4
+	github.com/cockroachdb/pebble/v2 v2.0.2
 	github.com/dustin/go-humanize v1.0.1
 	github.com/elastic/apm-aggregation v1.2.0
 	github.com/elastic/apm-data v1.16.0
@@ -60,20 +60,18 @@ require (
 require (
 	github.com/DataDog/zstd v1.5.6 // indirect
 	github.com/Microsoft/go-winio v0.6.2 // indirect
-	github.com/OneOfOne/xxhash v1.2.8 // indirect
 	github.com/armon/go-radix v1.0.0 // indirect
 	github.com/axiomhq/hyperloglog v0.2.0 // indirect
 	github.com/beorn7/perks v1.0.1 // indirect
-	github.com/cespare/xxhash v1.1.0 // indirect
+	github.com/cockroachdb/crlib v0.0.0-20241015224233-894974b3ad94 // indirect
 	github.com/cockroachdb/errors v1.11.3 // indirect
 	github.com/cockroachdb/fifo v0.0.0-20240816210425-c5d0cb0b6fc0 // indirect
 	github.com/cockroachdb/logtags v0.0.0-20230118201751-21c54148d20b // indirect
 	github.com/cockroachdb/pebble v1.1.2 // indirect
 	github.com/cockroachdb/redact v1.1.5 // indirect
+	github.com/cockroachdb/swiss v0.0.0-20240612210725-f4de07ae6964 // indirect
 	github.com/cockroachdb/tokenbucket v0.0.0-20230807174530-cc333fc44b06 // indirect
 	github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc // indirect
-	github.com/dgraph-io/ristretto v0.2.0 // indirect
-	github.com/dgryski/go-farm v0.0.0-20200201041132-a6ae2369ad13 // indirect
 	github.com/dgryski/go-metro v0.0.0-20211217172704-adc40b04c140 // indirect
 	github.com/dlclark/regexp2 v1.8.1 // indirect
 	github.com/docker/go-connections v0.5.0 // indirect
@@ -97,7 +95,7 @@ require (
 	github.com/go-logr/stdr v1.2.2 // indirect
 	github.com/go-ole/go-ole v1.2.6 // indirect
 	github.com/golang/protobuf v1.5.4 // indirect
-	github.com/golang/snappy v0.0.4 // indirect
+	github.com/golang/snappy v0.0.5-0.20231225225746-43d5d4cd4e0e // indirect
 	github.com/gomodule/redigo v1.8.9 // indirect
 	github.com/google/uuid v1.6.0 // indirect
 	github.com/grpc-ecosystem/go-grpc-middleware v1.3.0 // indirect

internal/beater/config/config_test.go

@@ -362,7 +362,6 @@ func TestUnpackConfig(t *testing.T) {
 						ESConfig:              elasticsearch.DefaultConfig(),
 						Interval:              1 * time.Minute,
 						IngestRateDecayFactor: 0.25,
-						StorageGCInterval:     5 * time.Minute,
 						StorageLimit:          "3GB",
 						StorageLimitParsed:    3000000000,
 						TTL:                   30 * time.Minute,
@@ -494,7 +493,6 @@ func TestUnpackConfig(t *testing.T) {
 						ESConfig:              elasticsearch.DefaultConfig(),
 						Interval:              2 * time.Minute,
 						IngestRateDecayFactor: 1.0,
-						StorageGCInterval:     5 * time.Minute,
 						StorageLimit:          "1GB",
 						StorageLimitParsed:    1000000000,
 						TTL:                   30 * time.Minute,

internal/beater/config/sampling.go

@@ -48,7 +48,6 @@ type TailSamplingConfig struct {
 	ESConfig              *elasticsearch.Config `config:"elasticsearch"`
 	Interval              time.Duration         `config:"interval" validate:"min=1s"`
 	IngestRateDecayFactor float64               `config:"ingest_rate_decay" validate:"min=0, max=1"`
-	StorageGCInterval     time.Duration         `config:"storage_gc_interval" validate:"min=1s"`
 	TTL                   time.Duration         `config:"ttl" validate:"min=1s"`
 	StorageLimit          string                `config:"storage_limit"`
 	StorageLimitParsed    uint64
@@ -151,7 +150,6 @@ func defaultTailSamplingConfig() TailSamplingConfig {
 		ESConfig:              elasticsearch.DefaultConfig(),
 		Interval:              1 * time.Minute,
 		IngestRateDecayFactor: 0.25,
-		StorageGCInterval:     5 * time.Minute,
 		TTL:                   30 * time.Minute,
 		StorageLimit:          "3GB",
 		DiscardOnWriteFailure: false,

systemtest/benchtest/expvar/metrics.go

@@ -167,7 +167,7 @@ func (c *Collector) accumulate(e expvar) {
 	c.processMetric(HeapAlloc, int64(e.HeapAlloc))
 	c.processMetric(HeapObjects, int64(e.HeapObjects))
 	c.processMetric(TBSLsmSize, e.TBSLsmSize)
-	c.processMetric(TBSVlogSize, e.TBSVlogSize)
+	c.processMetric(TBSVlogSize, e.TBSVlogSize) // FIXME: update to use new metrics
 }
 
 func (c *Collector) processMetric(m Metric, val int64) {

x-pack/apm-server/main.go

@@ -40,12 +40,9 @@ var (
 	// will hopefully disappear in the future, when agents no longer send unsampled transactions.
 	samplingMonitoringRegistry = monitoring.Default.GetRegistry("apm-server.sampling")
 
-	// badgerDB holds the badger database to use when tail-based sampling is configured.
-	badgerMu sync.Mutex
-	badgerDB *eventstorage.StorageManager
-
-	storageMu sync.Mutex
-	storage   *eventstorage.ManagedReadWriter
+	// db holds the database to use when tail-based sampling is configured.
+	dbMu sync.Mutex
+	db   *eventstorage.StorageManager
 
 	// samplerUUID is a UUID used to identify sampled trace ID documents
 	// published by this process.
@@ -117,11 +114,10 @@ func newTailSamplingProcessor(args beater.ServerParams) (*sampling.Processor, er
 	}
 
 	storageDir := paths.Resolve(paths.Data, tailSamplingStorageDir)
-	badgerDB, err = getBadgerDB(storageDir)
+	db, err := getDB(storageDir)
 	if err != nil {
-		return nil, fmt.Errorf("failed to get Badger database: %w", err)
+		return nil, fmt.Errorf("failed to get TBS database: %w", err)
 	}
-	readWriter := getStorage(badgerDB)
 
 	policies := make([]sampling.Policy, len(tailSamplingConfig.Policies))
 	for i, in := range tailSamplingConfig.Policies {
@@ -155,37 +151,25 @@ func newTailSamplingProcessor(args beater.ServerParams) (*sampling.Processor, er
 			UUID: samplerUUID.String(),
 		},
 		StorageConfig: sampling.StorageConfig{
-			DB:                    badgerDB,
-			Storage:               readWriter,
-			StorageDir:            storageDir,
-			StorageGCInterval:     tailSamplingConfig.StorageGCInterval,
+			DB:                    db,
 			StorageLimit:          tailSamplingConfig.StorageLimitParsed,
 			TTL:                   tailSamplingConfig.TTL,
 			DiscardOnWriteFailure: tailSamplingConfig.DiscardOnWriteFailure,
 		},
 	})
 }
 
-func getBadgerDB(storageDir string) (*eventstorage.StorageManager, error) {
-	badgerMu.Lock()
-	defer badgerMu.Unlock()
-	if badgerDB == nil {
+func getDB(storageDir string) (*eventstorage.StorageManager, error) {
+	dbMu.Lock()
+	defer dbMu.Unlock()
+	if db == nil {
 		sm, err := eventstorage.NewStorageManager(storageDir)
 		if err != nil {
 			return nil, err
 		}
-		badgerDB = sm
-	}
-	return badgerDB, nil
-}
-
-func getStorage(sm *eventstorage.StorageManager) *eventstorage.ManagedReadWriter {
-	storageMu.Lock()
-	defer storageMu.Unlock()
-	if storage == nil {
-		storage = sm.NewReadWriter()
+		db = sm
 	}
-	return storage
+	return db, nil
 }
 
 // runServerWithProcessors runs the APM Server and the given list of processors.
@@ -249,19 +233,19 @@ func wrapServer(args beater.ServerParams, runServer beater.RunServerFunc) (beate
 	return args, wrappedRunServer, nil
 }
 
-// closeBadger is called at process exit time to close the badger.DB opened
+// closeDB is called at process exit time to close the StorageManager opened
 // by the tail-based sampling processor constructor, if any. This is never
-// called concurrently with opening badger.DB/accessing the badgerDB global,
-// so it does not need to hold badgerMu.
-func closeBadger() error {
-	if badgerDB != nil {
-		return badgerDB.Close()
+// called concurrently with opening DB/accessing the db global,
+// so it does not need to hold dbMu.
+func closeDB() error {
+	if db != nil {
+		return db.Close()
 	}
 	return nil
 }
 
 func cleanup() error {
-	return closeBadger()
+	return closeDB()
 }
 
 func Main() error {

x-pack/apm-server/main_test.go

@@ -33,7 +33,7 @@ func TestMonitoring(t *testing.T) {
 	home := t.TempDir()
 	err := paths.InitPaths(&paths.Path{Home: home})
 	require.NoError(t, err)
-	defer closeBadger() // close badger.DB so data dir can be deleted on Windows
+	defer closeDB() // close DB so data dir can be deleted on Windows
 
 	cfg := config.DefaultConfig()
 	cfg.Sampling.Tail.Enabled = true

x-pack/apm-server/sampling/config.go

@@ -95,24 +95,16 @@ type DataStreamConfig struct {
 
 // StorageConfig holds Processor configuration related to event storage.
 type StorageConfig struct {
-	// DB holds the badger database in which event storage will be maintained.
+	// DB holds the StorageManager in which event storage will be maintained.
 	//
 	// DB will not be closed when the processor is closed.
 	DB *eventstorage.StorageManager
 
-	// Storage holds the read writers which provide sharded, locked access to storage.
-	//
-	// Storage lives outside processor lifecycle and will not be closed when processor
-	// is closed
-	Storage rw
-
-	// StorageDir holds the directory in which event storage will be maintained.
-	StorageDir string
-
-	// StorageGCInterval holds the amount of time between storage garbage collections.
-	StorageGCInterval time.Duration
+	// Storage overrides the default DB storage RW.
+	// For testing only.
+	Storage eventstorage.RW
 
-	// StorageLimit for the badger database, in bytes.
+	// StorageLimit for the TBS database, in bytes.
 	StorageLimit uint64
 
 	// TTL holds the amount of time before events and sampling decisions
@@ -242,15 +234,6 @@ func (config StorageConfig) validate() error {
 	if config.DB == nil {
 		return errors.New("DB unspecified")
 	}
-	if config.Storage == nil {
-		return errors.New("Storage unspecified")
-	}
-	if config.StorageDir == "" {
-		return errors.New("StorageDir unspecified")
-	}
-	if config.StorageGCInterval <= 0 {
-		return errors.New("StorageGCInterval unspecified or negative")
-	}
 	if config.TTL <= 0 {
 		return errors.New("TTL unspecified or negative")
 	}

x-pack/apm-server/sampling/config_test.go

@@ -73,15 +73,6 @@ func TestNewProcessorConfigInvalid(t *testing.T) {
 	assertInvalidConfigError("invalid storage config: DB unspecified")
 	config.DB = &eventstorage.StorageManager{}
 
-	assertInvalidConfigError("invalid storage config: Storage unspecified")
-	config.Storage = &eventstorage.ManagedReadWriter{}
-
-	assertInvalidConfigError("invalid storage config: StorageDir unspecified")
-	config.StorageDir = "tbs"
-
-	assertInvalidConfigError("invalid storage config: StorageGCInterval unspecified or negative")
-	config.StorageGCInterval = 1
-
 	assertInvalidConfigError("invalid storage config: TTL unspecified or negative")
 	config.TTL = 1
 }