feat: add Timeweb Cloud provider for Woodpecker CI autoscaler

- Implement timewebcloud provider with DeployAgent, RemoveAgent, ListDeployedAgentNames
- Add minimal HTTP API client for Timeweb Cloud (create/list/delete servers)
- Register provider in main.go with CLI flags
- Add timeweb-list and timeweb-tester utilities
- Include Dockerfile and docker-compose.yml for deployment
- Update DEPLOY.md with verified OS/preset IDs
This commit is contained in:
2026-05-16 13:09:07 +03:00
commit 191cdd108f
34 changed files with 8651 additions and 0 deletions

400
engine/autoscaler.go Normal file
View File

@@ -0,0 +1,400 @@
package engine
import (
"context"
"fmt"
"math"
"regexp"
"strings"
"time"
"github.com/rs/zerolog/log"
"go.woodpecker-ci.org/autoscaler/config"
"go.woodpecker-ci.org/autoscaler/engine/types"
"go.woodpecker-ci.org/autoscaler/server"
"go.woodpecker-ci.org/autoscaler/utils"
"go.woodpecker-ci.org/woodpecker/v3/woodpecker-go/woodpecker"
)
type Autoscaler struct {
client server.Client
agents []*woodpecker.Agent
config *config.Config
provider types.Provider
}
// NewAutoscaler creates a new Autoscaler instance.
// It takes in a Provider, Client and Config, and returns a configured
// Autoscaler struct.
func NewAutoscaler(p types.Provider, client server.Client, config *config.Config) Autoscaler {
return Autoscaler{
provider: p,
client: client,
config: config,
}
}
func (a *Autoscaler) loadAgents(_ context.Context) error {
a.agents = []*woodpecker.Agent{}
agents, err := a.client.AgentList()
if err != nil {
return fmt.Errorf("client.AgentList: %w", err)
}
r, err := regexp.Compile(fmt.Sprintf("pool-%s-agent-.*?", a.config.PoolID))
if err != nil {
return fmt.Errorf("could not create regex matcher for agent names by pool ID: %w", err)
}
for _, agent := range agents {
if r.MatchString(agent.Name) {
a.agents = append(a.agents, agent)
}
}
return nil
}
func (a *Autoscaler) getPoolAgents(excludeNoSchedule bool) []*woodpecker.Agent {
agents := make([]*woodpecker.Agent, 0)
for _, agent := range a.agents {
if excludeNoSchedule && agent.NoSchedule {
continue
}
agents = append(agents, agent)
}
return agents
}
func (a *Autoscaler) createAgents(ctx context.Context, amount int) error {
suffixLength := 4
reactivatedAgents := 0
// try to re-activate agents that are in no-schedule state
for i := 0; i < amount; i++ {
for _, agent := range a.agents {
if agent.NoSchedule {
log.Info().Str("agent", agent.Name).Msg("reactivate agent")
agent.NoSchedule = false
_, err := a.client.AgentUpdate(agent)
if err != nil {
return fmt.Errorf("client.AgentUpdate: %w", err)
}
reactivatedAgents++
}
}
}
// create new agents
for i := 0; i < amount-reactivatedAgents; i++ {
agent, err := a.client.AgentCreate(&woodpecker.Agent{
Name: fmt.Sprintf("pool-%s-agent-%s", a.config.PoolID, utils.RandomString(suffixLength)),
})
if err != nil {
return fmt.Errorf("client.AgentCreate: %w", err)
}
log.Info().Str("agent", agent.Name).Msg("deploying agent")
err = a.provider.DeployAgent(ctx, agent)
if err != nil {
return fmt.Errorf("types.DeployAgent: %w", err)
}
a.agents = append(a.agents, agent)
}
return nil
}
func (a *Autoscaler) drainAgents(_ context.Context, amount int) error {
for i := 0; i < amount; i++ {
for _, agent := range a.agents {
// agent is already marked for draining
if agent.NoSchedule {
continue
}
// agent has recently done work => not ready for draining
if time.Since(time.Unix(agent.LastWork, 0)) < a.config.AgentIdleTimeout {
continue
}
// agent has never contacted the server => not ready for draining
if agent.LastContact == 0 {
continue
}
log.Info().Str("agent", agent.Name).Msg("drain agent")
agent.NoSchedule = true
_, err := a.client.AgentUpdate(agent)
if err != nil {
return fmt.Errorf("client.AgentUpdate: %w", err)
}
break
}
}
return nil
}
func (a *Autoscaler) isAgentIdle(agent *woodpecker.Agent) (bool, error) {
tasks, err := a.client.AgentTasksList(agent.ID)
if err != nil {
return false, fmt.Errorf("client.AgentTasksList: %w", err)
}
// agent still has tasks => not idle
if len(tasks) > 0 {
return false, nil
}
// agent has done work recently => not idle
if time.Since(time.Unix(agent.LastWork, 0)) < a.config.AgentIdleTimeout {
return false, nil
}
return true, nil
}
func (a *Autoscaler) removeAgent(ctx context.Context, agent *woodpecker.Agent, reason string) error {
isIdle, err := a.isAgentIdle(agent)
if err != nil {
return err
}
if !isIdle {
log.Info().Str("agent", agent.Name).Msg("agent is still processing workload")
return nil
}
log.Info().Str("agent", agent.Name).Str("reason", reason).Msgf("removing agent")
err = a.provider.RemoveAgent(ctx, agent)
if err != nil {
return err
}
err = a.client.AgentDelete(agent.ID)
if err != nil {
return fmt.Errorf("client.AgentDelete: %w", err)
}
filteredAgents := make([]*woodpecker.Agent, 0)
for _, a := range a.agents {
if a.ID != agent.ID {
filteredAgents = append(filteredAgents, a)
}
}
a.agents = filteredAgents
return nil
}
func (a *Autoscaler) removeDrainedAgents(ctx context.Context) error {
for _, agent := range a.getPoolAgents(false) {
if !agent.NoSchedule {
continue
}
err := a.removeAgent(ctx, agent, "was drained")
if err != nil {
return err
}
}
return nil
}
func (a *Autoscaler) cleanupDanglingAgents(ctx context.Context) error {
woodpeckerAgents := a.getPoolAgents(false)
providerAgentNames, err := a.provider.ListDeployedAgentNames(ctx)
if err != nil {
return err
}
// remove agents that are not in the woodpecker agent list anymore
for _, agentName := range providerAgentNames {
found := false
for _, agent := range woodpeckerAgents {
if agent.Name == agentName {
found = true
break
}
}
if !found {
log.Info().Str("agent", agentName).Str("reason", "not found on woodpecker").Msg("remove agent")
if err := a.provider.RemoveAgent(ctx, &woodpecker.Agent{Name: agentName}); err != nil {
return fmt.Errorf("types.RemoveAgent: %w", err)
}
// remove agent from providerAgentNames
_providerAgentNames := make([]string, 0)
for _, a := range providerAgentNames {
if a != agentName {
_providerAgentNames = append(_providerAgentNames, a)
}
}
providerAgentNames = _providerAgentNames
}
}
// remove agents that do not exist on the provider anymore
for _, agent := range woodpeckerAgents {
found := false
for _, agentName := range providerAgentNames {
if agent.Name == agentName {
found = true
break
}
}
if !found {
log.Info().Str("agent", agent.Name).Str("reason", "not found on provider").Msg("remove agent")
if err = a.client.AgentDelete(agent.ID); err != nil {
return fmt.Errorf("client.AgentDelete: %w", err)
}
// remove agent from woodpeckerAgents
_woodpeckerAgents := make([]*woodpecker.Agent, 0)
for _, a := range a.agents {
if a.Name != agent.Name {
woodpeckerAgents = append(woodpeckerAgents, a)
}
}
a.agents = _woodpeckerAgents
}
}
return nil
}
func (a *Autoscaler) cleanupStaleAgents(ctx context.Context) error {
// remove agents that haven't contacted the server for a while (including agents that never contacted the server)
for _, agent := range a.getPoolAgents(false) {
if agent.NoSchedule {
continue
}
lastContact := agent.LastContact
// if agent has never contacted the server, use the creation time
if lastContact == 0 {
lastContact = agent.Created
}
if time.Since(time.Unix(lastContact, 0)) > a.config.AgentInactivityTimeout {
err := a.removeAgent(ctx, agent, "hasn't connected to the server for a while")
if err != nil {
return err
}
}
}
return nil
}
func (a *Autoscaler) getQueueInfo(_ context.Context) (freeTasks, runningTasks, pendingTasks int, err error) {
queueInfo, err := a.client.QueueInfo()
if err != nil {
return 0, 0, 0, fmt.Errorf("error from QueueInfo: %s", err.Error())
}
if a.config.FilterLabels == "" {
return queueInfo.Stats.Workers, queueInfo.Stats.Running, queueInfo.Stats.Pending, nil
}
labelFilterKey, labelFilterValue, ok := strings.Cut(a.config.FilterLabels, "=")
if !ok {
return 0, 0, 0, fmt.Errorf("invalid labels filter: %s", a.config.FilterLabels)
}
running := countTasksByLabel(queueInfo.Running, labelFilterKey, labelFilterValue)
pending := countTasksByLabel(queueInfo.Pending, labelFilterKey, labelFilterValue)
return queueInfo.Stats.Workers, running, pending, nil
}
func (a *Autoscaler) calcAgents(ctx context.Context) (float64, error) {
freeTasks, runningTasks, pendingTasks, err := a.getQueueInfo(ctx)
if err != nil {
return 0, err
}
log.Debug().Msgf("queue info: freeTasks = %v runningTasks = %v pendingTasks = %v", freeTasks, runningTasks, pendingTasks)
availableAgents := math.Ceil(float64(freeTasks+runningTasks) / float64((a.config.WorkflowsPerAgent)))
reqAgents := math.Ceil(float64(pendingTasks+runningTasks) / float64(a.config.WorkflowsPerAgent))
availablePoolAgents := len(a.getPoolAgents(true))
maxUp := float64(a.config.MaxAgents - availablePoolAgents)
maxDown := float64(availablePoolAgents - a.config.MinAgents)
reqPoolAgents := math.Ceil(reqAgents - (availableAgents + float64(availablePoolAgents)))
reqPoolAgents = math.Max(reqPoolAgents, -maxDown)
reqPoolAgents = math.Min(reqPoolAgents, maxUp)
log.Debug().Msgf("capacity info: agents = %v/%v pool = %v/%v limits = %v/%v", availableAgents, reqAgents, availablePoolAgents, reqPoolAgents, maxUp, maxDown)
return reqPoolAgents, nil
}
// Reconcile periodically checks the status of the agent pool and adjusts it to match
// the desired capacity based on the current queue state.
func (a *Autoscaler) Reconcile(ctx context.Context) error {
if err := a.loadAgents(ctx); err != nil {
return fmt.Errorf("loading agents failed: %w", err)
}
reqPoolAgents, err := a.calcAgents(ctx)
if err != nil {
return fmt.Errorf("calculating agents failed: %w", err)
}
if reqPoolAgents > 0 {
num := int(math.Abs(reqPoolAgents))
log.Debug().Msgf("starting %d additional agents", num)
if err := a.createAgents(ctx, num); err != nil {
return fmt.Errorf("creating agents failed: %w", err)
}
}
if reqPoolAgents < 0 {
num := int(math.Abs(reqPoolAgents))
log.Debug().Msgf("checking %d agents if ready for draining", num)
if err := a.drainAgents(ctx, num); err != nil {
return fmt.Errorf("draining agents failed: %w", err)
}
}
// cleanup agents that are only present at the provider or woodpecker
if err := a.cleanupDanglingAgents(ctx); err != nil {
return fmt.Errorf("cleaning up dangling agents failed: %w", err)
}
// cleanup agents that haven't contacted the server for a while
if err := a.cleanupStaleAgents(ctx); err != nil {
return fmt.Errorf("cleaning up stale agents failed: %w", err)
}
// remove agents that are drained
if err := a.removeDrainedAgents(ctx); err != nil {
return fmt.Errorf("removing drained agents failed: %w", err)
}
return nil
}
func countTasksByLabel(jobs []woodpecker.Task, labelKey, labelValue string) int {
count := 0
for _, job := range jobs {
val, exists := job.Labels[labelKey]
if exists && val == labelValue {
count++
}
}
return count
}

546
engine/autoscaler_test.go Normal file
View File

@@ -0,0 +1,546 @@
package engine
import (
"testing"
"time"
"github.com/rs/zerolog"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/mock"
"go.woodpecker-ci.org/autoscaler/config"
mocks_provider "go.woodpecker-ci.org/autoscaler/engine/types/mocks"
mocks_server "go.woodpecker-ci.org/autoscaler/server/mocks"
"go.woodpecker-ci.org/woodpecker/v3/woodpecker-go/woodpecker"
)
type MockClient struct {
workers int
running int
pending int
waitingOnDeps int
woodpecker.Client
}
func (m MockClient) QueueInfo() (*woodpecker.Info, error) {
info := &woodpecker.Info{}
info.Stats.Workers = m.workers
info.Stats.Running = m.running
info.Stats.Pending = m.pending
info.Stats.WaitingOnDeps = m.waitingOnDeps
info.Pending = []woodpecker.Task{
{
Labels: map[string]string{
"arch": "amd64",
},
},
}
info.Running = []woodpecker.Task{
{
Labels: map[string]string{
"arch": "amd64",
},
},
}
return info, nil
}
func Test_calcAgents(t *testing.T) {
zerolog.SetGlobalLevel(zerolog.ErrorLevel)
t.Run("should create new agent (MinAgents)", func(t *testing.T) {
autoscaler := Autoscaler{client: &MockClient{
pending: 0,
}, config: &config.Config{
WorkflowsPerAgent: 1,
MaxAgents: 2,
MinAgents: 1,
}}
value, _ := autoscaler.calcAgents(t.Context())
assert.Equal(t, float64(1), value)
})
t.Run("should create single agent", func(t *testing.T) {
autoscaler := Autoscaler{client: &MockClient{
pending: 2,
}, config: &config.Config{
WorkflowsPerAgent: 5,
MaxAgents: 3,
}}
value, _ := autoscaler.calcAgents(t.Context())
assert.Equal(t, float64(1), value)
})
t.Run("should create multiple agents", func(t *testing.T) {
autoscaler := Autoscaler{client: &MockClient{
pending: 6,
}, config: &config.Config{
WorkflowsPerAgent: 5,
MaxAgents: 3,
}}
value, _ := autoscaler.calcAgents(t.Context())
assert.Equal(t, float64(2), value)
})
t.Run("should create new agent (MaxAgents)", func(t *testing.T) {
autoscaler := Autoscaler{client: &MockClient{
pending: 2,
}, config: &config.Config{
WorkflowsPerAgent: 1,
MaxAgents: 2,
}, agents: []*woodpecker.Agent{
{Name: "pool-1-agent-1234"},
}}
value, _ := autoscaler.calcAgents(t.Context())
assert.Equal(t, float64(1), value)
})
t.Run("should not create new agent (availableAgents)", func(t *testing.T) {
autoscaler := Autoscaler{client: &MockClient{
workers: 2,
pending: 2,
}, config: &config.Config{
WorkflowsPerAgent: 1,
MaxAgents: 2,
}}
value, _ := autoscaler.calcAgents(t.Context())
assert.Equal(t, float64(0), value)
})
}
func Test_getQueueInfo(t *testing.T) {
zerolog.SetGlobalLevel(zerolog.ErrorLevel)
t.Run("should not filter", func(t *testing.T) {
autoscaler := Autoscaler{
client: &MockClient{
pending: 2,
},
config: &config.Config{},
}
free, running, pending, _ := autoscaler.getQueueInfo(t.Context())
assert.Equal(t, 0, free)
assert.Equal(t, 0, running)
assert.Equal(t, 2, pending)
})
t.Run("should filter one by label", func(t *testing.T) {
autoscaler := Autoscaler{
client: &MockClient{
pending: 2,
},
config: &config.Config{
FilterLabels: "arch=amd64",
},
}
free, running, pending, _ := autoscaler.getQueueInfo(t.Context())
assert.Equal(t, 0, free)
assert.Equal(t, 1, running)
assert.Equal(t, 1, pending)
})
t.Run("should filter all by label", func(t *testing.T) {
autoscaler := Autoscaler{
client: &MockClient{
pending: 2,
},
config: &config.Config{
FilterLabels: "arch=arm64",
},
}
free, running, pending, _ := autoscaler.getQueueInfo(t.Context())
assert.Equal(t, 0, free)
assert.Equal(t, 0, running)
assert.Equal(t, 0, pending)
})
}
func Test_getPoolAgents(t *testing.T) {
autoscaler := Autoscaler{
agents: []*woodpecker.Agent{
{ID: 1, Name: "pool-1-agent-1", NoSchedule: false},
{ID: 2, Name: "pool-1-agent-2", NoSchedule: true},
{ID: 3, Name: "pool-1-agent-3", NoSchedule: false},
},
}
agents := autoscaler.getPoolAgents(false)
assert.Equal(t, 3, len(agents))
agents = autoscaler.getPoolAgents(true)
assert.Equal(t, 2, len(agents))
}
func Test_createAgents(t *testing.T) {
zerolog.SetGlobalLevel(zerolog.ErrorLevel)
t.Run("should create a new agent", func(t *testing.T) {
ctx := t.Context()
client := mocks_server.NewMockClient(t)
provider := mocks_provider.NewMockProvider(t)
autoscaler := Autoscaler{
client: client,
provider: provider,
config: &config.Config{
PoolID: "1",
},
}
client.On("AgentCreate", mock.Anything).Return(&woodpecker.Agent{Name: "pool-1-agent-1"}, nil)
provider.On("DeployAgent", ctx, mock.Anything).Return(nil)
err := autoscaler.createAgents(ctx, 1)
assert.NoError(t, err)
})
t.Run("should reuse an no-schedule agent first before creating a new one", func(t *testing.T) {
ctx := t.Context()
client := mocks_server.NewMockClient(t)
provider := mocks_provider.NewMockProvider(t)
autoscaler := Autoscaler{
client: client,
provider: provider,
agents: []*woodpecker.Agent{
{
ID: 1,
NoSchedule: true,
},
},
config: &config.Config{
PoolID: "1",
},
}
client.On("AgentUpdate", mock.MatchedBy(func(agent *woodpecker.Agent) bool {
return agent.ID == 1 && agent.NoSchedule == false
})).Return(nil, nil)
client.On("AgentCreate", mock.Anything).Return(&woodpecker.Agent{Name: "pool-1-agent-1"}, nil)
provider.On("DeployAgent", ctx, mock.Anything).Return(nil)
err := autoscaler.createAgents(ctx, 2)
assert.NoError(t, err)
})
}
func Test_cleanupDanglingAgents(t *testing.T) {
t.Run("should remove agent that is only present on woodpecker (not provider)", func(t *testing.T) {
ctx := t.Context()
client := mocks_server.NewMockClient(t)
provider := mocks_provider.NewMockProvider(t)
autoscaler := Autoscaler{
agents: []*woodpecker.Agent{
{ID: 1, Name: "pool-1-agent-1", NoSchedule: false},
},
provider: provider,
client: client,
}
provider.On("ListDeployedAgentNames", mock.Anything).Return(nil, nil)
client.On("AgentDelete", int64(1)).Return(nil)
err := autoscaler.cleanupDanglingAgents(ctx)
assert.NoError(t, err)
})
t.Run("should remove agent that is only present on provider (not woodpecker)", func(t *testing.T) {
ctx := t.Context()
client := mocks_server.NewMockClient(t)
provider := mocks_provider.NewMockProvider(t)
autoscaler := Autoscaler{
agents: []*woodpecker.Agent{
{ID: 1, Name: "pool-1-agent-1", NoSchedule: false},
},
provider: provider,
client: client,
}
provider.On("ListDeployedAgentNames", mock.Anything).Return([]string{"pool-1-agent-1", "pool-1-agent-2"}, nil)
provider.On("RemoveAgent", mock.Anything, mock.MatchedBy(func(agent *woodpecker.Agent) bool {
return agent.Name == "pool-1-agent-2"
})).Return(nil)
err := autoscaler.cleanupDanglingAgents(ctx)
assert.NoError(t, err)
})
}
func Test_cleanupStaleAgents(t *testing.T) {
t.Run("should remove agent that never connected (last contact = 0) in over 15 minutes", func(t *testing.T) {
ctx := t.Context()
client := mocks_server.NewMockClient(t)
provider := mocks_provider.NewMockProvider(t)
autoscaler := Autoscaler{
agents: []*woodpecker.Agent{
{
ID: 1,
Name: "active agent",
NoSchedule: false,
Created: time.Now().Add(-time.Minute * 20).Unix(), // created 20 minutes ago
LastContact: time.Now().Add(-time.Minute * 5).Unix(), // last contact 5 minutes ago
},
{
ID: 2,
Name: "never contacted agent",
NoSchedule: false,
Created: time.Now().Add(-time.Minute * 20).Unix(), // created 20 minutes ago
LastContact: 0, // never contacted
},
},
provider: provider,
client: client,
config: &config.Config{
AgentInactivityTimeout: time.Minute * 15,
},
}
client.On("AgentTasksList", int64(2)).Return(nil, nil)
client.On("AgentDelete", int64(2)).Return(nil)
provider.On("RemoveAgent", mock.Anything, mock.MatchedBy(func(agent *woodpecker.Agent) bool {
return agent.ID == 2
})).Return(nil)
err := autoscaler.cleanupStaleAgents(ctx)
assert.NoError(t, err)
})
t.Run("should remove agent that has lost connection for more than 15 minutes", func(t *testing.T) {
ctx := t.Context()
client := mocks_server.NewMockClient(t)
provider := mocks_provider.NewMockProvider(t)
autoscaler := Autoscaler{
agents: []*woodpecker.Agent{
{
ID: 1,
Name: "active agent",
NoSchedule: false,
Created: time.Now().Add(-time.Minute * 20).Unix(), // created 20 minutes ago
LastContact: time.Now().Add(-time.Minute * 5).Unix(), // last contact 5 minutes ago
},
{
ID: 2,
Name: "stale agent",
NoSchedule: false,
Created: time.Now().Add(-time.Minute * 20).Unix(), // created 20 minutes ago
LastContact: time.Now().Add(-time.Minute * 20).Unix(), // last contact 20 minutes ago
},
},
provider: provider,
client: client,
config: &config.Config{
AgentInactivityTimeout: time.Minute * 15,
},
}
client.On("AgentTasksList", int64(2)).Return(nil, nil)
client.On("AgentDelete", int64(2)).Return(nil)
provider.On("RemoveAgent", mock.Anything, mock.MatchedBy(func(agent *woodpecker.Agent) bool {
return agent.ID == 2
})).Return(nil)
err := autoscaler.cleanupStaleAgents(ctx)
assert.NoError(t, err)
})
}
func Test_isAgentIdle(t *testing.T) {
t.Run("should return false if agent has tasks", func(t *testing.T) {
client := mocks_server.NewMockClient(t)
autoscaler := Autoscaler{
client: client,
config: &config.Config{
AgentIdleTimeout: time.Minute * 15,
},
}
client.On("AgentTasksList", int64(1)).Return([]*woodpecker.Task{
{ID: "1"},
}, nil)
idle, err := autoscaler.isAgentIdle(&woodpecker.Agent{
ID: 1,
Name: "pool-1-agent-1",
NoSchedule: false,
})
assert.NoError(t, err)
assert.False(t, idle)
})
t.Run("should return false if agent has done work recently", func(t *testing.T) {
client := mocks_server.NewMockClient(t)
autoscaler := Autoscaler{
client: client,
config: &config.Config{
AgentIdleTimeout: time.Minute * 15,
},
}
client.On("AgentTasksList", int64(1)).Return(nil, nil)
idle, err := autoscaler.isAgentIdle(&woodpecker.Agent{
ID: 1,
Name: "pool-1-agent-1",
NoSchedule: false,
LastWork: time.Now().Add(-time.Minute * 10).Unix(),
})
assert.NoError(t, err)
assert.False(t, idle)
})
t.Run("should return true if agent is idle", func(t *testing.T) {
client := mocks_server.NewMockClient(t)
autoscaler := Autoscaler{
client: client,
config: &config.Config{
AgentIdleTimeout: time.Minute * 15,
},
}
client.On("AgentTasksList", int64(1)).Return(nil, nil) // no tasks
idle, err := autoscaler.isAgentIdle(&woodpecker.Agent{
ID: 1,
Name: "pool-1-agent-1",
NoSchedule: false,
LastWork: time.Now().Add(-time.Minute * 20).Unix(), // last work 20 minutes ago
})
assert.NoError(t, err)
assert.True(t, idle)
})
}
func Test_drainAgents(t *testing.T) {
t.Run("should drain agents and skip no-schedule ones", func(t *testing.T) {
ctx := t.Context()
client := mocks_server.NewMockClient(t)
provider := mocks_provider.NewMockProvider(t)
autoscaler := Autoscaler{
agents: []*woodpecker.Agent{
{ID: 1, Name: "pool-1-agent-1", NoSchedule: false, LastContact: time.Now().Add(-time.Minute * 2).Unix()},
{ID: 2, Name: "pool-1-agent-2", NoSchedule: true, LastContact: time.Now().Add(-time.Minute * 2).Unix()},
{ID: 3, Name: "pool-1-agent-3", NoSchedule: true, LastContact: time.Now().Add(-time.Minute * 2).Unix()},
{ID: 4, Name: "pool-1-agent-4", NoSchedule: false, LastContact: time.Now().Add(-time.Minute * 2).Unix()},
},
provider: provider,
client: client,
config: &config.Config{
AgentIdleTimeout: time.Minute * 15,
},
}
client.On("AgentUpdate", mock.MatchedBy(func(agent *woodpecker.Agent) bool {
return (agent.ID == 1 || agent.ID == 4) && agent.NoSchedule == true
})).Return(nil, nil)
err := autoscaler.drainAgents(ctx, 2)
assert.NoError(t, err)
assert.True(t, autoscaler.agents[0].NoSchedule)
assert.True(t, autoscaler.agents[3].NoSchedule)
})
t.Run("should not remove an agent that never connected", func(t *testing.T) {
ctx := t.Context()
client := mocks_server.NewMockClient(t)
provider := mocks_provider.NewMockProvider(t)
autoscaler := Autoscaler{
agents: []*woodpecker.Agent{
{ID: 1, Name: "pool-1-agent-1", NoSchedule: false, LastContact: 0},
},
provider: provider,
client: client,
config: &config.Config{
AgentIdleTimeout: time.Minute * 15,
},
}
err := autoscaler.drainAgents(ctx, 1)
assert.NoError(t, err)
assert.False(t, autoscaler.agents[0].NoSchedule)
})
t.Run("should not remove an agent that has recently done some work", func(t *testing.T) {
ctx := t.Context()
client := mocks_server.NewMockClient(t)
provider := mocks_provider.NewMockProvider(t)
autoscaler := Autoscaler{
agents: []*woodpecker.Agent{
{
ID: 1,
Name: "pool-1-agent-1",
NoSchedule: false,
LastContact: time.Now().Add(-time.Minute * 2).Unix(), // last contact 2 minutes ago
LastWork: time.Now().Add(-time.Minute * 5).Unix(), // last work 5 minutes ago
},
},
provider: provider,
client: client,
config: &config.Config{
AgentIdleTimeout: time.Minute * 15,
},
}
err := autoscaler.drainAgents(ctx, 1)
assert.NoError(t, err)
assert.False(t, autoscaler.agents[0].NoSchedule)
})
}
func Test_removeDrainedAgents(t *testing.T) {
t.Run("should remove agent", func(t *testing.T) {
ctx := t.Context()
client := mocks_server.NewMockClient(t)
provider := mocks_provider.NewMockProvider(t)
autoscaler := Autoscaler{
agents: []*woodpecker.Agent{
{ID: 1, Name: "pool-1-agent-1", NoSchedule: false},
{ID: 2, Name: "pool-1-agent-2", NoSchedule: true},
{ID: 3, Name: "pool-1-agent-3", NoSchedule: false},
},
provider: provider,
client: client,
config: &config.Config{
AgentIdleTimeout: time.Minute * 15,
},
}
client.On("AgentTasksList", int64(2)).Return(nil, nil)
provider.On("RemoveAgent", mock.Anything, mock.MatchedBy(func(agent *woodpecker.Agent) bool {
return agent.ID == 2
})).Return(nil)
client.On("AgentDelete", int64(2)).Return(nil)
err := autoscaler.removeDrainedAgents(ctx)
assert.NoError(t, err)
})
t.Run("should not remove agent with tasks", func(t *testing.T) {
ctx := t.Context()
client := mocks_server.NewMockClient(t)
provider := mocks_provider.NewMockProvider(t)
autoscaler := Autoscaler{
agents: []*woodpecker.Agent{
{ID: 1, Name: "pool-1-agent-1", NoSchedule: false},
{ID: 2, Name: "pool-1-agent-2", NoSchedule: true},
{ID: 3, Name: "pool-1-agent-3", NoSchedule: false},
},
provider: provider,
client: client,
}
client.On("AgentTasksList", int64(2)).Return([]*woodpecker.Task{
{ID: "1"},
}, nil)
err := autoscaler.removeDrainedAgents(ctx)
assert.NoError(t, err)
})
}

9
engine/const.go Normal file
View File

@@ -0,0 +1,9 @@
package engine
import "fmt"
var (
LabelPrefix = "wp.autoscaler/"
LabelPool = fmt.Sprintf("%spool", LabelPrefix)
LabelImage = fmt.Sprintf("%simage", LabelPrefix)
)

View File

@@ -0,0 +1,116 @@
package cloudinit
import (
"bytes"
"fmt"
"strings"
"text/template"
"go.woodpecker-ci.org/autoscaler/config"
"go.woodpecker-ci.org/woodpecker/v3/woodpecker-go/woodpecker"
)
// RenderUserDataTemplate renders the user data template for an Agent
// using the provided configuration.
func RenderUserDataTemplate(config *config.Config, agent *woodpecker.Agent, tmpl *template.Template) (string, error) {
var err error
switch {
case tmpl != nil:
case config.UserData != "":
tmpl, err = template.New("user-data").Parse(config.UserData)
default:
tmpl, err = template.New("user-data").Parse(CloudInitUserDataUbuntuDefault)
}
if err != nil {
return "", fmt.Errorf("template.New.Parse %w", err)
}
params := struct {
Image string
Environment map[string]string
}{
Image: config.Image,
Environment: map[string]string{
"WOODPECKER_SERVER": config.GRPCAddress,
"WOODPECKER_AGENT_SECRET": agent.Token,
"WOODPECKER_MAX_WORKFLOWS": fmt.Sprintf("%d", config.WorkflowsPerAgent),
},
}
if config.GRPCSecure {
params.Environment["WOODPECKER_GRPC_SECURE"] = "true"
}
for key, value := range config.Environment {
params.Environment[key] = value
}
params.Environment["WOODPECKER_AGENT_LABELS"] = genExtraAgentLabels(config.ExtraAgentLabels)
var userData bytes.Buffer
if err := tmpl.Execute(&userData, params); err != nil {
return "", err
}
return userData.String(), nil
}
func genExtraAgentLabels(conf map[string]string) string {
out := make([]string, 0, len(conf))
for k, v := range conf {
out = append(out, fmt.Sprintf("%s=%s", k, v))
}
return strings.Join(out, ",")
}
// editorconfig-checker-disable
var CloudInitUserDataUbuntuDefault = `
#cloud-config
package_reboot_if_required: false
package_update: true
package_upgrade: false
groups:
- docker
system_info:
default_user:
groups: [ docker ]
apt:
sources:
docker.list:
keyid: 9DC858229FC7DD38854AE2D88D81803C0EBFCD88
keyserver: https://download.docker.com/linux/ubuntu/gpg
source: deb [signed-by=$KEY_FILE] https://download.docker.com/linux/ubuntu $RELEASE stable
packages:
- docker-ce
- docker-compose-plugin
- binfmt-support
- qemu-user-static
write_files:
- path: /root/docker-compose.yml
content: |
# docker-compose.yml
version: '3'
services:
woodpecker-agent:
image: {{ .Image }}
restart: always
volumes:
- /var/run/docker.sock:/var/run/docker.sock
environment:
{{- range $key, $value := .Environment }}
- {{ $key }}={{ $value }}
{{- end }}
runcmd:
- sh -xc "cd /root; docker compose up -d"
final_message: "The system is finally up, after $UPTIME seconds"
` // editorconfig-checker-enable

View File

@@ -0,0 +1,65 @@
package cloudinit_test
import (
"testing"
"text/template"
"github.com/stretchr/testify/assert"
"go.woodpecker-ci.org/autoscaler/config"
"go.woodpecker-ci.org/autoscaler/engine/inits/cloudinit"
"go.woodpecker-ci.org/woodpecker/v3/woodpecker-go/woodpecker"
)
var testUserDataStr = `
image: {{ .Image }}
environment:
{{- range $key, $value := .Environment }}
- {{ $key }}={{ $value }}
{{- end }}
`
var testUserDataTmpl = template.Must(template.New("test").Parse(testUserDataStr))
func TestRenderUserDataTemplate(t *testing.T) {
config := &config.Config{
Image: "test-image",
GRPCAddress: "test-address",
GRPCSecure: false,
Environment: map[string]string{
"FOO": "bar",
},
}
agent := &woodpecker.Agent{
Token: "test-token",
}
userData, err := cloudinit.RenderUserDataTemplate(config, agent, testUserDataTmpl)
assert.NoError(t, err)
assert.Contains(t, userData, "test-image")
assert.Contains(t, userData, "bar")
assert.Contains(t, userData, "WOODPECKER_SERVER=test-address")
assert.Contains(t, userData, "WOODPECKER_AGENT_SECRET=test-token")
}
func TestRenderUserDataTemplate_Secure(t *testing.T) {
config := &config.Config{
GRPCSecure: true,
}
agent := &woodpecker.Agent{}
userData, err := cloudinit.RenderUserDataTemplate(config, agent, testUserDataTmpl)
assert.NoError(t, err)
assert.Contains(t, userData, "WOODPECKER_GRPC_SECURE=true")
}
func TestRenderUserDataTemplate_Error(t *testing.T) {
config := &config.Config{}
agent := &woodpecker.Agent{}
tmpl := template.Must(template.New("test").Parse("{{.Missing}}"))
_, err := cloudinit.RenderUserDataTemplate(config, agent, tmpl)
assert.Error(t, err)
}

View File

@@ -0,0 +1,215 @@
// Code generated by mockery; DO NOT EDIT.
// github.com/vektra/mockery
// template: testify
package mocks
import (
"context"
mock "github.com/stretchr/testify/mock"
"go.woodpecker-ci.org/woodpecker/v3/woodpecker-go/woodpecker"
)
// NewMockProvider creates a new instance of MockProvider. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
// The first argument is typically a *testing.T value.
func NewMockProvider(t interface {
mock.TestingT
Cleanup(func())
}) *MockProvider {
mock := &MockProvider{}
mock.Mock.Test(t)
t.Cleanup(func() { mock.AssertExpectations(t) })
return mock
}
// MockProvider is an autogenerated mock type for the Provider type
type MockProvider struct {
mock.Mock
}
type MockProvider_Expecter struct {
mock *mock.Mock
}
func (_m *MockProvider) EXPECT() *MockProvider_Expecter {
return &MockProvider_Expecter{mock: &_m.Mock}
}
// DeployAgent provides a mock function for the type MockProvider
func (_mock *MockProvider) DeployAgent(context1 context.Context, agent *woodpecker.Agent) error {
ret := _mock.Called(context1, agent)
if len(ret) == 0 {
panic("no return value specified for DeployAgent")
}
var r0 error
if returnFunc, ok := ret.Get(0).(func(context.Context, *woodpecker.Agent) error); ok {
r0 = returnFunc(context1, agent)
} else {
r0 = ret.Error(0)
}
return r0
}
// MockProvider_DeployAgent_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'DeployAgent'
type MockProvider_DeployAgent_Call struct {
*mock.Call
}
// DeployAgent is a helper method to define mock.On call
// - context1 context.Context
// - agent *woodpecker.Agent
func (_e *MockProvider_Expecter) DeployAgent(context1 interface{}, agent interface{}) *MockProvider_DeployAgent_Call {
return &MockProvider_DeployAgent_Call{Call: _e.mock.On("DeployAgent", context1, agent)}
}
func (_c *MockProvider_DeployAgent_Call) Run(run func(context1 context.Context, agent *woodpecker.Agent)) *MockProvider_DeployAgent_Call {
_c.Call.Run(func(args mock.Arguments) {
var arg0 context.Context
if args[0] != nil {
arg0 = args[0].(context.Context)
}
var arg1 *woodpecker.Agent
if args[1] != nil {
arg1 = args[1].(*woodpecker.Agent)
}
run(
arg0,
arg1,
)
})
return _c
}
func (_c *MockProvider_DeployAgent_Call) Return(err error) *MockProvider_DeployAgent_Call {
_c.Call.Return(err)
return _c
}
func (_c *MockProvider_DeployAgent_Call) RunAndReturn(run func(context1 context.Context, agent *woodpecker.Agent) error) *MockProvider_DeployAgent_Call {
_c.Call.Return(run)
return _c
}
// ListDeployedAgentNames provides a mock function for the type MockProvider
func (_mock *MockProvider) ListDeployedAgentNames(context1 context.Context) ([]string, error) {
ret := _mock.Called(context1)
if len(ret) == 0 {
panic("no return value specified for ListDeployedAgentNames")
}
var r0 []string
var r1 error
if returnFunc, ok := ret.Get(0).(func(context.Context) ([]string, error)); ok {
return returnFunc(context1)
}
if returnFunc, ok := ret.Get(0).(func(context.Context) []string); ok {
r0 = returnFunc(context1)
} else {
if ret.Get(0) != nil {
r0 = ret.Get(0).([]string)
}
}
if returnFunc, ok := ret.Get(1).(func(context.Context) error); ok {
r1 = returnFunc(context1)
} else {
r1 = ret.Error(1)
}
return r0, r1
}
// MockProvider_ListDeployedAgentNames_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'ListDeployedAgentNames'
type MockProvider_ListDeployedAgentNames_Call struct {
*mock.Call
}
// ListDeployedAgentNames is a helper method to define mock.On call
// - context1 context.Context
func (_e *MockProvider_Expecter) ListDeployedAgentNames(context1 interface{}) *MockProvider_ListDeployedAgentNames_Call {
return &MockProvider_ListDeployedAgentNames_Call{Call: _e.mock.On("ListDeployedAgentNames", context1)}
}
func (_c *MockProvider_ListDeployedAgentNames_Call) Run(run func(context1 context.Context)) *MockProvider_ListDeployedAgentNames_Call {
_c.Call.Run(func(args mock.Arguments) {
var arg0 context.Context
if args[0] != nil {
arg0 = args[0].(context.Context)
}
run(
arg0,
)
})
return _c
}
func (_c *MockProvider_ListDeployedAgentNames_Call) Return(strings []string, err error) *MockProvider_ListDeployedAgentNames_Call {
_c.Call.Return(strings, err)
return _c
}
func (_c *MockProvider_ListDeployedAgentNames_Call) RunAndReturn(run func(context1 context.Context) ([]string, error)) *MockProvider_ListDeployedAgentNames_Call {
_c.Call.Return(run)
return _c
}
// RemoveAgent provides a mock function for the type MockProvider
func (_mock *MockProvider) RemoveAgent(context1 context.Context, agent *woodpecker.Agent) error {
ret := _mock.Called(context1, agent)
if len(ret) == 0 {
panic("no return value specified for RemoveAgent")
}
var r0 error
if returnFunc, ok := ret.Get(0).(func(context.Context, *woodpecker.Agent) error); ok {
r0 = returnFunc(context1, agent)
} else {
r0 = ret.Error(0)
}
return r0
}
// MockProvider_RemoveAgent_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'RemoveAgent'
type MockProvider_RemoveAgent_Call struct {
*mock.Call
}
// RemoveAgent is a helper method to define mock.On call
// - context1 context.Context
// - agent *woodpecker.Agent
func (_e *MockProvider_Expecter) RemoveAgent(context1 interface{}, agent interface{}) *MockProvider_RemoveAgent_Call {
return &MockProvider_RemoveAgent_Call{Call: _e.mock.On("RemoveAgent", context1, agent)}
}
func (_c *MockProvider_RemoveAgent_Call) Run(run func(context1 context.Context, agent *woodpecker.Agent)) *MockProvider_RemoveAgent_Call {
_c.Call.Run(func(args mock.Arguments) {
var arg0 context.Context
if args[0] != nil {
arg0 = args[0].(context.Context)
}
var arg1 *woodpecker.Agent
if args[1] != nil {
arg1 = args[1].(*woodpecker.Agent)
}
run(
arg0,
arg1,
)
})
return _c
}
func (_c *MockProvider_RemoveAgent_Call) Return(err error) *MockProvider_RemoveAgent_Call {
_c.Call.Return(err)
return _c
}
func (_c *MockProvider_RemoveAgent_Call) RunAndReturn(run func(context1 context.Context, agent *woodpecker.Agent) error) *MockProvider_RemoveAgent_Call {
_c.Call.Return(run)
return _c
}

13
engine/types/provider.go Normal file
View File

@@ -0,0 +1,13 @@
package types
import (
"context"
"go.woodpecker-ci.org/woodpecker/v3/woodpecker-go/woodpecker"
)
type Provider interface {
DeployAgent(context.Context, *woodpecker.Agent) error
RemoveAgent(context.Context, *woodpecker.Agent) error
ListDeployedAgentNames(context.Context) ([]string, error)
}