TBS: Optimize performance for instances with more CPU and memory (elastic#17254)

To improve performance for larger instances e.g. 16GB+ memory on cloud, a few tweaks are done:
- DB cache size is autoscaled to memory limit
- Enable concurrent compaction to increase CPU utilization

benchstat 32gb GCP on ECH, 0.01 sample rate. +60% throughput after PR.

Author: carsonip

internal/beater/beater.go

@@ -240,6 +240,14 @@ func (s *Runner) Run(ctx context.Context) error {
 		)
 	}
 
+	if s.config.Sampling.Tail.Enabled && s.config.Sampling.Tail.DatabaseCacheSize == 0 {
+		// 1GB=16MB, 2GB=24MB, 4GB=40MB, ..., 32GB=264MB, 64GB=520MB
+		s.config.Sampling.Tail.DatabaseCacheSize = uint64(linearScaledValue(8<<20, memLimitGB, 8<<20))
+		s.logger.Infof("Sampling.Tail.DatabaseCacheSize set to %d based on %0.1fgb of memory",
+			s.config.Sampling.Tail.DatabaseCacheSize, memLimitGB,
+		)
+	}
+
 	// Send config to telemetry.
 	recordAPMServerConfig(s.config)
 

internal/beater/config/sampling.go

@@ -63,6 +63,9 @@ type TailSamplingConfig struct {
 
 	DiscardOnWriteFailure bool `config:"discard_on_write_failure"`
 
+	// DatabaseCacheSize is cache size in bytes for tail-sampling database.
+	DatabaseCacheSize uint64 `config:"database_cache_size"`
+
 	esConfigured bool
 }
 

x-pack/apm-server/main.go

@@ -108,7 +108,7 @@ func newTailSamplingProcessor(args beater.ServerParams) (*sampling.Processor, er
 	}
 
 	storageDir := paths.Resolve(paths.Data, tailSamplingStorageDir)
-	db, err := getDB(storageDir, args.MeterProvider, args.Logger)
+	db, err := getDB(storageDir, tailSamplingConfig.DatabaseCacheSize, args.MeterProvider, args.Logger)
 	if err != nil {
 		return nil, fmt.Errorf("failed to get tail-sampling database: %w", err)
 	}
@@ -154,11 +154,13 @@ func newTailSamplingProcessor(args beater.ServerParams) (*sampling.Processor, er
 	}, args.Logger)
 }
 
-func getDB(storageDir string, mp metric.MeterProvider, logger *logp.Logger) (*eventstorage.StorageManager, error) {
+func getDB(storageDir string, cacheSize uint64, mp metric.MeterProvider, logger *logp.Logger) (*eventstorage.StorageManager, error) {
 	dbMu.Lock()
 	defer dbMu.Unlock()
 	if db == nil {
-		var opts []eventstorage.StorageManagerOptions
+		opts := []eventstorage.StorageManagerOptions{
+			eventstorage.WithDBCacheSize(cacheSize),
+		}
 		if mp != nil {
 			opts = append(opts, eventstorage.WithMeterProvider(mp))
 		}

x-pack/apm-server/sampling/eventstorage/pebble.go

@@ -38,8 +38,10 @@ func eventComparer() *pebble.Comparer {
 	return &comparer
 }
 
-func OpenEventPebble(storageDir string, logger *logp.Logger) (*pebble.DB, error) {
+func OpenEventPebble(storageDir string, cacheSize uint64, logger *logp.Logger) (*pebble.DB, error) {
 	// Option values are picked and validated in https://github.com/elastic/apm-server/issues/15568
+	cache := pebble.NewCache(int64(cacheSize))
+	defer cache.Unref()
 	opts := &pebble.Options{
 		FormatMajorVersion: pebble.FormatColumnarBlocks,
 		Logger:             logger.Named(logs.Sampling),
@@ -53,13 +55,19 @@ func OpenEventPebble(storageDir string, logger *logp.Logger) (*pebble.DB, error)
 			},
 		},
 		Comparer: eventComparer(),
+		Cache:    cache,
+		MaxConcurrentCompactions: func() int {
+			return 2
+		}, // Better utilizes CPU on larger instances
 	}
 	return pebble.Open(filepath.Join(storageDir, "event"), opts)
 }
 
-func OpenDecisionPebble(storageDir string, logger *logp.Logger) (*pebble.DB, error) {
+func OpenDecisionPebble(storageDir string, cacheSize uint64, logger *logp.Logger) (*pebble.DB, error) {
 	// Option values are picked and validated in https://github.com/elastic/apm-server/issues/15568
-	return pebble.Open(filepath.Join(storageDir, "decision"), &pebble.Options{
+	cache := pebble.NewCache(int64(cacheSize))
+	defer cache.Unref()
+	opts := &pebble.Options{
 		FormatMajorVersion: pebble.FormatColumnarBlocks,
 		Logger:             logger.Named(logs.Sampling),
 		MemTableSize:       2 << 20, // big memtables are slow to scan, and significantly slow the hot path
@@ -71,5 +79,10 @@ func OpenDecisionPebble(storageDir string, logger *logp.Logger) (*pebble.DB, err
 				FilterType:   pebble.TableFilter,
 			},
 		},
-	})
+		Cache: cache,
+		MaxConcurrentCompactions: func() int {
+			return 2
+		}, // Better utilizes CPU on larger instances
+	}
+	return pebble.Open(filepath.Join(storageDir, "decision"), opts)
 }

x-pack/apm-server/sampling/eventstorage/prefix_test.go

@@ -19,7 +19,7 @@ import (
 )
 
 func newEventPebble(t *testing.T) *pebble.DB {
-	db, err := eventstorage.OpenEventPebble(t.TempDir(), logptest.NewTestingLogger(t, ""))
+	db, err := eventstorage.OpenEventPebble(t.TempDir(), 8<<20, logptest.NewTestingLogger(t, ""))
 	require.NoError(t, err)
 	t.Cleanup(func() {
 		db.Close()
@@ -28,7 +28,7 @@ func newEventPebble(t *testing.T) *pebble.DB {
 }
 
 func newDecisionPebble(t *testing.T) *pebble.DB {
-	db, err := eventstorage.OpenDecisionPebble(t.TempDir(), logptest.NewTestingLogger(t, ""))
+	db, err := eventstorage.OpenDecisionPebble(t.TempDir(), 8<<20, logptest.NewTestingLogger(t, ""))
 	require.NoError(t, err)
 	t.Cleanup(func() {
 		db.Close()

x-pack/apm-server/sampling/eventstorage/storage_manager.go

@@ -88,6 +88,13 @@ func WithGetDiskUsage(getDiskUsage func() (DiskUsage, error)) StorageManagerOpti
 	}
 }
 
+// WithDBCacheSize sets the total size in bytes of in-memory cache of all databases managed by StorageManager.
+func WithDBCacheSize(size uint64) StorageManagerOptions {
+	return func(sm *StorageManager) {
+		sm.dbCacheSize = size
+	}
+}
+
 // DiskUsage is the struct returned by getDiskUsage.
 type DiskUsage struct {
 	UsedBytes, TotalBytes uint64
@@ -96,8 +103,9 @@ type DiskUsage struct {
 // StorageManager encapsulates pebble.DB.
 // It assumes exclusive access to pebble DB at storageDir.
 type StorageManager struct {
-	storageDir string
-	logger     *logp.Logger
+	storageDir  string
+	dbCacheSize uint64
+	logger      *logp.Logger
 
 	eventDB         *pebble.DB
 	decisionDB      *pebble.DB
@@ -155,6 +163,7 @@ func NewStorageManager(storageDir string, logger *logp.Logger, opts ...StorageMa
 				TotalBytes: usage.TotalBytes,
 			}, err
 		},
+		dbCacheSize: 16 << 20, // default to 16MB cache shared between event and decision DB
 	}
 	sm.getDBSize = func() uint64 {
 		return sm.eventDB.Metrics().DiskSpaceUsage() + sm.decisionDB.Metrics().DiskSpaceUsage()
@@ -179,13 +188,14 @@ func NewStorageManager(storageDir string, logger *logp.Logger, opts ...StorageMa
 
 // reset initializes db and storage.
 func (sm *StorageManager) reset() error {
-	eventDB, err := OpenEventPebble(sm.storageDir, sm.logger)
+	// Configured db cache size is split between event DB and decision DB
+	eventDB, err := OpenEventPebble(sm.storageDir, sm.dbCacheSize/2, sm.logger)
 	if err != nil {
 		return fmt.Errorf("open event db error: %w", err)
 	}
 	sm.eventDB = eventDB
 
-	decisionDB, err := OpenDecisionPebble(sm.storageDir, sm.logger)
+	decisionDB, err := OpenDecisionPebble(sm.storageDir, sm.dbCacheSize/2, sm.logger)
 	if err != nil {
 		return fmt.Errorf("open decision db error: %w", err)
 	}