Use strong typing and allow atomic change of the max runners (#24)

* Use strong typing and allow atomic change of the max runners

* prepare type change in dockerscaleset
This commit is contained in:
Nikola Jokic
2025-11-26 12:10:39 +01:00
committed by GitHub
parent 6d1c317b90
commit 7442885c99
7 changed files with 51 additions and 61 deletions
+8 -12
View File
@@ -371,7 +371,7 @@ func (c *Client) GetRunnerScaleSet(ctx context.Context, runnerGroupID int, runne
return &runnerScaleSetList.RunnerScaleSets[0], nil
}
func (c *Client) GetRunnerScaleSetByID(ctx context.Context, runnerScaleSetID int) (*RunnerScaleSet, error) {
func (c *Client) GetRunnerScaleSetByID(ctx context.Context, runnerScaleSetID uint64) (*RunnerScaleSet, error) {
path := fmt.Sprintf("/%s/%d", scaleSetEndpoint, runnerScaleSetID)
req, err := c.newActionsServiceRequest(ctx, http.MethodGet, path, nil)
if err != nil {
@@ -540,7 +540,7 @@ func (c *Client) DeleteRunnerScaleSet(ctx context.Context, runnerScaleSetID int)
}
// GetMessage fetches a message from the runner scale set message queue.
func (c *Client) GetMessage(ctx context.Context, messageQueueURL, messageQueueAccessToken string, lastMessageID int64, maxCapacity int) (*RunnerScaleSetMessage, error) {
func (c *Client) GetMessage(ctx context.Context, messageQueueURL, messageQueueAccessToken string, lastMessageID uint64, maxCapacity uint32) (*RunnerScaleSetMessage, error) {
u, err := url.Parse(messageQueueURL)
if err != nil {
return nil, fmt.Errorf("failed to parse message queue url: %w", err)
@@ -548,14 +548,10 @@ func (c *Client) GetMessage(ctx context.Context, messageQueueURL, messageQueueAc
if lastMessageID > 0 {
q := u.Query()
q.Set("lastMessageId", strconv.FormatInt(lastMessageID, 10))
q.Set("lastMessageId", strconv.FormatUint(lastMessageID, 10))
u.RawQuery = q.Encode()
}
if maxCapacity < 0 {
return nil, fmt.Errorf("maxCapacity must be greater than or equal to 0")
}
req, err := http.NewRequestWithContext(ctx, http.MethodGet, u.String(), nil)
if err != nil {
return nil, fmt.Errorf("failed to create new request with context: %w", err)
@@ -564,7 +560,7 @@ func (c *Client) GetMessage(ctx context.Context, messageQueueURL, messageQueueAc
req.Header.Set("Accept", "application/json; api-version=6.0-preview")
req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", messageQueueAccessToken))
req.Header.Set("User-Agent", c.userAgent.Load())
req.Header.Set(HeaderScaleSetMaxCapacity, strconv.Itoa(maxCapacity))
req.Header.Set(HeaderScaleSetMaxCapacity, strconv.Itoa(int(maxCapacity)))
resp, err := c.do(req)
if err != nil {
@@ -608,7 +604,7 @@ func (c *Client) GetMessage(ctx context.Context, messageQueueURL, messageQueueAc
return message, nil
}
func (c *Client) DeleteMessage(ctx context.Context, messageQueueURL, messageQueueAccessToken string, messageID int64) error {
func (c *Client) DeleteMessage(ctx context.Context, messageQueueURL, messageQueueAccessToken string, messageID uint64) error {
u, err := url.Parse(messageQueueURL)
if err != nil {
return fmt.Errorf("failed to parse message queue url: %w", err)
@@ -655,7 +651,7 @@ func (c *Client) DeleteMessage(ctx context.Context, messageQueueURL, messageQueu
}
}
func (c *Client) CreateMessageSession(ctx context.Context, runnerScaleSetID int, owner string) (*RunnerScaleSetSession, error) {
func (c *Client) CreateMessageSession(ctx context.Context, runnerScaleSetID uint64, owner string) (*RunnerScaleSetSession, error) {
path := fmt.Sprintf("/%s/%d/sessions", scaleSetEndpoint, runnerScaleSetID)
newSession := &RunnerScaleSetSession{
@@ -676,12 +672,12 @@ func (c *Client) CreateMessageSession(ctx context.Context, runnerScaleSetID int,
return createdSession, nil
}
func (c *Client) DeleteMessageSession(ctx context.Context, runnerScaleSetID int, sessionID uuid.UUID) error {
func (c *Client) DeleteMessageSession(ctx context.Context, runnerScaleSetID uint64, sessionID uuid.UUID) error {
path := fmt.Sprintf("/%s/%d/sessions/%s", scaleSetEndpoint, runnerScaleSetID, sessionID.String())
return c.doSessionRequest(ctx, http.MethodDelete, path, nil, http.StatusNoContent, nil)
}
func (c *Client) RefreshMessageSession(ctx context.Context, runnerScaleSetID int, sessionID uuid.UUID) (*RunnerScaleSetSession, error) {
func (c *Client) RefreshMessageSession(ctx context.Context, runnerScaleSetID uint64, sessionID uuid.UUID) (*RunnerScaleSetSession, error) {
path := fmt.Sprintf("/%s/%d/sessions/%s", scaleSetEndpoint, runnerScaleSetID, sessionID.String())
refreshedSession := &RunnerScaleSetSession{}
if err := c.doSessionRequest(ctx, http.MethodPatch, path, nil, http.StatusOK, refreshedSession); err != nil {
+1 -6
View File
@@ -429,7 +429,7 @@ func TestGetRunnerGroupByName(t *testing.T) {
}
t.Run("Get RunnerGroup by Name", func(t *testing.T) {
runnerGroupID := 1
var runnerGroupID uint64 = 1
runnerGroupName := "test-runner-group"
want := &RunnerGroup{
ID: runnerGroupID,
@@ -1204,11 +1204,6 @@ func TestGetMessage(t *testing.T) {
client, err := NewClient(server.configURLForOrg("my-org"), auth)
require.NoError(t, err)
_, err = client.GetMessage(ctx, server.URL, token, 0, -1)
require.Error(t, err)
// Ensure we don't send requests with negative capacity
assert.False(t, errors.Is(err, &ActionsError{}))
_, err = client.GetMessage(ctx, server.URL, token, 0, 0)
assert.Error(t, err)
var expectedErr *ActionsError
+2 -2
View File
@@ -12,8 +12,8 @@ import (
type Config struct {
RegistrationURL string
MaxRunners int
MinRunners int
MaxRunners uint32
MinRunners uint32
ScaleSetName string
RunnerGroup string
GitHubApp scaleset.GitHubAppAuth
+5 -5
View File
@@ -19,8 +19,8 @@ import (
func init() {
flags := cmd.Flags()
flags.StringVar(&cfg.RegistrationURL, "url", "", "REQUIRED: URL where to register your scale set (e.g. https://github.com/org/repo)")
flags.IntVar(&cfg.MaxRunners, "max-runners", 10, "Maximum number of runners")
flags.IntVar(&cfg.MinRunners, "min-runners", 0, "Minimum number of runners")
flags.Uint32Var(&cfg.MaxRunners, "max-runners", 10, "Maximum number of runners")
flags.Uint32Var(&cfg.MinRunners, "min-runners", 0, "Minimum number of runners")
flags.StringVar(&cfg.ScaleSetName, "name", "", "REQUIRED: Name of your scale set")
flags.StringVar(&cfg.RunnerGroup, "runner-group", scaleset.DefaultRunnerGroup, "Name of the runner group your scale set should belong to")
flags.StringVar(&cfg.GitHubApp.ClientID, "app-client-id", "", "GitHub App client id")
@@ -140,8 +140,8 @@ func run(ctx context.Context, c Config) error {
logger.Info("Initializing listener")
listener, err := listener.New(scalesetClient, listener.Config{
ScaleSetID: scaleSet.ID,
MinRunners: c.MinRunners,
MaxRunners: c.MaxRunners,
MinRunners: int(c.MinRunners),
MaxRunners: int(c.MaxRunners),
Logger: logger.WithGroup("listener"),
})
if err != nil {
@@ -158,7 +158,7 @@ func run(ctx context.Context, c Config) error {
minRunners: c.MinRunners,
dockerClient: dockerClient,
scalesetClient: scalesetClient,
scaleSetID: scaleSet.ID,
scaleSetID: uint64(scaleSet.ID),
}
defer scaler.shutdown(context.WithoutCancel(ctx))
+3 -3
View File
@@ -16,16 +16,16 @@ import (
type Scaler struct {
runners runnerState
runnerImage string
scaleSetID int
scaleSetID uint64
dockerClient *dockerclient.Client
scalesetClient *scaleset.Client
minRunners int
minRunners uint32
logger *slog.Logger
}
func (a *Scaler) HandleDesiredRunnerCount(ctx context.Context, count int) (int, error) {
currentCount := a.runners.count()
targetRunnerCount := min(a.minRunners + count)
targetRunnerCount := min(int(a.minRunners) + count)
switch {
case targetRunnerCount == currentCount:
+19 -20
View File
@@ -9,6 +9,7 @@ import (
"log/slog"
"net/http"
"os"
"sync/atomic"
"time"
"github.com/actions/scaleset"
@@ -20,9 +21,9 @@ const (
)
type Config struct {
ScaleSetID int
MinRunners int
MaxRunners int
ScaleSetID uint64
MinRunners uint32
MaxRunners uint32
Logger *slog.Logger
}
@@ -38,12 +39,6 @@ func (c *Config) Validate() error {
if c.ScaleSetID == 0 {
return errors.New("scaleSetID is required")
}
if c.MinRunners < 0 {
return errors.New("minRunners must be greater than or equal to 0")
}
if c.MaxRunners < 0 {
return errors.New("maxRunners must be greater than or equal to 0")
}
if c.MaxRunners > 0 && c.MinRunners > c.MaxRunners {
return errors.New("minRunners must be less than or equal to maxRunners")
}
@@ -55,12 +50,11 @@ type Listener struct {
client *scaleset.Client
// Configuration for the listener
scaleSetID int
minRunners int
maxRunners int
scaleSetID uint64
maxRunners atomic.Uint32
// lastMessageID keeps track of the last processed message ID
lastMessageID int64
lastMessageID uint64
// hostname of the current machine
hostname string
// session represents the current message session
@@ -70,6 +64,10 @@ type Listener struct {
logger *slog.Logger
}
func (l *Listener) SetMaxRunners(count uint32) {
l.maxRunners.Store(count)
}
func New(client *scaleset.Client, config Config) (*Listener, error) {
if client == nil {
return nil, errors.New("client is required")
@@ -85,20 +83,21 @@ func New(client *scaleset.Client, config Config) (*Listener, error) {
config.Logger.Info("Failed to get hostname, fallback to uuid", "uuid", hostname, "error", err)
}
return &Listener{
listener := &Listener{
client: client,
scaleSetID: config.ScaleSetID,
minRunners: config.MinRunners,
maxRunners: config.MaxRunners,
hostname: hostname,
logger: config.Logger,
}, nil
}
listener.maxRunners.Store(uint32(config.MaxRunners))
return listener, nil
}
type Scaler interface {
HandleJobStarted(ctx context.Context, jobInfo *scaleset.JobStarted) error
HandleJobCompleted(ctx context.Context, jobInfo *scaleset.JobCompleted) error
HandleDesiredRunnerCount(ctx context.Context, count int) (int, error)
HandleDesiredRunnerCount(ctx context.Context, count uint64) (int, error)
}
func (l *Listener) Run(ctx context.Context, handler Scaler) error {
@@ -235,7 +234,7 @@ func (l *Listener) getMessage(ctx context.Context) (*scaleset.RunnerScaleSetMess
l.session.MessageQueueURL,
l.session.MessageQueueAccessToken,
l.lastMessageID,
l.maxRunners,
l.maxRunners.Load(),
)
if err == nil { // if NO error
return msg, nil
@@ -257,7 +256,7 @@ func (l *Listener) getMessage(ctx context.Context) (*scaleset.RunnerScaleSetMess
l.session.MessageQueueURL,
l.session.MessageQueueAccessToken,
l.lastMessageID,
l.maxRunners,
l.maxRunners.Load(),
)
if err != nil { // if error
return nil, fmt.Errorf("failed to get next message after message session refresh: %w", err)
+13 -13
View File
@@ -71,9 +71,9 @@ type Label struct {
}
type RunnerGroup struct {
ID int `json:"id"`
ID uint64 `json:"id"`
Name string `json:"name"`
Size int64 `json:"size"`
Size uint64 `json:"size"`
IsDefault bool `json:"isDefaultGroup"`
}
@@ -83,9 +83,9 @@ type RunnerGroupList struct {
}
type RunnerScaleSet struct {
ID int `json:"id,omitempty"`
ID uint64 `json:"id,omitempty"`
Name string `json:"name,omitempty"`
RunnerGroupID int `json:"runnerGroupId,omitempty"`
RunnerGroupID uint64 `json:"runnerGroupId,omitempty"`
RunnerGroupName string `json:"runnerGroupName,omitempty"`
Labels []Label `json:"labels,omitempty"`
RunnerSetting RunnerSetting `json:"RunnerSetting,omitempty"`
@@ -100,14 +100,14 @@ type RunnerScaleSetJitRunnerSetting struct {
}
type RunnerScaleSetMessage struct {
MessageID int64 `json:"messageId"`
MessageID uint64 `json:"messageId"`
MessageType string `json:"messageType"`
Body string `json:"body"`
Statistics *RunnerScaleSetStatistic `json:"statistics"`
}
type runnerScaleSetsResponse struct {
Count int `json:"count"`
Count uint64 `json:"count"`
RunnerScaleSets []RunnerScaleSet `json:"value"`
}
@@ -121,13 +121,13 @@ type RunnerScaleSetSession struct {
}
type RunnerScaleSetStatistic struct {
TotalAvailableJobs int `json:"totalAvailableJobs"`
TotalAcquiredJobs int `json:"totalAcquiredJobs"`
TotalAssignedJobs int `json:"totalAssignedJobs"`
TotalRunningJobs int `json:"totalRunningJobs"`
TotalRegisteredRunners int `json:"totalRegisteredRunners"`
TotalBusyRunners int `json:"totalBusyRunners"`
TotalIdleRunners int `json:"totalIdleRunners"`
TotalAvailableJobs uint64 `json:"totalAvailableJobs"`
TotalAcquiredJobs uint64 `json:"totalAcquiredJobs"`
TotalAssignedJobs uint64 `json:"totalAssignedJobs"`
TotalRunningJobs uint64 `json:"totalRunningJobs"`
TotalRegisteredRunners uint64 `json:"totalRegisteredRunners"`
TotalBusyRunners uint64 `json:"totalBusyRunners"`
TotalIdleRunners uint64 `json:"totalIdleRunners"`
}
type RunnerSetting struct {