mirror of
https://github.com/actions/actions-runner-controller.git
synced 2025-12-12 04:26:51 +00:00
Propagate max capacity information to the actions back-end (#3431)
This commit is contained in:
@@ -29,6 +29,9 @@ const (
|
||||
apiVersionQueryParam = "api-version=6.0-preview"
|
||||
)
|
||||
|
||||
// Header used to propagate capacity information to the back-end
|
||||
const HeaderScaleSetMaxCapacity = "X-ScaleSetMaxCapacity"
|
||||
|
||||
//go:generate mockery --inpackage --name=ActionsService
|
||||
type ActionsService interface {
|
||||
GetRunnerScaleSet(ctx context.Context, runnerGroupId int, runnerScaleSetName string) (*RunnerScaleSet, error)
|
||||
@@ -45,7 +48,7 @@ type ActionsService interface {
|
||||
AcquireJobs(ctx context.Context, runnerScaleSetId int, messageQueueAccessToken string, requestIds []int64) ([]int64, error)
|
||||
GetAcquirableJobs(ctx context.Context, runnerScaleSetId int) (*AcquirableJobList, error)
|
||||
|
||||
GetMessage(ctx context.Context, messageQueueUrl, messageQueueAccessToken string, lastMessageId int64) (*RunnerScaleSetMessage, error)
|
||||
GetMessage(ctx context.Context, messageQueueUrl, messageQueueAccessToken string, lastMessageId int64, maxCapacity int) (*RunnerScaleSetMessage, error)
|
||||
DeleteMessage(ctx context.Context, messageQueueUrl, messageQueueAccessToken string, messageId int64) error
|
||||
|
||||
GenerateJitRunnerConfig(ctx context.Context, jitRunnerSetting *RunnerScaleSetJitRunnerSetting, scaleSetId int) (*RunnerScaleSetJitRunnerConfig, error)
|
||||
@@ -104,6 +107,8 @@ type Client struct {
|
||||
proxyFunc ProxyFunc
|
||||
}
|
||||
|
||||
var _ ActionsService = &Client{}
|
||||
|
||||
type ProxyFunc func(req *http.Request) (*url.URL, error)
|
||||
|
||||
type ClientOption func(*Client)
|
||||
@@ -543,7 +548,7 @@ func (c *Client) DeleteRunnerScaleSet(ctx context.Context, runnerScaleSetId int)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *Client) GetMessage(ctx context.Context, messageQueueUrl, messageQueueAccessToken string, lastMessageId int64) (*RunnerScaleSetMessage, error) {
|
||||
func (c *Client) GetMessage(ctx context.Context, messageQueueUrl, messageQueueAccessToken string, lastMessageId int64, maxCapacity int) (*RunnerScaleSetMessage, error) {
|
||||
u, err := url.Parse(messageQueueUrl)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
@@ -555,6 +560,10 @@ func (c *Client) GetMessage(ctx context.Context, messageQueueUrl, messageQueueAc
|
||||
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, err
|
||||
@@ -563,6 +572,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.String())
|
||||
req.Header.Set(HeaderScaleSetMaxCapacity, strconv.Itoa(maxCapacity))
|
||||
|
||||
resp, err := c.Do(req)
|
||||
if err != nil {
|
||||
|
||||
@@ -5,6 +5,7 @@ import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"net/http"
|
||||
"strconv"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
@@ -35,7 +36,7 @@ func TestGetMessage(t *testing.T) {
|
||||
client, err := actions.NewClient(s.configURLForOrg("my-org"), auth)
|
||||
require.NoError(t, err)
|
||||
|
||||
got, err := client.GetMessage(ctx, s.URL, token, 0)
|
||||
got, err := client.GetMessage(ctx, s.URL, token, 0, 10)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, want, got)
|
||||
})
|
||||
@@ -52,7 +53,7 @@ func TestGetMessage(t *testing.T) {
|
||||
client, err := actions.NewClient(s.configURLForOrg("my-org"), auth)
|
||||
require.NoError(t, err)
|
||||
|
||||
got, err := client.GetMessage(ctx, s.URL, token, 1)
|
||||
got, err := client.GetMessage(ctx, s.URL, token, 1, 10)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, want, got)
|
||||
})
|
||||
@@ -76,7 +77,7 @@ func TestGetMessage(t *testing.T) {
|
||||
)
|
||||
require.NoError(t, err)
|
||||
|
||||
_, err = client.GetMessage(ctx, server.URL, token, 0)
|
||||
_, err = client.GetMessage(ctx, server.URL, token, 0, 10)
|
||||
assert.NotNil(t, err)
|
||||
assert.Equalf(t, actualRetry, expectedRetry, "A retry was expected after the first request but got: %v", actualRetry)
|
||||
})
|
||||
@@ -89,7 +90,7 @@ func TestGetMessage(t *testing.T) {
|
||||
client, err := actions.NewClient(server.configURLForOrg("my-org"), auth)
|
||||
require.NoError(t, err)
|
||||
|
||||
_, err = client.GetMessage(ctx, server.URL, token, 0)
|
||||
_, err = client.GetMessage(ctx, server.URL, token, 0, 10)
|
||||
require.NotNil(t, err)
|
||||
|
||||
var expectedErr *actions.MessageQueueTokenExpiredError
|
||||
@@ -108,7 +109,7 @@ func TestGetMessage(t *testing.T) {
|
||||
client, err := actions.NewClient(server.configURLForOrg("my-org"), auth)
|
||||
require.NoError(t, err)
|
||||
|
||||
_, err = client.GetMessage(ctx, server.URL, token, 0)
|
||||
_, err = client.GetMessage(ctx, server.URL, token, 0, 10)
|
||||
require.NotNil(t, err)
|
||||
assert.Equal(t, want.Error(), err.Error())
|
||||
})
|
||||
@@ -122,9 +123,35 @@ func TestGetMessage(t *testing.T) {
|
||||
client, err := actions.NewClient(server.configURLForOrg("my-org"), auth)
|
||||
require.NoError(t, err)
|
||||
|
||||
_, err = client.GetMessage(ctx, server.URL, token, 0)
|
||||
_, err = client.GetMessage(ctx, server.URL, token, 0, 10)
|
||||
assert.NotNil(t, err)
|
||||
})
|
||||
|
||||
t.Run("Capacity error handling", func(t *testing.T) {
|
||||
server := newActionsServer(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
hc := r.Header.Get(actions.HeaderScaleSetMaxCapacity)
|
||||
c, err := strconv.Atoi(hc)
|
||||
require.NoError(t, err)
|
||||
assert.GreaterOrEqual(t, c, 0)
|
||||
|
||||
w.WriteHeader(http.StatusBadRequest)
|
||||
w.Header().Set("Content-Type", "text/plain")
|
||||
}))
|
||||
|
||||
client, err := actions.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, &actions.ActionsError{}))
|
||||
|
||||
_, err = client.GetMessage(ctx, server.URL, token, 0, 0)
|
||||
assert.Error(t, err)
|
||||
var expectedErr *actions.ActionsError
|
||||
assert.ErrorAs(t, err, &expectedErr)
|
||||
assert.Equal(t, http.StatusBadRequest, expectedErr.StatusCode)
|
||||
})
|
||||
}
|
||||
|
||||
func TestDeleteMessage(t *testing.T) {
|
||||
|
||||
@@ -47,7 +47,6 @@ func (e *ActionsError) IsException(target string) bool {
|
||||
if ex, ok := e.Err.(*ActionsExceptionError); ok {
|
||||
return strings.Contains(ex.ExceptionName, target)
|
||||
}
|
||||
|
||||
return false
|
||||
}
|
||||
|
||||
|
||||
@@ -259,7 +259,7 @@ func (f *FakeClient) GetAcquirableJobs(ctx context.Context, runnerScaleSetId int
|
||||
return f.getAcquirableJobsResult.AcquirableJobList, f.getAcquirableJobsResult.err
|
||||
}
|
||||
|
||||
func (f *FakeClient) GetMessage(ctx context.Context, messageQueueUrl, messageQueueAccessToken string, lastMessageId int64) (*actions.RunnerScaleSetMessage, error) {
|
||||
func (f *FakeClient) GetMessage(ctx context.Context, messageQueueUrl, messageQueueAccessToken string, lastMessageId int64, maxCapacity int) (*actions.RunnerScaleSetMessage, error) {
|
||||
return f.getMessageResult.RunnerScaleSetMessage, f.getMessageResult.err
|
||||
}
|
||||
|
||||
|
||||
@@ -186,25 +186,25 @@ func (_m *MockActionsService) GetAcquirableJobs(ctx context.Context, runnerScale
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
// GetMessage provides a mock function with given fields: ctx, messageQueueUrl, messageQueueAccessToken, lastMessageId
|
||||
func (_m *MockActionsService) GetMessage(ctx context.Context, messageQueueUrl string, messageQueueAccessToken string, lastMessageId int64) (*RunnerScaleSetMessage, error) {
|
||||
ret := _m.Called(ctx, messageQueueUrl, messageQueueAccessToken, lastMessageId)
|
||||
// GetMessage provides a mock function with given fields: ctx, messageQueueUrl, messageQueueAccessToken, lastMessageId, maxCapacity
|
||||
func (_m *MockActionsService) GetMessage(ctx context.Context, messageQueueUrl string, messageQueueAccessToken string, lastMessageId int64, maxCapacity int) (*RunnerScaleSetMessage, error) {
|
||||
ret := _m.Called(ctx, messageQueueUrl, messageQueueAccessToken, lastMessageId, maxCapacity)
|
||||
|
||||
var r0 *RunnerScaleSetMessage
|
||||
var r1 error
|
||||
if rf, ok := ret.Get(0).(func(context.Context, string, string, int64) (*RunnerScaleSetMessage, error)); ok {
|
||||
return rf(ctx, messageQueueUrl, messageQueueAccessToken, lastMessageId)
|
||||
if rf, ok := ret.Get(0).(func(context.Context, string, string, int64, int) (*RunnerScaleSetMessage, error)); ok {
|
||||
return rf(ctx, messageQueueUrl, messageQueueAccessToken, lastMessageId, maxCapacity)
|
||||
}
|
||||
if rf, ok := ret.Get(0).(func(context.Context, string, string, int64) *RunnerScaleSetMessage); ok {
|
||||
r0 = rf(ctx, messageQueueUrl, messageQueueAccessToken, lastMessageId)
|
||||
if rf, ok := ret.Get(0).(func(context.Context, string, string, int64, int) *RunnerScaleSetMessage); ok {
|
||||
r0 = rf(ctx, messageQueueUrl, messageQueueAccessToken, lastMessageId, maxCapacity)
|
||||
} else {
|
||||
if ret.Get(0) != nil {
|
||||
r0 = ret.Get(0).(*RunnerScaleSetMessage)
|
||||
}
|
||||
}
|
||||
|
||||
if rf, ok := ret.Get(1).(func(context.Context, string, string, int64) error); ok {
|
||||
r1 = rf(ctx, messageQueueUrl, messageQueueAccessToken, lastMessageId)
|
||||
if rf, ok := ret.Get(1).(func(context.Context, string, string, int64, int) error); ok {
|
||||
r1 = rf(ctx, messageQueueUrl, messageQueueAccessToken, lastMessageId, maxCapacity)
|
||||
} else {
|
||||
r1 = ret.Error(1)
|
||||
}
|
||||
|
||||
@@ -67,25 +67,25 @@ func (_m *MockSessionService) DeleteMessage(ctx context.Context, messageId int64
|
||||
return r0
|
||||
}
|
||||
|
||||
// GetMessage provides a mock function with given fields: ctx, lastMessageId
|
||||
func (_m *MockSessionService) GetMessage(ctx context.Context, lastMessageId int64) (*RunnerScaleSetMessage, error) {
|
||||
ret := _m.Called(ctx, lastMessageId)
|
||||
// GetMessage provides a mock function with given fields: ctx, lastMessageId, maxCapacity
|
||||
func (_m *MockSessionService) GetMessage(ctx context.Context, lastMessageId int64, maxCapacity int) (*RunnerScaleSetMessage, error) {
|
||||
ret := _m.Called(ctx, lastMessageId, maxCapacity)
|
||||
|
||||
var r0 *RunnerScaleSetMessage
|
||||
var r1 error
|
||||
if rf, ok := ret.Get(0).(func(context.Context, int64) (*RunnerScaleSetMessage, error)); ok {
|
||||
return rf(ctx, lastMessageId)
|
||||
if rf, ok := ret.Get(0).(func(context.Context, int64, int) (*RunnerScaleSetMessage, error)); ok {
|
||||
return rf(ctx, lastMessageId, maxCapacity)
|
||||
}
|
||||
if rf, ok := ret.Get(0).(func(context.Context, int64) *RunnerScaleSetMessage); ok {
|
||||
r0 = rf(ctx, lastMessageId)
|
||||
if rf, ok := ret.Get(0).(func(context.Context, int64, int) *RunnerScaleSetMessage); ok {
|
||||
r0 = rf(ctx, lastMessageId, maxCapacity)
|
||||
} else {
|
||||
if ret.Get(0) != nil {
|
||||
r0 = ret.Get(0).(*RunnerScaleSetMessage)
|
||||
}
|
||||
}
|
||||
|
||||
if rf, ok := ret.Get(1).(func(context.Context, int64) error); ok {
|
||||
r1 = rf(ctx, lastMessageId)
|
||||
if rf, ok := ret.Get(1).(func(context.Context, int64, int) error); ok {
|
||||
r1 = rf(ctx, lastMessageId, maxCapacity)
|
||||
} else {
|
||||
r1 = ret.Error(1)
|
||||
}
|
||||
|
||||
@@ -7,7 +7,7 @@ import (
|
||||
|
||||
//go:generate mockery --inpackage --name=SessionService
|
||||
type SessionService interface {
|
||||
GetMessage(ctx context.Context, lastMessageId int64) (*RunnerScaleSetMessage, error)
|
||||
GetMessage(ctx context.Context, lastMessageId int64, maxCapacity int) (*RunnerScaleSetMessage, error)
|
||||
DeleteMessage(ctx context.Context, messageId int64) error
|
||||
AcquireJobs(ctx context.Context, requestIds []int64) ([]int64, error)
|
||||
io.Closer
|
||||
|
||||
Reference in New Issue
Block a user