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:
@@ -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
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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))
|
||||
|
||||
@@ -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
@@ -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)
|
||||
|
||||
@@ -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 {
|
||||
|
||||
Reference in New Issue
Block a user