Condition the OSS code to migrate the changes to ent

agent/agent.go

@@ -74,12 +74,12 @@ func (a *Agent) Run(ctx context.Context) error {
 	defer close(policyEvalCh)
 
 	limiter := policy.NewLimiter(policy.DefaultLimiterTimeout,
-		a.config.PolicyEval.Workers["horizontal"],
-		a.config.PolicyEval.Workers["cluster"])
+		a.config.PolicyEval.Workers)
 
 	if err := a.setupPolicyManager(limiter); err != nil {
 		return fmt.Errorf("failed to setup policy manager: %v", err)
 	}
+
 	go a.policyManager.Run(ctx, policyEvalCh)
 
 	a.initEnt(ctx, a.entReload)
@@ -89,48 +89,6 @@ func (a *Agent) Run(ctx context.Context) error {
 	return nil
 }
 
-func (a *Agent) setupPolicyManager(limiter *policy.Limiter) error {
-
-	// Create our processor, a shared method for performing basic policy
-	// actions.
-	cfgDefaults := policy.ConfigDefaults{
-		DefaultEvaluationInterval: a.config.Policy.DefaultEvaluationInterval,
-		DefaultCooldown:           a.config.Policy.DefaultCooldown,
-	}
-	policyProcessor := policy.NewProcessor(&cfgDefaults, a.getNomadAPMNames())
-
-	// Setup our initial default policy source which is Nomad.
-	sources := map[policy.SourceName]policy.Source{}
-	for _, s := range a.config.Policy.Sources {
-		if s.Enabled == nil || !*s.Enabled {
-			continue
-		}
-
-		switch policy.SourceName(s.Name) {
-		case policy.SourceNameNomad:
-			sources[policy.SourceNameNomad] = nomadPolicy.NewNomadSource(a.logger, a.NomadClient, policyProcessor)
-		case policy.SourceNameFile:
-			// Only setup the file source if operators have configured a
-			// scaling policy directory to read from.
-			if a.config.Policy.Dir != "" {
-				sources[policy.SourceNameFile] = filePolicy.NewFileSource(a.logger, a.config.Policy.Dir, policyProcessor)
-			}
-		}
-	}
-
-	// TODO: Once full policy source reload is implemented this should probably
-	// be just a warning.
-	if len(sources) == 0 {
-		return errors.New("no policy source available")
-	}
-
-	a.policySources = sources
-	a.policyManager = policy.NewManager(a.logger, a.policySources,
-		a.pluginManager, a.config.Telemetry.CollectionInterval, limiter)
-
-	return nil
-}
-
 func (a *Agent) stop() {
 	// Kill all the plugins.
 	if a.pluginManager != nil {
@@ -189,6 +147,48 @@ func (a *Agent) reload() {
 	}
 }
 
+func (a *Agent) setupPolicyManager(limiter *policy.Limiter) error {
+
+	// Create our processor, a shared method for performing basic policy
+	// actions.
+	cfgDefaults := policy.ConfigDefaults{
+		DefaultEvaluationInterval: a.config.Policy.DefaultEvaluationInterval,
+		DefaultCooldown:           a.config.Policy.DefaultCooldown,
+	}
+	policyProcessor := policy.NewProcessor(&cfgDefaults, a.getNomadAPMNames())
+
+	// Setup our initial default policy source which is Nomad.
+	sources := map[policy.SourceName]policy.Source{}
+	for _, s := range a.config.Policy.Sources {
+		if s.Enabled == nil || !*s.Enabled {
+			continue
+		}
+
+		switch policy.SourceName(s.Name) {
+		case policy.SourceNameNomad:
+			sources[policy.SourceNameNomad] = nomadPolicy.NewNomadSource(a.logger, a.NomadClient, policyProcessor)
+		case policy.SourceNameFile:
+			// Only setup the file source if operators have configured a
+			// scaling policy directory to read from.
+			if a.config.Policy.Dir != "" {
+				sources[policy.SourceNameFile] = filePolicy.NewFileSource(a.logger, a.config.Policy.Dir, policyProcessor)
+			}
+		}
+	}
+
+	// TODO: Once full policy source reload is implemented this should probably
+	// be just a warning.
+	if len(sources) == 0 {
+		return errors.New("no policy source available")
+	}
+
+	a.policySources = sources
+	a.policyManager = policy.NewManager(a.logger, a.policySources,
+		a.pluginManager, a.config.Telemetry.CollectionInterval, limiter)
+
+	return nil
+}
+
 // handleSignals blocks until the agent receives an exit signal.
 func (a *Agent) handleSignals() {
 

agent/agent_oss.go

@@ -6,7 +6,9 @@
 
 package agent
 
-import "context"
+import (
+	"context"
+)
 
 func (a *Agent) initEnt(ctx context.Context, reload <-chan any) {
 	go func() {

policy/check.go

@@ -41,7 +41,7 @@ type checkRunner struct {
 }
 
 // NewCheckHandler returns a new checkHandler instance.
-func NewCheckRunner(config *CheckRunnerConfig, c *sdk.ScalingPolicyCheck) *checkRunner {
+func newCheckRunner(config *CheckRunnerConfig, c *sdk.ScalingPolicyCheck) *checkRunner {
 	return &checkRunner{
 		log:            config.Log,
 		check:          c,
@@ -51,10 +51,10 @@ func NewCheckRunner(config *CheckRunnerConfig, c *sdk.ScalingPolicyCheck) *check
 	}
 }
 
-// GetNewCountFromStrategy begins the execution of the checks and returns
+// getNewCountFromStrategy begins the execution of the checks and returns
 // and action containing the instance count after applying the strategy to the
 // metrics.
-func (ch *checkRunner) GetNewCountFromStrategy(ctx context.Context, currentCount int64,
+func (ch *checkRunner) getNewCountFromStrategy(ctx context.Context, currentCount int64,
 	metrics sdk.TimestampedMetrics) (sdk.ScalingAction, error) {
 	ch.log.Debug("calculating new count", "current count", currentCount)
 
@@ -131,8 +131,8 @@ func (ch *checkRunner) runStrategy(ctx context.Context, currentCount int64, ms s
 	return *runResp.Action, nil
 }
 
-// QueryMetrics wraps the apm.Query call to provide operational functionality.
-func (ch *checkRunner) QueryMetrics(ctx context.Context) (sdk.TimestampedMetrics, error) {
+// queryMetrics wraps the apm.Query call to provide operational functionality.
+func (ch *checkRunner) queryMetrics(ctx context.Context) (sdk.TimestampedMetrics, error) {
 	ch.log.Debug("querying source", "query", ch.check.Query, "source", ch.check.Source)
 
 	// Trigger a metric measure to track latency of the call.
@@ -176,3 +176,26 @@ func (ch *checkRunner) QueryMetrics(ctx context.Context) (sdk.TimestampedMetrics
 
 	return ms, nil
 }
+
+func (ch *checkRunner) group() string {
+	return ch.check.Group
+}
+
+func (ch *checkRunner) runCheckAndCapCount(ctx context.Context, currentCount int64) (sdk.ScalingAction, error) {
+	ch.log.Debug("received policy check for evaluation")
+
+	metrics, err := ch.queryMetrics(ctx)
+	if err != nil {
+		return sdk.ScalingAction{}, fmt.Errorf("failed to query source: %v", err)
+	}
+
+	action, err := ch.getNewCountFromStrategy(ctx, currentCount, metrics)
+	if err != nil {
+		return sdk.ScalingAction{}, fmt.Errorf("failed get count from metrics: %v", err)
+
+	}
+
+	action.CapCount(ch.policy.Min, ch.policy.Max)
+
+	return action, nil
+}

policy/check_test.go

@@ -148,13 +148,13 @@ func TestCheckHandler_getNewCountFromMetrics(t *testing.T) {
 				},
 			}
 
-			runner := NewCheckRunner(&CheckRunnerConfig{
+			runner := newCheckRunner(&CheckRunnerConfig{
 				Log:            hclog.NewNullLogger(),
 				StrategyRunner: sr,
 				Policy:         tt.policy,
 			}, ch)
 
-			action, err := runner.GetNewCountFromStrategy(context.Background(), 3, tt.metrics)
+			action, err := runner.getNewCountFromStrategy(context.Background(), 3, tt.metrics)
 			must.Eq(t, tt.expectedAction, action)
 			must.Eq(t, tt.expError, errors.Unwrap(err))
 
@@ -226,15 +226,15 @@ func TestCheckHandler_runAPMQuery(t *testing.T) {
 				},
 			}
 
-			handler := NewCheckRunner(&CheckRunnerConfig{
+			handler := newCheckRunner(&CheckRunnerConfig{
 				Log:           hclog.NewNullLogger(),
 				MetricsGetter: mockLooker,
 				Policy: &sdk.ScalingPolicy{
 					ID: "testPolicy",
 				},
 			}, check)
 
-			result, err := handler.QueryMetrics(context.Background())
+			result, err := handler.queryMetrics(context.Background())
 			must.Eq(t, tc.expResult, result)
 			must.True(t, errors.Is(err, tc.expErr))
 		})