Skip to content

🚀 Speculative agent autoscaling #664

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 2 commits into
base: main
Choose a base branch
from
+181 −11
Open

🚀 Speculative agent autoscaling #664

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
1131063
Fix issue with speculative runs not scaling agents
Shweta-Deshpande Nov 4, 2025
9bf3a74
Add mocks and unit test
Shweta-Deshpande Nov 12, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions go.mod
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,12 +5,15 @@ go 1.25.4
require (
github.com/go-logr/logr v1.4.3
github.com/go-logr/zapr v1.3.0
github.com/golang/mock v1.6.0
github.com/google/go-cmp v0.7.0
github.com/hashicorp/go-slug v0.16.7
github.com/hashicorp/go-tfe v1.93.0
github.com/onsi/ginkgo/v2 v2.23.4
github.com/onsi/gomega v1.36.3
github.com/prometheus/client_golang v1.22.0
github.com/stretchr/testify v1.11.1
go.uber.org/mock v0.4.0
go.uber.org/zap v1.27.0
k8s.io/api v0.34.1
k8s.io/apimachinery v0.34.1
Expand Down
13 changes: 13 additions & 0 deletions go.sum
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,8 @@ github.com/go-task/slim-sprig/v3 v3.0.0 h1:sUs3vkvUymDpBKi3qH1YSqBQk9+9D/8M2mN1v
github.com/go-task/slim-sprig/v3 v3.0.0/go.mod h1:W848ghGpv3Qj3dhTPRyJypKRiqCdHZiAzKg9hl15HA8=
github.com/gogo/protobuf v1.3.2 h1:Ov1cvc58UF3b5XjBnZv7+opcTcQFZebYjWzi34vdm4Q=
github.com/gogo/protobuf v1.3.2/go.mod h1:P1XiOD3dCwIKUDQYPy72D8LYyHL2YPYrpS2s69NZV8Q=
github.com/golang/mock v1.6.0 h1:ErTB+efbowRARo13NNdxyJji2egdxLGQhRaY+DUumQc=
github.com/golang/mock v1.6.0/go.mod h1:p6yTPP+5HYm5mzsMV8JkE6ZKdX+/wYM6Hr+LicevLPs=
github.com/google/btree v1.1.3 h1:CVpQJjYgC4VbzxeGVHfvZrv1ctoYCAI8vbl07Fcxlyg=
github.com/google/btree v1.1.3/go.mod h1:qOPhT0dTNdNzV6Z/lhRX0YXUafgPLFUh+gZMl761Gm4=
github.com/google/gnostic-models v0.7.0 h1:qwTtogB15McXDaNqTZdzPJRHvaVJlAl+HVQnLmJEJxo=
Expand Down Expand Up @@ -136,10 +138,13 @@ github.com/x448/float16 v0.8.4 h1:qLwI1I70+NjRFUR3zs1JPUCgaCXSh3SW62uAKT1mSBM=
github.com/x448/float16 v0.8.4/go.mod h1:14CWIYCyZA/cWjXOioeEpHeN/83MdbZDRQHoFcYsOfg=
github.com/yuin/goldmark v1.1.27/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=
github.com/yuin/goldmark v1.2.1/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=
github.com/yuin/goldmark v1.3.5/go.mod h1:mwnBkeHKe2W/ZEtQ+71ViKU8L12m81fl3OWwC1Zlc8k=
go.uber.org/automaxprocs v1.6.0 h1:O3y2/QNTOdbF+e/dpXNNW7Rx2hZ4sTIPyybbxyNqTUs=
go.uber.org/automaxprocs v1.6.0/go.mod h1:ifeIMSnPZuznNm6jmdzmU3/bfk01Fe2fotchwEFJ8r8=
go.uber.org/goleak v1.3.0 h1:2K3zAYmnTNqV73imy9J1T3WC+gmCePx2hEGkimedGto=
go.uber.org/goleak v1.3.0/go.mod h1:CoHD4mav9JJNrW/WLlf7HGZPjdw8EucARQHekz1X6bE=
go.uber.org/mock v0.4.0 h1:VcM4ZOtdbR4f6VXfiOpwpVJDL6lCReaZ6mw31wqh7KU=
go.uber.org/mock v0.4.0/go.mod h1:a6FSlNadKUHUa9IP5Vyt1zh4fC7uAwxMutEAscFbkZc=
go.uber.org/multierr v1.11.0 h1:blXXJkSxSSfBVBlC76pxqeO+LN3aDfLQo+309xJstO0=
go.uber.org/multierr v1.11.0/go.mod h1:20+QtiLqy0Nd6FdQB9TLXag12DsQkrbs3htMFfDN80Y=
go.uber.org/zap v1.27.0 h1:aJMhYGrd5QSmlpLMr2MftRKl7t8J8PTZPA732ud/XR8=
Expand All @@ -153,24 +158,31 @@ golang.org/x/crypto v0.0.0-20191011191535-87dc89f01550/go.mod h1:yigFU9vqHzYiE8U
golang.org/x/crypto v0.0.0-20200622213623-75b288015ac9/go.mod h1:LzIPMQfyMNhhGPhUkYOs5KpL4U8rLKemX1yGLhDgUto=
golang.org/x/mod v0.2.0/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA=
golang.org/x/mod v0.3.0/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA=
golang.org/x/mod v0.4.2/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA=
golang.org/x/net v0.0.0-20190404232315-eb5bcb51f2a3/go.mod h1:t9HGtf8HONx5eT2rtn7q6eTqICYqUVnKs3thJo3Qplg=
golang.org/x/net v0.0.0-20190620200207-3b0461eec859/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s=
golang.org/x/net v0.0.0-20200226121028-0de0cce0169b/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s=
golang.org/x/net v0.0.0-20201021035429-f5854403a974/go.mod h1:sp8m0HH+o8qH0wwXwYZr8TS3Oi6o0r6Gce1SSxlDquU=
golang.org/x/net v0.0.0-20210405180319-a5a99cb37ef4/go.mod h1:p54w0d4576C0XHj96bSt6lcn1PtDYWL6XObtHCRCNQM=
golang.org/x/net v0.38.0 h1:vRMAPTMaeGqVhG5QyLJHqNDwecKTomGeqbnfZyKlBI8=
golang.org/x/net v0.38.0/go.mod h1:ivrbrMbzFq5J41QOQh0siUuly180yBYtLp+CKbEaFx8=
golang.org/x/oauth2 v0.28.0 h1:CrgCKl8PPAVtLnU3c+EDw6x11699EWlsDeWNWKdIOkc=
golang.org/x/oauth2 v0.28.0/go.mod h1:onh5ek6nERTohokkhCD/y2cV4Do3fxFHFuAejCkRWT8=
golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
golang.org/x/sync v0.0.0-20190911185100-cd5d95a43a6e/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
golang.org/x/sync v0.0.0-20201020160332-67f06af15bc9/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
golang.org/x/sync v0.0.0-20210220032951-036812b2e83c/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
golang.org/x/sync v0.17.0 h1:l60nONMj9l5drqw6jlhIELNv9I0A4OFgRsG9k2oT9Ug=
golang.org/x/sync v0.17.0/go.mod h1:9KTHXmSnoGruLpwFjVSX0lNNA75CykiMECbovNTZqGI=
golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
golang.org/x/sys v0.0.0-20190412213103-97732733099d/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
golang.org/x/sys v0.0.0-20200930185726-fdedc70b468f/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
golang.org/x/sys v0.0.0-20201119102817-f84b799fce68/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
golang.org/x/sys v0.0.0-20210330210617-4fbd30eecc44/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
golang.org/x/sys v0.0.0-20210510120138-977fb7262007/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
golang.org/x/sys v0.32.0 h1:s77OFDvIQeibCmezSnk/q6iAfkdiQaJi4VzroCFrN20=
golang.org/x/sys v0.32.0/go.mod h1:BJP2sWEmIv4KK5OTEluFJCKSidICx8ciO85XgH3Ak8k=
golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo=
golang.org/x/term v0.30.0 h1:PQ39fJZ+mfadBm0y5WlL4vlM7Sx1Hgf13sMIY2+QS9Y=
golang.org/x/term v0.30.0/go.mod h1:NYYFdzHoI5wRh/h5tDMdMqCqPJZEuNqVR5xJLd/n67g=
golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
Expand All @@ -183,6 +195,7 @@ golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGm
golang.org/x/tools v0.0.0-20191119224855-298f0cb1881e/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo=
golang.org/x/tools v0.0.0-20200619180055-7c47624df98f/go.mod h1:EkVYQZoAsY45+roYkvgYkIh4xh/qjgUK9TdY2XT94GE=
golang.org/x/tools v0.0.0-20210106214847-113979e3529a/go.mod h1:emZCQorbCU4vsT4fOWvOPXz4eW1wZW4PmDk9uLelYpA=
golang.org/x/tools v0.1.1/go.mod h1:o0xws9oXOQQZyjljx8fwUC0k7L1pTE6eaCbjGeHmOkk=
golang.org/x/tools v0.31.0 h1:0EedkvKDbh+qistFTd0Bcwe/YLh4vHwWEkiI0toFIBU=
golang.org/x/tools v0.31.0/go.mod h1:naFTU+Cev749tSJRXJlna0T3WxKvb1kWEx15xA4SdmQ=
golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
Expand Down
30 changes: 21 additions & 9 deletions internal/controller/agentpool_controller_autoscaling.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,10 @@ import (
appv1alpha2 "github.com/hashicorp/hcp-terraform-operator/api/v1alpha2"
)

type AgentPoolControllerAutoscaling interface {
pendingWorkspaceRuns(ctx context.Context, ap *agentPoolInstance) (int32, error)
}

// userInteractionRunStatuses contains run statuses that require user interaction.
var userInteractionRunStatuses = map[tfc.RunStatus]struct{}{
tfc.RunCostEstimated: {},
Expand Down Expand Up @@ -61,9 +65,10 @@ func matchWildcardName(wildcard string, str string) bool {
}
}

// pendingWorkspaceRuns returns the number of workspaces with pending runs for a given agent pool.
// pendingWorkspaceRuns returns the number of agents needed to execute current pending runs for a given agent pool.
// If there are no plan-only runs in the list of current pending runs for a workspace this functoion returns the number of workspaces.
Copy link
Collaborator

@arybolovlev arybolovlev Nov 17, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Shall we stick to the number of pending runs only in the description? I would suggest either removing the second line or rephrasing it, as it’s confusing in its current form.

// This function is compatible with HCP Terraform and TFE version v202409-1 and later.
func pendingWorkspaceRuns(ctx context.Context, ap *agentPoolInstance) (int32, error) {
func (ap *agentPoolInstance) pendingWorkspaceRuns(ctx context.Context) (int32, error) {
runs := map[string]struct{}{}
awaitingUserInteractionRuns := map[string]int{} // Track runs awaiting user interaction by status for future metrics
listOpts := &tfc.RunListForOrganizationOptions{
Expand All @@ -74,8 +79,9 @@ func pendingWorkspaceRuns(ctx context.Context, ap *agentPoolInstance) (int32, er
PageNumber: initPageNumber,
},
}

planOnlyRunCount := 0
for {
ap.log.Info("Fetching runs for organization", "org", ap.instance.Spec.Organization, "page", listOpts.PageNumber)
Copy link
Collaborator

@arybolovlev arybolovlev Nov 17, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Let's remove this message; it seems to be redundant here.

runsList, err := ap.tfClient.Client.Runs.ListForOrganization(ctx, ap.instance.Spec.Organization, listOpts)
if err != nil {
return 0, err
Expand All @@ -87,6 +93,11 @@ func pendingWorkspaceRuns(ctx context.Context, ap *agentPoolInstance) (int32, er
awaitingUserInteractionRuns[string(run.Status)]++
continue
}
// Count plan-only runs separately so agents can scale up and execute runs parallely
if run.PlanOnly {
planOnlyRunCount++
continue
}
runs[run.Workspace.ID] = struct{}{}
}
if runsList.NextPage == 0 {
Expand All @@ -97,13 +108,14 @@ func pendingWorkspaceRuns(ctx context.Context, ap *agentPoolInstance) (int32, er

// TODO:
// Add metric(s) for runs awaiting user interaction

return int32(len(runs)), nil
agentsCount := len(runs) + planOnlyRunCount
ap.log.Info("Workspaces and plan-only runs count", "msg", fmt.Sprintf("Workspaces: %+v Plan-only runs: %d Total agents: %d", runs, planOnlyRunCount, agentsCount))
Copy link
Collaborator

@arybolovlev arybolovlev Nov 17, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

For all log messages, we keep a context. For instance, for the autoscaling part, all info and error messages should follow the following template:

ap.log.Info("Reconcile Agent Autoscaling", "msg", <MESSAGE>)
ap.log.Error(err, "Reconcile Agent Autoscaling", "msg", <MESSAGE>)

For simplicity, the messages could as simple as Pending runs apply/plan-only: %d/%d.

The summary is logged below once this function returns, and it currently logs the number of workspaces with pending runs. It will be worth rephrasing that message and keeping the name 'runs' rather than 'agents'. Agents are calculated later.

return int32(agentsCount), nil
}

// computeRequiredAgents is a legacy algorithm that is used to compute the number of agents needed.
// It is used when the TFE version is less than v202409-1.
func computeRequiredAgents(ctx context.Context, ap *agentPoolInstance) (int32, error) {
func (ap *agentPoolInstance) computeRequiredAgents(ctx context.Context) (int32, error) {
required := 0
// NOTE:
// - Two maps are used here to simplify target workspace searching by ID, name, and wildcard.
Expand Down Expand Up @@ -242,7 +254,7 @@ func (r *AgentPoolReconciler) reconcileAgentAutoscaling(ctx context.Context, ap

requiredAgents, err := func() (int32, error) {
if ap.tfClient.Client.IsCloud() {
return pendingWorkspaceRuns(ctx, ap)
return ap.pendingWorkspaceRuns(ctx)
}
tfeVersion := ap.tfClient.Client.RemoteTFEVersion()
useRunsEndpoint, err := validateTFEVersion(tfeVersion)
Expand All @@ -256,10 +268,10 @@ func (r *AgentPoolReconciler) reconcileAgentAutoscaling(ctx context.Context, ap
// It now allows retrieving a list of runs for the organization.
if useRunsEndpoint {
ap.log.Info("Reconcile Agent Autoscaling", "msg", fmt.Sprintf("Proceeding with the new algorithm based on the detected TFE version %s", tfeVersion))
return pendingWorkspaceRuns(ctx, ap)
return ap.pendingWorkspaceRuns(ctx)
}
ap.log.Info("Reconcile Agent Autoscaling", "msg", fmt.Sprintf("Proceeding with the legacy algorithm based to the detected TFE version %s", tfeVersion))
return computeRequiredAgents(ctx, ap)
return ap.computeRequiredAgents(ctx)
}()
if err != nil {
ap.log.Error(err, "Reconcile Agent Autoscaling", "msg", "Failed to get agents needed")
Expand Down
96 changes: 94 additions & 2 deletions internal/controller/agentpool_controller_autoscaling_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,17 +4,25 @@
package controller

import (
"context"
"errors"
"fmt"
"testing"
"time"

"github.com/go-logr/logr"
tfc "github.com/hashicorp/go-tfe"
"github.com/hashicorp/go-tfe/mocks"
Copy link
Collaborator

@arybolovlev arybolovlev Nov 17, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This confuses me. Why do we need to import mocks from the go-tfe package?

appv1alpha2 "github.com/hashicorp/hcp-terraform-operator/api/v1alpha2"
"github.com/hashicorp/hcp-terraform-operator/internal/pointer"
. "github.com/onsi/ginkgo/v2"
. "github.com/onsi/gomega"
gomock "go.uber.org/mock/gomock"
corev1 "k8s.io/api/core/v1"
"k8s.io/apimachinery/pkg/api/errors"
k8sapierrors "k8s.io/apimachinery/pkg/api/errors"
metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"

"github.com/stretchr/testify/assert"
)

var _ = Describe("Agent Pool controller", Ordered, func() {
Expand Down Expand Up @@ -77,7 +85,7 @@ var _ = Describe("Agent Pool controller", Ordered, func() {
Expect(k8sClient.Delete(ctx, instance)).To(Succeed())
Eventually(func() bool {
err := k8sClient.Get(ctx, namespacedName, instance)
return errors.IsNotFound(err)
return k8sapierrors.IsNotFound(err)
}).Should(BeTrue())
})

Expand Down Expand Up @@ -194,3 +202,87 @@ var _ = Describe("Agent Pool controller", Ordered, func() {
})
})
})

func TestPendingWorkspaceRuns(t *testing.T) {
tests := []struct {
name string
mockRuns []*tfc.Run
mockErr error
expectedCount int32
expectError bool
}{
{
name: "returns error from client",
mockErr: errors.New("api error"),
expectedCount: 0,
expectError: true,
},
{
name: "counts plan-only runs",
mockRuns: []*tfc.Run{
{ID: "run1", PlanOnly: true, Status: tfc.RunPlanning, Workspace: &tfc.Workspace{ID: "ws1"}},
{ID: "run2", PlanOnly: true, Status: tfc.RunPlanning, Workspace: &tfc.Workspace{ID: "ws2"}},
},
expectedCount: 2,
expectError: false,
},
{
name: "skips user interaction runs",
mockRuns: []*tfc.Run{
{ID: "run1", PlanOnly: false, Status: tfc.RunPlanned, Workspace: &tfc.Workspace{ID: "ws1"}},
{ID: "run2", PlanOnly: false, Status: tfc.RunPolicyOverride, Workspace: &tfc.Workspace{ID: "ws2"}},
},
expectedCount: 0,
expectError: false,
},
{
name: "counts normal pending runs",
mockRuns: []*tfc.Run{
{ID: "run1", PlanOnly: false, Status: tfc.RunPlanning, Workspace: &tfc.Workspace{ID: "ws1"}},
{ID: "run2", PlanOnly: false, Status: tfc.RunPlanning, Workspace: &tfc.Workspace{ID: "ws2"}},
},
expectedCount: 2,
expectError: false,
},
{
name: "mix of plan-only and normal runs",
mockRuns: []*tfc.Run{
{ID: "run1", PlanOnly: true, Status: tfc.RunPlanning, Workspace: &tfc.Workspace{ID: "ws1"}},
{ID: "run2", PlanOnly: false, Status: tfc.RunPlanning, Workspace: &tfc.Workspace{ID: "ws2"}},
},
expectedCount: 2,
expectError: false,
},
}

for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
ctrl := gomock.NewController(t)
defer ctrl.Finish()

mockRuns := mocks.NewMockRuns(ctrl)
mockRuns.EXPECT().
ListForOrganization(gomock.Any(), "test-org", gomock.Any()).
Return(&tfc.RunList{Items: tt.mockRuns, Pagination: &tfc.Pagination{NextPage: 0}}, tt.mockErr)

ap := &agentPoolInstance{
tfClient: HCPTerraformClient{Client: &tfc.Client{Runs: mockRuns}},
instance: appv1alpha2.AgentPool{
Spec: appv1alpha2.AgentPoolSpec{
Name: "test-pool",
Organization: "test-org",
},
},
log: logr.Logger{},
}

count, err := ap.pendingWorkspaceRuns(context.Background())
if tt.expectError {
assert.Error(t, err)
} else {
assert.NoError(t, err)
assert.Equal(t, tt.expectedCount, count)
}
})
}
}
50 changes: 50 additions & 0 deletions mocks/agentpool_controller_autoscaling_mocks.go
View file
Open in desktop

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

Loading