diff --git a/client.go b/client.go index 309b367..f4d6607 100644 --- a/client.go +++ b/client.go @@ -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 { diff --git a/client_test.go b/client_test.go index ce1a04a..2e63288 100644 --- a/client_test.go +++ b/client_test.go @@ -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 diff --git a/examples/dockerscaleset/config.go b/examples/dockerscaleset/config.go index 8412bc6..2796f14 100644 --- a/examples/dockerscaleset/config.go +++ b/examples/dockerscaleset/config.go @@ -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 diff --git a/examples/dockerscaleset/main.go b/examples/dockerscaleset/main.go index 35495f1..0b5308f 100644 --- a/examples/dockerscaleset/main.go +++ b/examples/dockerscaleset/main.go @@ -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)) diff --git a/examples/dockerscaleset/scaler.go b/examples/dockerscaleset/scaler.go index ad4449c..a92e0cb 100644 --- a/examples/dockerscaleset/scaler.go +++ b/examples/dockerscaleset/scaler.go @@ -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: diff --git a/listener/listener.go b/listener/listener.go index d9d8a1b..3555612 100644 --- a/listener/listener.go +++ b/listener/listener.go @@ -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) diff --git a/types.go b/types.go index 462564e..fec367c 100644 --- a/types.go +++ b/types.go @@ -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 {