Adaptive (memory-feedback) staged-startup pacing

cmd/security-agent/subcommands/start/command.go

@@ -15,8 +15,6 @@ import (
 	_ "net/http/pprof" // Blank import used because this isn't directly used in this file
 	"os"
 	"os/signal"
-	"runtime"
-	"runtime/debug"
 	"strings"
 	"syscall"
 	"time"
@@ -77,6 +75,7 @@ import (
 	"github.com/DataDog/datadog-agent/pkg/util/defaultpaths"
 	"github.com/DataDog/datadog-agent/pkg/util/fxutil"
 	"github.com/DataDog/datadog-agent/pkg/util/profiling"
+	"github.com/DataDog/datadog-agent/pkg/util/stagedstart"
 	"github.com/DataDog/datadog-agent/pkg/util/startstop"
 	"github.com/DataDog/datadog-agent/pkg/version"
 )
@@ -211,23 +210,18 @@ func Commands(globalParams *command.GlobalParams) []*cobra.Command {
 }
 
 // stagedStartPace spreads security-agent subsystem startup when staged startup
-// is enabled: it reclaims the transient memory allocated by prior initialization
-// and pauses briefly before starting the next subsystem, so their startup memory
-// peaks do not stack into a single spike (both within this process and relative
-// to the other agent processes starting at the same time). When staged startup
-// is disabled this is a no-op.
+// is enabled: before starting the named subsystem it reclaims transient memory
+// from prior initialization and waits (a fixed interval, or until this process's
+// memory settles in adaptive mode) so their startup memory peaks do not stack
+// into a single spike, both within this process and relative to the other agent
+// processes starting at the same time. When staged startup is disabled this is a
+// no-op.
 func stagedStartPace(cfg config.Component, logger log.Component, name string) {
-	if !cfg.GetBool("staged_start.enabled") {
-		return
-	}
-	if cfg.GetBool("staged_start.free_os_memory") {
-		runtime.GC()
-		debug.FreeOSMemory()
-	}
-	if interval := cfg.GetDuration("staged_start.stage_interval"); interval > 0 {
-		logger.Infof("staged startup: pacing %s startup by %s", name, interval)
-		time.Sleep(interval)
-	}
+	pacer := stagedstart.NewPacer(stagedstart.ConfigFromReader(cfg),
+		func(f string, a ...interface{}) { logger.Infof(f, a...) },
+		func(f string, a ...interface{}) { logger.Warnf(f, a...) },
+	)
+	pacer.Pace(context.Background(), name)
 }
 
 // start will start the security-agent.

comp/core/startupsequencer/impl/startupsequencer.go

@@ -8,14 +8,13 @@ package startupsequencerimpl
 
 import (
 	"context"
-	"runtime"
-	"runtime/debug"
 	"sync"
 	"time"
 
 	config "github.com/DataDog/datadog-agent/comp/core/config"
 	log "github.com/DataDog/datadog-agent/comp/core/log/def"
 	startupsequencer "github.com/DataDog/datadog-agent/comp/core/startupsequencer/def"
+	"github.com/DataDog/datadog-agent/pkg/util/stagedstart"
 )
 
 // Requires defines the dependencies of the startupsequencer component.
@@ -35,10 +34,9 @@ type deferredStart struct {
 }
 
 type sequencer struct {
-	log          log.Component
-	enabled      bool
-	interval     time.Duration
-	freeOSMemory bool
+	log     log.Component
+	enabled bool
+	pacer   *stagedstart.Pacer
 
 	mu       sync.Mutex
 	begun    bool
@@ -47,12 +45,15 @@ type sequencer struct {
 
 // NewComponent returns the staged startup sequencer.
 func NewComponent(reqs Requires) Provides {
+	cfg := stagedstart.ConfigFromReader(reqs.Config)
 	return Provides{
 		Comp: &sequencer{
-			log:          reqs.Log,
-			enabled:      reqs.Config.GetBool("staged_start.enabled"),
-			interval:     reqs.Config.GetDuration("staged_start.stage_interval"),
-			freeOSMemory: reqs.Config.GetBool("staged_start.free_os_memory"),
+			log:     reqs.Log,
+			enabled: cfg.Enabled,
+			pacer: stagedstart.NewPacer(cfg,
+				func(f string, a ...interface{}) { reqs.Log.Infof(f, a...) },
+				func(f string, a ...interface{}) { reqs.Log.Warnf(f, a...) },
+			),
 		},
 	}
 }
@@ -94,34 +95,32 @@ func (s *sequencer) Begin(ctx context.Context) {
 }
 
 func (s *sequencer) run(ctx context.Context, deferred [startupsequencer.NumStages][]deferredStart) {
-	s.log.Infof("staged startup: beginning (stage interval %s)", s.interval)
+	// Flatten stages into a single ordered list. Stages are only a coarse
+	// priority (critical-path subsystems first); the pacer keeps consecutive
+	// items from allocating on top of one another, so the exact order within
+	// that priority does not affect the memory peak.
+	var items []deferredStart
 	for stage := startupsequencer.Stage(0); stage < startupsequencer.NumStages; stage++ {
-		for _, d := range deferred[stage] {
-			if ctx.Err() != nil {
-				s.log.Infof("staged startup: aborted before %q (context cancelled)", d.name)
-				return
-			}
-			start := time.Now()
-			if err := d.fn(ctx); err != nil {
-				s.log.Errorf("staged startup: %q (stage %d) failed: %v", d.name, stage, err)
-			} else {
-				s.log.Debugf("staged startup: started %q (stage %d) in %s", d.name, stage, time.Since(start))
-			}
-		}
+		items = append(items, deferred[stage]...)
+	}
 
-		// Return transient memory allocated during this stage to the OS before
-		// the next stage allocates, keeping the peak RSS close to steady state.
-		if s.freeOSMemory {
-			runtime.GC()
-			debug.FreeOSMemory()
+	s.log.Info("staged startup: beginning")
+	for i, d := range items {
+		if ctx.Err() != nil {
+			s.log.Infof("staged startup: aborted before %q (context cancelled)", d.name)
+			return
+		}
+		start := time.Now()
+		if err := d.fn(ctx); err != nil {
+			s.log.Errorf("staged startup: %q failed: %v", d.name, err)
+		} else {
+			s.log.Debugf("staged startup: started %q in %s", d.name, time.Since(start))
 		}
 
-		if stage < startupsequencer.NumStages-1 {
-			select {
-			case <-ctx.Done():
-				return
-			case <-time.After(s.interval):
-			}
+		// Pace before releasing the next item (reclaims transient memory and,
+		// in adaptive mode, waits until this item's allocation settles).
+		if i < len(items)-1 {
+			s.pacer.Pace(ctx, d.name)
 		}
 	}
 	s.log.Info("staged startup: complete")

comp/core/startupsequencer/impl/startupsequencer_test.go

@@ -17,19 +17,21 @@ import (
 
 	logmock "github.com/DataDog/datadog-agent/comp/core/log/mock"
 	startupsequencer "github.com/DataDog/datadog-agent/comp/core/startupsequencer/def"
+	"github.com/DataDog/datadog-agent/pkg/util/stagedstart"
 )
 
+// newTestSequencer builds a sequencer whose pacer is a no-op (no delay, no
+// reclaim) so ordering/inline tests run instantly.
 func newTestSequencer(t *testing.T, enabled bool) *sequencer {
+	cfg := stagedstart.Config{Enabled: enabled}
 	return &sequencer{
-		log:          logmock.New(t),
-		enabled:      enabled,
-		interval:     time.Millisecond,
-		freeOSMemory: false,
+		log:     logmock.New(t),
+		enabled: enabled,
+		pacer:   stagedstart.NewPacer(cfg, nil, nil),
 	}
 }
 
-// When disabled, Defer must run the work synchronously and propagate its error,
-// so a Defer call from an OnStart hook is identical to running the work inline.
+// When disabled, Defer must run the work synchronously and propagate its error.
 func TestDisabledRunsInline(t *testing.T) {
 	s := newTestSequencer(t, false)
 
@@ -45,7 +47,6 @@ func TestDisabledRunsInline(t *testing.T) {
 	err = s.Defer(startupsequencer.StageChecks, "y", func(context.Context) error { return sentinel })
 	assert.ErrorIs(t, err, sentinel, "error must propagate when staging is disabled")
 
-	// Begin is a no-op when disabled.
 	s.Begin(context.Background())
 }
 
@@ -66,7 +67,6 @@ func TestEnabledRunsInStageOrder(t *testing.T) {
 	}
 
 	done := make(chan struct{})
-	// Register out of stage order to prove the sequencer orders by stage.
 	require.NoError(t, s.Defer(startupsequencer.StageBackground, "background", func(ctx context.Context) error {
 		_ = record("background")(ctx)
 		close(done)
@@ -75,7 +75,6 @@ func TestEnabledRunsInStageOrder(t *testing.T) {
 	require.NoError(t, s.Defer(startupsequencer.StageCritical, "critical", record("critical")))
 	require.NoError(t, s.Defer(startupsequencer.StageIngest, "ingest", record("ingest")))
 
-	// Nothing should have run before Begin.
 	mu.Lock()
 	assert.Empty(t, order, "no deferred work should run before Begin")
 	mu.Unlock()
@@ -93,8 +92,7 @@ func TestEnabledRunsInStageOrder(t *testing.T) {
 	assert.Equal(t, []string{"critical", "ingest", "background"}, order)
 }
 
-// Work registered after the sequence has begun must still run (inline), not be
-// silently dropped.
+// Work registered after the sequence has begun must still run (inline).
 func TestLateRegistrationRunsInline(t *testing.T) {
 	s := newTestSequencer(t, true)
 	s.Begin(context.Background())
@@ -107,10 +105,12 @@ func TestLateRegistrationRunsInline(t *testing.T) {
 	assert.True(t, ran, "work registered after Begin should run inline")
 }
 
-// A cancelled context must stop the sequence rather than running later stages.
+// A cancelled context must stop the sequence rather than running later items.
 func TestContextCancellationStops(t *testing.T) {
 	s := newTestSequencer(t, true)
-	s.interval = time.Hour // ensure we block between stages
+	// A long pacer interval makes the sequencer block between items so we can
+	// cancel mid-sequence.
+	s.pacer = stagedstart.NewPacer(stagedstart.Config{Enabled: true, Interval: time.Hour}, nil, nil)
 
 	first := make(chan struct{})
 	require.NoError(t, s.Defer(startupsequencer.StageCritical, "first", func(context.Context) error {
@@ -126,11 +126,9 @@ func TestContextCancellationStops(t *testing.T) {
 	ctx, cancel := context.WithCancel(context.Background())
 	s.Begin(ctx)
 
-	<-first  // first stage ran; sequencer now sleeping before stage 2
-	cancel() // cancel during the inter-stage wait
+	<-first  // first item ran; sequencer now paused before the second
+	cancel() // cancel during the inter-item wait
 
-	// The second stage is gated behind a one-hour inter-stage wait, so the only
-	// way it could run is if cancellation were ignored.
 	time.Sleep(50 * time.Millisecond)
-	assert.False(t, secondRan, "later stage must not run after cancellation")
+	assert.False(t, secondRan, "later item must not run after cancellation")
 }

pkg/config/setup/common_settings.go

@@ -83,6 +83,23 @@ func initCoreAgentFull(config pkgconfigmodel.Setup) {
 	config.BindEnvAndSetDefault("staged_start.stage_interval", 30*time.Second)
 	// If true, transient memory is returned to the OS between staged startup stages.
 	config.BindEnvAndSetDefault("staged_start.free_os_memory", true)
+	// Pacing mode between staged startup steps: "interval" (fixed stage_interval
+	// delay) or "adaptive" (advance as this process's memory settles, bounded by
+	// step_max so a process under memory pressure proceeds instead of hanging).
+	// NOTE: temporarily defaulted to "adaptive" on this branch so SMP exercises
+	// it; revisit the shipping default before merge.
+	config.BindEnvAndSetDefault("staged_start.mode", "adaptive")
+	// Adaptive mode: minimum wait per step (lets a step's async work begin).
+	config.BindEnvAndSetDefault("staged_start.step_min", 1*time.Second)
+	// Adaptive mode: hard cap per step. On timeout the next step proceeds anyway
+	// (and a warning is logged) so startup never hangs on memory that won't settle.
+	config.BindEnvAndSetDefault("staged_start.step_max", 15*time.Second)
+	// Adaptive mode: a step is "settled" once retained memory grows less than
+	// settle_threshold_bytes over this window.
+	config.BindEnvAndSetDefault("staged_start.settle_window", 2*time.Second)
+	config.BindEnvAndSetDefault("staged_start.settle_threshold_bytes", int64(2*1024*1024))
+	// Adaptive mode: how often to sample memory while waiting for a step to settle.
+	config.BindEnvAndSetDefault("staged_start.poll_interval", 250*time.Millisecond)
 
 	// If true, then new version of disk v2 check will be used.
 	// Otherwise, the python version of disk check will be used.

pkg/system-probe/api/module/loader.go

@@ -11,8 +11,6 @@ import (
 	"fmt"
 	"maps"
 	"net/http"
-	"runtime"
-	"runtime/debug"
 	"runtime/pprof"
 	"sync"
 	"time"
@@ -21,6 +19,7 @@ import (
 	rcclient "github.com/DataDog/datadog-agent/comp/remote-config/rcclient/def"
 	sysconfigtypes "github.com/DataDog/datadog-agent/pkg/system-probe/config/types"
 	"github.com/DataDog/datadog-agent/pkg/util/log"
+	"github.com/DataDog/datadog-agent/pkg/util/stagedstart"
 )
 
 var l *loader
@@ -90,24 +89,25 @@ func Register(cfg *sysconfigtypes.Config, httpMux *http.ServeMux, factories []*F
 	// Staged startup: creating a module loads its eBPF assets, which allocates a
 	// large but transient amount of scratch (BTF parsing, bytecode, CO-RE
 	// relocations). Loading every module back-to-back stacks those transients
-	// into a single startup memory peak. When enabled, reclaim each module's
-	// scratch and pause briefly between modules so the peak stays close to the
-	// steady-state footprint. The total added delay is bounded by one stage
-	// interval regardless of how many modules are enabled.
-	stagedStart := deps.CoreConfig.GetBool("staged_start.enabled")
-	freeOSMemory := deps.CoreConfig.GetBool("staged_start.free_os_memory")
-	var moduleDelay time.Duration
-	if stagedStart && len(enabledModulesFactories) > 1 {
-		moduleDelay = deps.CoreConfig.GetDuration("staged_start.stage_interval") / time.Duration(len(enabledModulesFactories))
+	// into a single startup memory peak. When enabled, the pacer reclaims each
+	// module's scratch and waits between modules (a fixed slice of the stage
+	// interval, or until memory settles in adaptive mode) so the peak stays
+	// close to the steady-state footprint.
+	pacerCfg := stagedstart.ConfigFromReader(deps.CoreConfig)
+	// In fixed-interval mode, spread the whole stage interval across the modules
+	// (rather than waiting a full interval per module). Adaptive mode self-bounds
+	// per step via step_max and ignores this.
+	if !pacerCfg.Adaptive && len(enabledModulesFactories) > 1 {
+		pacerCfg.Interval /= time.Duration(len(enabledModulesFactories))
 	}
+	pacer := stagedstart.NewPacer(pacerCfg,
+		func(f string, a ...interface{}) { log.Infof(f, a...) },
+		func(f string, a ...interface{}) { log.Warnf(f, a...) },
+	)
 
 	for i, factory := range enabledModulesFactories {
-		if stagedStart && i > 0 {
-			if freeOSMemory {
-				runtime.GC()
-				debug.FreeOSMemory()
-			}
-			time.Sleep(moduleDelay)
+		if i > 0 {
+			pacer.Pace(context.Background(), string(enabledModulesFactories[i-1].Name))
 		}
 
 		var err error