TBS: set default sampling.tail.storage_limit to 0 but limit disk usage to 90%

internal/beater/config/config_test.go

@@ -362,8 +362,9 @@ func TestUnpackConfig(t *testing.T) {
 						ESConfig:              elasticsearch.DefaultConfig(),
 						Interval:              1 * time.Minute,
 						IngestRateDecayFactor: 0.25,
-						StorageLimit:          "3GB",
-						StorageLimitParsed:    3000000000,
+						StorageLimit:          "",
+						StorageLimitParsed:    0,
+						DiskThresholdRatio:    0.9,
 						TTL:                   30 * time.Minute,
 					},
 				},
@@ -410,11 +411,12 @@ func TestUnpackConfig(t *testing.T) {
 					},
 				},
 				"sampling.tail": map[string]interface{}{
-					"enabled":           false,
-					"policies":          []map[string]interface{}{{"sample_rate": 0.5}},
-					"interval":          "2m",
-					"ingest_rate_decay": 1.0,
-					"storage_limit":     "1GB",
+					"enabled":              false,
+					"policies":             []map[string]interface{}{{"sample_rate": 0.5}},
+					"interval":             "2m",
+					"ingest_rate_decay":    1.0,
+					"storage_limit":        "1GB",
+					"disk_threshold_ratio": 0.8,
 				},
 				"data_streams": map[string]interface{}{
 					"namespace": "foo",
@@ -495,6 +497,7 @@ func TestUnpackConfig(t *testing.T) {
 						IngestRateDecayFactor: 1.0,
 						StorageLimit:          "1GB",
 						StorageLimitParsed:    1000000000,
+						DiskThresholdRatio:    0.8,
 						TTL:                   30 * time.Minute,
 					},
 				},

internal/beater/config/sampling.go

@@ -49,8 +49,23 @@ type TailSamplingConfig struct {
 	Interval              time.Duration         `config:"interval" validate:"min=1s"`
 	IngestRateDecayFactor float64               `config:"ingest_rate_decay" validate:"min=0, max=1"`
 	TTL                   time.Duration         `config:"ttl" validate:"min=1s"`
-	StorageLimit          string                `config:"storage_limit"`
-	StorageLimitParsed    uint64
+
+	// StorageLimit is the user-configured tail-sampling database size limit.
+	// 0 means unlimited storage.
+	StorageLimit       string `config:"storage_limit"`
+	StorageLimitParsed uint64
+
+	// DiskThresholdRatio controls the proportion of the disk to be filled at max, irrespective of db size.
+	// e.g. 0.9 means the last 10% of disk should not be written to.
+	//
+	// This is additional to StorageLimit which controls db size,
+	// i.e. if both StorageLimit and DiskThresholdRatio are set, both are enforced.
+	// If StorageLimit is 0 and any error occurs when getting filesystem stats,
+	// a fallback StorageLimit will be applied.
+	// To disable the fallback to allow unlimited db size and disk threshold under all circumstances,
+	// set both StorageLimit and DiskThresholdRatio to 0.
+	DiskThresholdRatio float64 `config:"disk_threshold_ratio"  validate:"min=0, max=1"`
+
 	DiscardOnWriteFailure bool `config:"discard_on_write_failure"`
 
 	esConfigured bool
@@ -90,15 +105,16 @@ func (c *TailSamplingConfig) Unpack(in *config.C) error {
 		err = errors.Wrap(err, "error unpacking config")
 		return nil
 	}
-	limit, err := humanize.ParseBytes(cfg.StorageLimit)
-	if err != nil {
-		return err
+	if cfg.StorageLimit != "" {
+		limit, err := humanize.ParseBytes(cfg.StorageLimit)
+		if err != nil {
+			return err
+		}
+		cfg.StorageLimitParsed = limit
 	}
-	cfg.StorageLimitParsed = limit
 	cfg.Enabled = in.Enabled()
 	*c = TailSamplingConfig(cfg)
 	c.esConfigured = in.HasField("elasticsearch")
-	c.StorageLimitParsed = limit
 	err = errors.Wrap(c.Validate(), "invalid config")
 	return nil
 }
@@ -151,13 +167,9 @@ func defaultTailSamplingConfig() TailSamplingConfig {
 		Interval:              1 * time.Minute,
 		IngestRateDecayFactor: 0.25,
 		TTL:                   30 * time.Minute,
-		StorageLimit:          "3GB",
+		StorageLimit:          "",
+		DiskThresholdRatio:    0.9,
 		DiscardOnWriteFailure: false,
 	}
-	parsed, err := humanize.ParseBytes(cfg.StorageLimit)
-	if err != nil {
-		panic(err)
-	}
-	cfg.StorageLimitParsed = parsed
 	return cfg
 }

x-pack/apm-server/main.go

@@ -147,7 +147,7 @@ func newTailSamplingProcessor(args beater.ServerParams) (*sampling.Processor, er
 		},
 		StorageConfig: sampling.StorageConfig{
 			DB:                    db,
-			Storage:               db.NewReadWriter(tailSamplingConfig.StorageLimitParsed),
+			Storage:               db.NewReadWriter(tailSamplingConfig.StorageLimitParsed, tailSamplingConfig.DiskThresholdRatio),
 			TTL:                   tailSamplingConfig.TTL,
 			DiscardOnWriteFailure: tailSamplingConfig.DiscardOnWriteFailure,
 		},

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

@@ -16,6 +16,7 @@ import (
 	"time"
 
 	"github.com/cockroachdb/pebble/v2"
+	"github.com/cockroachdb/pebble/v2/vfs"
 	"go.opentelemetry.io/otel/metric"
 	"golang.org/x/sync/errgroup"
 
@@ -46,6 +47,10 @@ const (
 
 	// diskUsageFetchInterval is how often disk usage is fetched which is equivalent to how long disk usage is cached.
 	diskUsageFetchInterval = 1 * time.Second
+
+	// dbStorageLimitFallback is the default fallback storage limit in bytes
+	// that applies when disk threshold cannot be enforced due to an error.
+	dbStorageLimitFallback = 5 << 30
 )
 
 type StorageManagerOptions func(*StorageManager)
@@ -62,6 +67,19 @@ func WithMeterProvider(mp metric.MeterProvider) StorageManagerOptions {
 	}
 }
 
+// WithGetDiskUsage configures getDiskUsage function used by StorageManager.
+// For testing only.
+func WithGetDiskUsage(getDiskUsage func() (DiskUsage, error)) StorageManagerOptions {
+	return func(sm *StorageManager) {
+		sm.getDiskUsage = getDiskUsage
+	}
+}
+
+// DiskUsage is the struct returned by getDiskUsage.
+type DiskUsage struct {
+	UsedBytes, TotalBytes uint64
+}
+
 // StorageManager encapsulates pebble.DB.
 // It assumes exclusive access to pebble DB at storageDir.
 type StorageManager struct {
@@ -83,6 +101,13 @@ type StorageManager struct {
 	// cachedDBSize is a cached result of db size.
 	cachedDBSize atomic.Uint64
 
+	getDiskUsage func() (DiskUsage, error)
+	// diskStat is disk usage statistics about the disk only, not related to the databases.
+	diskStat struct {
+		used, total atomic.Uint64
+	}
+	diskStatFailed atomic.Bool
+
 	// runCh acts as a mutex to ensure only 1 Run is actively running per StorageManager.
 	// as it is possible that 2 separate Run are created by 2 TBS processors during a hot reload.
 	runCh chan struct{}
@@ -100,6 +125,13 @@ func NewStorageManager(storageDir string, opts ...StorageManagerOptions) (*Stora
 		runCh:      make(chan struct{}, 1),
 		logger:     logp.NewLogger(logs.Sampling),
 		codec:      ProtobufCodec{},
+		getDiskUsage: func() (DiskUsage, error) {
+			usage, err := vfs.Default.GetDiskUsage(storageDir)
+			return DiskUsage{
+				UsedBytes:  usage.UsedBytes,
+				TotalBytes: usage.TotalBytes,
+			}, err
+		},
 	}
 	for _, opt := range opts {
 		opt(sm)
@@ -211,6 +243,28 @@ func (sm *StorageManager) dbSize() uint64 {
 
 func (sm *StorageManager) updateDiskUsage() {
 	sm.cachedDBSize.Store(sm.eventDB.Metrics().DiskSpaceUsage() + sm.decisionDB.Metrics().DiskSpaceUsage())
+
+	if sm.diskStatFailed.Load() {
+		// Skip GetDiskUsage under the assumption that
+		// it will always get the same error if GetDiskUsage ever returns one,
+		// such that it does not keep logging GetDiskUsage errors.
+		return
+	}
+	usage, err := sm.getDiskUsage()
+	if err != nil {
+		sm.logger.With(logp.Error(err)).Warn("failed to get disk usage")
+		sm.diskStatFailed.Store(true)
+		sm.diskStat.total.Store(0) // setting total to 0 to disable any running disk threshold checks
+		return
+	}
+	sm.diskStat.used.Store(usage.UsedBytes)
+	sm.diskStat.total.Store(usage.TotalBytes)
+}
+
+// diskUsed returns the actual used disk space in bytes.
+// Not to be confused with dbSize which is specific to database.
+func (sm *StorageManager) diskUsed() uint64 {
+	return sm.diskStat.used.Load()
 }
 
 // runDiskUsageLoop runs a loop that updates cached disk usage regularly.
@@ -341,30 +395,61 @@ func (sm *StorageManager) WriteSubscriberPosition(data []byte) error {
 // NewUnlimitedReadWriter returns a read writer with no storage limit.
 // For testing only.
 func (sm *StorageManager) NewUnlimitedReadWriter() StorageLimitReadWriter {
-	return sm.NewReadWriter(0)
+	return sm.NewReadWriter(0, 0)
 }
 
-// NewReadWriter returns a read writer with storage limit.
-func (sm *StorageManager) NewReadWriter(storageLimit uint64) StorageLimitReadWriter {
+// NewReadWriter returns a read writer configured with storage limit and disk threshold ratio.
+func (sm *StorageManager) NewReadWriter(storageLimit uint64, diskThresholdRatio float64) StorageLimitReadWriter {
 	splitRW := SplitReadWriter{
 		eventRW:    sm.eventStorage.NewReadWriter(),
 		decisionRW: sm.decisionStorage.NewReadWriter(),
 	}
 
-	dbStorageLimit := func() uint64 {
-		return storageLimit
-	}
-	if storageLimit == 0 {
-		sm.logger.Infof("setting database storage limit to unlimited")
+	var dbStorageLimit func() uint64
+	if sm.diskStatFailed.Load() && storageLimit == 0 && diskThresholdRatio > 0 {
+		dbStorageLimit = func() uint64 {
+			return dbStorageLimitFallback
+		}
+		sm.logger.Warnf("failed to get disk usage; overriding storage_limit to fallback default %.1fgb", float64(dbStorageLimitFallback))
 	} else {
-		sm.logger.Infof("setting database storage limit to %.1fgb", float64(storageLimit))
+		dbStorageLimit = func() uint64 {
+			return storageLimit
+		}
+		if storageLimit == 0 {
+			sm.logger.Infof("setting database storage limit to unlimited")
+		} else {
+			sm.logger.Infof("setting database storage limit to %.1fgb", float64(storageLimit))
+		}
 	}
 
 	// To limit db size to storage_limit
 	dbStorageLimitChecker := NewStorageLimitCheckerFunc(sm.dbSize, dbStorageLimit)
 	dbStorageLimitRW := NewStorageLimitReadWriter("database storage limit", dbStorageLimitChecker, splitRW)
 
-	return dbStorageLimitRW
+	// diskThreshold returns max used disk space in bytes.
+	// After which, writes should be rejected.
+	var diskThreshold func() uint64
+	if sm.diskStatFailed.Load() {
+		diskThreshold = func() uint64 {
+			return 0 // unlimited
+		}
+		sm.logger.Warnf("failed to get disk usage; disabling disk threshold check")
+	} else {
+		diskThreshold = func() uint64 {
+			return uint64(float64(sm.diskStat.total.Load()) * diskThresholdRatio)
+		}
+		sm.logger.Infof("setting disk threshold ratio to %.2f", diskThresholdRatio)
+	}
+
+	// To limit actual disk usage percentage to diskThresholdRatio
+	diskThresholdChecker := NewStorageLimitCheckerFunc(sm.diskUsed, diskThreshold)
+	diskThresholdRW := NewStorageLimitReadWriter(
+		fmt.Sprintf("disk usage ratio exceeding threshold of %.2f", diskThresholdRatio),
+		diskThresholdChecker,
+		dbStorageLimitRW,
+	)
+
+	return diskThresholdRW
 }
 
 // wrapNonNilErr only wraps an error with format if the error is not nil.