Propagate max capacity information to the actions back-end (#3431)

This commit is contained in:
Nikola Jokic
2024-04-16 14:00:40 +02:00
committed by GitHub
parent 8075e5ee74
commit 4ee49fee14
18 changed files with 147 additions and 100 deletions

View File

@@ -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 {

View File

@@ -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) {

View File

@@ -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
}

View File

@@ -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
}

View File

@@ -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)
}

View File

@@ -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)
}

View File

@@ -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