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

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 {

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

examples/dockerscaleset/config.go

@@ -12,8 +12,8 @@ import (
 
 type Config struct {
 	RegistrationURL string
-	MaxRunners      int
+	MaxRunners      uint32
-	MinRunners      int
+	MinRunners      uint32
 	ScaleSetName    string
 	RunnerGroup     string
 	GitHubApp       scaleset.GitHubAppAuth

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.Uint32Var(&cfg.MaxRunners, "max-runners", 10, "Maximum number of runners")
-	flags.IntVar(&cfg.MinRunners, "min-runners", 0, "Minimum 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,
+		MinRunners: int(c.MinRunners),
-		MaxRunners: c.MaxRunners,
+		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))

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:

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
+	ScaleSetID uint64
-	MinRunners int
+	MinRunners uint32
-	MaxRunners int
+	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
+	scaleSetID uint64
-	minRunners int
+	maxRunners atomic.Uint32
-	maxRunners int
 
 	// 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)

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"`
+	TotalAvailableJobs     uint64 `json:"totalAvailableJobs"`
-	TotalAcquiredJobs      int `json:"totalAcquiredJobs"`
+	TotalAcquiredJobs      uint64 `json:"totalAcquiredJobs"`
-	TotalAssignedJobs      int `json:"totalAssignedJobs"`
+	TotalAssignedJobs      uint64 `json:"totalAssignedJobs"`
-	TotalRunningJobs       int `json:"totalRunningJobs"`
+	TotalRunningJobs       uint64 `json:"totalRunningJobs"`
-	TotalRegisteredRunners int `json:"totalRegisteredRunners"`
+	TotalRegisteredRunners uint64 `json:"totalRegisteredRunners"`
-	TotalBusyRunners       int `json:"totalBusyRunners"`
+	TotalBusyRunners       uint64 `json:"totalBusyRunners"`
-	TotalIdleRunners       int `json:"totalIdleRunners"`
+	TotalIdleRunners       uint64 `json:"totalIdleRunners"`
 }
 
 type RunnerSetting struct {