agent/core: Allow setting goal CU in tests

pkg/agent/core/dumpstate.go

@@ -22,25 +22,24 @@ func shallowCopy[T any](ptr *T) *T {
 //
 // It implements json.Marshaler.
 type StateDump struct {
-	internal state
+	internal state[*StdAlgorithm]
 }
 
 func (d StateDump) MarshalJSON() ([]byte, error) {
 	return json.Marshal(d.internal)
 }
 
 // Dump produces a JSON-serializable copy of the State
-func (s *State) Dump() StateDump {
+func DumpState(s *State[*StdAlgorithm]) StateDump {
 	return StateDump{
-		internal: state{
+		internal: state[*StdAlgorithm]{
 			Debug:                s.internal.Debug,
 			Config:               s.internal.Config,
 			VM:                   s.internal.VM,
 			Plugin:               s.internal.Plugin.deepCopy(),
 			Monitor:              s.internal.Monitor.deepCopy(),
 			NeonVM:               s.internal.NeonVM.deepCopy(),
-			Metrics:              shallowCopy[SystemMetrics](s.internal.Metrics),
-			LFCMetrics:           shallowCopy[LFCMetrics](s.internal.LFCMetrics),
+			Metrics:              shallowCopy[StdAlgorithm](s.internal.Metrics),
 			TargetRevision:       s.internal.TargetRevision,
 			LastDesiredResources: s.internal.LastDesiredResources,
 		},

pkg/agent/core/goalcu.go

@@ -12,49 +12,102 @@ import (
 	"github.com/neondatabase/autoscaling/pkg/api"
 )
 
-type scalingGoal struct {
-	hasAllMetrics bool
-	goalCU        uint32
+// AlgorithmState abstracts over providers of "goal CU" calculation.
+//
+// This interface exists so that our unit tests for State.NextActions() don't need to be rewritten
+// when the algorithm changes.
+//
+// StdAlgorithm is expected to be used for anything outside unit tests.
+type AlgorithmState interface {
+	// CalculateGoalCU returns the desired compute units to scale to, plus any log fields with
+	// additional useful information.
+	CalculateGoalCU(
+		warn func(string),
+		cfg api.ScalingConfig,
+		computeUnit api.Resources,
+	) (ScalingGoal, []zap.Field)
+
+	// LatestAPIMetrics returns the api.Metrics that should be sent to the scheduler for
+	// prioritization of live migration.
+	LatestAPIMetrics() *api.Metrics
+
+	// ScalingConfigUpdated gives the algorithm a chance to update its internal state to reflect
+	// that the scaling config has been updated.
+	//
+	// In practice, this is only used by StdAlgorithm to reset its LFC metrics when they become
+	// disabled, so that we don't accidentally use stale metrics if/when it's re-enabled.
+	ScalingConfigUpdated(api.ScalingConfig)
 }
 
-func calculateGoalCU(
+type ScalingGoal struct {
+	HasAllMetrics bool
+	GoalCU        uint32
+}
+
+// StdAlgorithm is the standard implementation of AlgorithmState for usage in production.
+type StdAlgorithm struct {
+	System *SystemMetrics
+	LFC    *LFCMetrics
+}
+
+func DefaultAlgorithm() *StdAlgorithm {
+	return &StdAlgorithm{
+		System: nil,
+		LFC:    nil,
+	}
+}
+
+// LatestAPIMetrics implements AlgorithmState
+func (m *StdAlgorithm) LatestAPIMetrics() *api.Metrics {
+	if m.System != nil {
+		return lo.ToPtr(m.System.ToAPI())
+	} else {
+		return nil
+	}
+}
+
+func (m *StdAlgorithm) ScalingConfigUpdated(conf api.ScalingConfig) {
+	// Make sure that if LFC metrics are disabled & later enabled, we don't make decisions based on
+	// stale data.
+	if !*conf.EnableLFCMetrics {
+		m.LFC = nil
+	}
+}
+
+// CalculateGoalCU implements AlgorithmState
+func (m *StdAlgorithm) CalculateGoalCU(
 	warn func(string),
 	cfg api.ScalingConfig,
 	computeUnit api.Resources,
-	systemMetrics *SystemMetrics,
-	lfcMetrics *LFCMetrics,
-) (scalingGoal, []zap.Field) {
-	hasAllMetrics := systemMetrics != nil && (!*cfg.EnableLFCMetrics || lfcMetrics != nil)
-	if !hasAllMetrics {
-		warn("Making scaling decision without all required metrics available")
-	}
+) (ScalingGoal, []zap.Field) {
+	hasAllMetrics := m.System != nil && (!*cfg.EnableLFCMetrics || m.LFC != nil)
 
 	var lfcGoalCU, cpuGoalCU, memGoalCU, memTotalGoalCU uint32
 	var logFields []zap.Field
 
 	var wss *api.Bytes // estimated working set size
 
-	if lfcMetrics != nil {
+	if m.LFC != nil {
 		var lfcLogFunc func(zapcore.ObjectEncoder) error
-		lfcGoalCU, wss, lfcLogFunc = calculateLFCGoalCU(warn, cfg, computeUnit, *lfcMetrics)
+		lfcGoalCU, wss, lfcLogFunc = calculateLFCGoalCU(warn, cfg, computeUnit, *m.LFC)
 		if lfcLogFunc != nil {
 			logFields = append(logFields, zap.Object("lfc", zapcore.ObjectMarshalerFunc(lfcLogFunc)))
 		}
 	}
 
-	if systemMetrics != nil {
-		cpuGoalCU = calculateCPUGoalCU(cfg, computeUnit, *systemMetrics)
+	if m.System != nil {
+		cpuGoalCU = calculateCPUGoalCU(cfg, computeUnit, *m.System)
 
-		memGoalCU = calculateMemGoalCU(cfg, computeUnit, *systemMetrics)
+		memGoalCU = calculateMemGoalCU(cfg, computeUnit, *m.System)
 	}
 
-	if systemMetrics != nil && wss != nil {
-		memTotalGoalCU = calculateMemTotalGoalCU(cfg, computeUnit, *systemMetrics, *wss)
+	if m.System != nil && wss != nil {
+		memTotalGoalCU = calculateMemTotalGoalCU(cfg, computeUnit, *m.System, *wss)
 	}
 
 	goalCU := max(cpuGoalCU, memGoalCU, memTotalGoalCU, lfcGoalCU)
 
-	return scalingGoal{hasAllMetrics: hasAllMetrics, goalCU: goalCU}, logFields
+	return ScalingGoal{HasAllMetrics: hasAllMetrics, GoalCU: goalCU}, logFields
 }
 
 // For CPU: