Moving to scaleset client for the controller (#4390)

2026-03-13 14:36:41 +01:00
parent 1d9f626c53
commit f99c6eda0b
27 changed files with 2695 additions and 360 deletions
@@ -0,0 +1,230 @@
+package fake
+
+import (
+	"context"
+
+	"github.com/actions/actions-runner-controller/controllers/actions.github.com/multiclient"
+	"github.com/actions/scaleset"
+)
+
+// ClientOption is a functional option for configuring a fake Client
+type ClientOption func(*Client)
+
+// WithGetRunnerScaleSet configures the result of GetRunnerScaleSet
+func WithGetRunnerScaleSet(result *scaleset.RunnerScaleSet, err error) ClientOption {
+	return func(c *Client) {
+		c.getRunnerScaleSetResult.RunnerScaleSet = result
+		c.getRunnerScaleSetResult.err = err
+	}
+}
+
+// WithGetRunnerScaleSetByID configures the result of GetRunnerScaleSetByID
+func WithGetRunnerScaleSetByID(result *scaleset.RunnerScaleSet, err error) ClientOption {
+	return func(c *Client) {
+		c.getRunnerScaleSetByIDResult.RunnerScaleSet = result
+		c.getRunnerScaleSetByIDResult.err = err
+	}
+}
+
+// WithGetRunnerGroupByName configures the result of GetRunnerGroupByName
+func WithGetRunnerGroupByName(result *scaleset.RunnerGroup, err error) ClientOption {
+	return func(c *Client) {
+		c.getRunnerGroupByNameResult.RunnerGroup = result
+		c.getRunnerGroupByNameResult.err = err
+	}
+}
+
+// WithGetRunnerGroupByNameFunc configures a function to handle GetRunnerGroupByName calls dynamically
+func WithGetRunnerGroupByNameFunc(fn func(context.Context, string) (*scaleset.RunnerGroup, error)) ClientOption {
+	return func(c *Client) {
+		c.getRunnerGroupByNameFunc = fn
+	}
+}
+
+// WithCreateRunnerScaleSet configures the result of CreateRunnerScaleSet
+func WithCreateRunnerScaleSet(result *scaleset.RunnerScaleSet, err error) ClientOption {
+	return func(c *Client) {
+		c.createRunnerScaleSetResult.RunnerScaleSet = result
+		c.createRunnerScaleSetResult.err = err
+	}
+}
+
+// WithUpdateRunnerScaleSet configures the result of UpdateRunnerScaleSet
+func WithUpdateRunnerScaleSet(result *scaleset.RunnerScaleSet, err error) ClientOption {
+	return func(c *Client) {
+		c.updateRunnerScaleSetResult.RunnerScaleSet = result
+		c.updateRunnerScaleSetResult.err = err
+	}
+}
+
+// WithDeleteRunnerScaleSet configures the result of DeleteRunnerScaleSet
+func WithDeleteRunnerScaleSet(err error) ClientOption {
+	return func(c *Client) {
+		c.deleteRunnerScaleSetResult.err = err
+	}
+}
+
+// WithRemoveRunner configures the result of RemoveRunner
+func WithRemoveRunner(err error) ClientOption {
+	return func(c *Client) {
+		c.removeRunnerResult.err = err
+	}
+}
+
+// WithGenerateJitRunnerConfig configures the result of GenerateJitRunnerConfig
+func WithGenerateJitRunnerConfig(result *scaleset.RunnerScaleSetJitRunnerConfig, err error) ClientOption {
+	return func(c *Client) {
+		c.generateJitRunnerConfigResult.RunnerScaleSetJitRunnerConfig = result
+		c.generateJitRunnerConfigResult.err = err
+	}
+}
+
+// WithGetRunnerByName configures the result of GetRunnerByName
+func WithGetRunnerByName(result *scaleset.RunnerReference, err error) ClientOption {
+	return func(c *Client) {
+		c.getRunnerByNameResult.RunnerReference = result
+		c.getRunnerByNameResult.err = err
+	}
+}
+
+// WithGetRunner configures the result of GetRunner
+func WithGetRunner(result *scaleset.RunnerReference, err error) ClientOption {
+	return func(c *Client) {
+		c.getRunnerResult.RunnerReference = result
+		c.getRunnerResult.err = err
+	}
+}
+
+// WithSystemInfo configures the SystemInfo
+func WithSystemInfo(info scaleset.SystemInfo) ClientOption {
+	return func(c *Client) {
+		c.systemInfo = info
+	}
+}
+
+// WithUpdateRunnerScaleSetFunc configures a function to handle UpdateRunnerScaleSet calls dynamically
+func WithUpdateRunnerScaleSetFunc(fn func(context.Context, int, *scaleset.RunnerScaleSet) (*scaleset.RunnerScaleSet, error)) ClientOption {
+	return func(c *Client) {
+		c.updateRunnerScaleSetFunc = fn
+	}
+}
+
+// Client implements multiclient.Client interface for testing
+type Client struct {
+	systemInfo               scaleset.SystemInfo
+	updateRunnerScaleSetFunc func(context.Context, int, *scaleset.RunnerScaleSet) (*scaleset.RunnerScaleSet, error)
+
+	getRunnerScaleSetResult struct {
+		*scaleset.RunnerScaleSet
+		err error
+	}
+	getRunnerScaleSetByIDResult struct {
+		*scaleset.RunnerScaleSet
+		err error
+	}
+	getRunnerGroupByNameResult struct {
+		*scaleset.RunnerGroup
+		err error
+	}
+	getRunnerGroupByNameFunc   func(context.Context, string) (*scaleset.RunnerGroup, error)
+	createRunnerScaleSetResult struct {
+		*scaleset.RunnerScaleSet
+		err error
+	}
+	updateRunnerScaleSetResult struct {
+		*scaleset.RunnerScaleSet
+		err error
+	}
+	deleteRunnerScaleSetResult struct {
+		err error
+	}
+	removeRunnerResult struct {
+		err error
+	}
+	generateJitRunnerConfigResult struct {
+		*scaleset.RunnerScaleSetJitRunnerConfig
+		err error
+	}
+	getRunnerByNameResult struct {
+		*scaleset.RunnerReference
+		err error
+	}
+	getRunnerResult struct {
+		*scaleset.RunnerReference
+		err error
+	}
+	messageSessionClientResult struct {
+		*scaleset.MessageSessionClient
+		err error
+	}
+}
+
+// Compile-time interface check
+var _ multiclient.Client = (*Client)(nil)
+
+// NewClient creates a new fake Client with the given options
+func NewClient(opts ...ClientOption) *Client {
+	c := &Client{}
+	for _, opt := range opts {
+		opt(c)
+	}
+	return c
+}
+
+func (c *Client) SetSystemInfo(info scaleset.SystemInfo) {
+	c.systemInfo = info
+}
+
+func (c *Client) SystemInfo() scaleset.SystemInfo {
+	return c.systemInfo
+}
+
+func (c *Client) MessageSessionClient(ctx context.Context, runnerScaleSetID int, owner string, options ...scaleset.HTTPOption) (*scaleset.MessageSessionClient, error) {
+	return c.messageSessionClientResult.MessageSessionClient, c.messageSessionClientResult.err
+}
+
+func (c *Client) GenerateJitRunnerConfig(ctx context.Context, jitRunnerSetting *scaleset.RunnerScaleSetJitRunnerSetting, scaleSetID int) (*scaleset.RunnerScaleSetJitRunnerConfig, error) {
+	return c.generateJitRunnerConfigResult.RunnerScaleSetJitRunnerConfig, c.generateJitRunnerConfigResult.err
+}
+
+func (c *Client) GetRunner(ctx context.Context, runnerID int) (*scaleset.RunnerReference, error) {
+	return c.getRunnerResult.RunnerReference, c.getRunnerResult.err
+}
+
+func (c *Client) GetRunnerByName(ctx context.Context, runnerName string) (*scaleset.RunnerReference, error) {
+	return c.getRunnerByNameResult.RunnerReference, c.getRunnerByNameResult.err
+}
+
+func (c *Client) RemoveRunner(ctx context.Context, runnerID int64) error {
+	return c.removeRunnerResult.err
+}
+
+func (c *Client) GetRunnerGroupByName(ctx context.Context, runnerGroup string) (*scaleset.RunnerGroup, error) {
+	if c.getRunnerGroupByNameFunc != nil {
+		return c.getRunnerGroupByNameFunc(ctx, runnerGroup)
+	}
+	return c.getRunnerGroupByNameResult.RunnerGroup, c.getRunnerGroupByNameResult.err
+}
+
+func (c *Client) GetRunnerScaleSet(ctx context.Context, runnerGroupID int, runnerScaleSetName string) (*scaleset.RunnerScaleSet, error) {
+	return c.getRunnerScaleSetResult.RunnerScaleSet, c.getRunnerScaleSetResult.err
+}
+
+func (c *Client) GetRunnerScaleSetByID(ctx context.Context, runnerScaleSetID int) (*scaleset.RunnerScaleSet, error) {
+	return c.getRunnerScaleSetByIDResult.RunnerScaleSet, c.getRunnerScaleSetByIDResult.err
+}
+
+func (c *Client) CreateRunnerScaleSet(ctx context.Context, runnerScaleSet *scaleset.RunnerScaleSet) (*scaleset.RunnerScaleSet, error) {
+	return c.createRunnerScaleSetResult.RunnerScaleSet, c.createRunnerScaleSetResult.err
+}
+
+func (c *Client) UpdateRunnerScaleSet(ctx context.Context, runnerScaleSetID int, runnerScaleSet *scaleset.RunnerScaleSet) (*scaleset.RunnerScaleSet, error) {
+	if c.updateRunnerScaleSetFunc != nil {
+		return c.updateRunnerScaleSetFunc(ctx, runnerScaleSetID, runnerScaleSet)
+	}
+	return c.updateRunnerScaleSetResult.RunnerScaleSet, c.updateRunnerScaleSetResult.err
+}
+
+func (c *Client) DeleteRunnerScaleSet(ctx context.Context, runnerScaleSetID int) error {
+	return c.deleteRunnerScaleSetResult.err
+}
@@ -0,0 +1,53 @@
+package fake
+
+import (
+	"context"
+
+	"github.com/actions/actions-runner-controller/controllers/actions.github.com/multiclient"
+)
+
+// MultiClientOption is a functional option for configuring a fake MultiClient
+type MultiClientOption func(*MultiClient)
+
+// WithClient configures the client that GetClientFor will return
+func WithClient(c multiclient.Client) MultiClientOption {
+	return func(mc *MultiClient) {
+		mc.client = c
+	}
+}
+
+// WithGetClientForError configures an error that GetClientFor will return
+func WithGetClientForError(err error) MultiClientOption {
+	return func(mc *MultiClient) {
+		mc.getClientForErr = err
+	}
+}
+
+// MultiClient implements multiclient.MultiClient interface for testing
+type MultiClient struct {
+	client          multiclient.Client
+	getClientForErr error
+}
+
+// Compile-time interface check
+var _ multiclient.MultiClient = (*MultiClient)(nil)
+
+// NewMultiClient creates a new fake MultiClient with the given options
+func NewMultiClient(opts ...MultiClientOption) *MultiClient {
+	mc := &MultiClient{}
+	for _, opt := range opts {
+		opt(mc)
+	}
+	// Default behavior: if no client configured, return a default NewClient()
+	if mc.client == nil {
+		mc.client = NewClient()
+	}
+	return mc
+}
+
+func (mc *MultiClient) GetClientFor(ctx context.Context, opts *multiclient.ClientForOptions) (multiclient.Client, error) {
+	if mc.getClientForErr != nil {
+		return nil, mc.getClientForErr
+	}
+	return mc.client, nil
+}
@@ -0,0 +1,177 @@
+package multiclient
+
+import (
+	"context"
+	"crypto/sha256"
+	"crypto/x509"
+	"fmt"
+	"net/http"
+	"net/url"
+	"strconv"
+	"sync"
+
+	"github.com/actions/actions-runner-controller/apis/actions.github.com/v1alpha1/appconfig"
+	"github.com/actions/actions-runner-controller/build"
+	"github.com/actions/actions-runner-controller/github/actions"
+	"github.com/actions/scaleset"
+	"github.com/google/uuid"
+)
+
+type MultiClient interface {
+	GetClientFor(ctx context.Context, opts *ClientForOptions) (Client, error)
+}
+
+type Scaleset struct {
+	mu      sync.Mutex
+	clients map[string]*multiClientEntry
+}
+
+type multiClientEntry struct {
+	client  *scaleset.Client
+	rootCAs *x509.CertPool
+}
+
+func NewScaleset() *Scaleset {
+	return &Scaleset{
+		clients: make(map[string]*multiClientEntry),
+	}
+}
+
+type Client interface {
+	SetSystemInfo(info scaleset.SystemInfo)
+	SystemInfo() scaleset.SystemInfo
+
+	MessageSessionClient(ctx context.Context, runnerScaleSetID int, owner string, options ...scaleset.HTTPOption) (*scaleset.MessageSessionClient, error)
+
+	GenerateJitRunnerConfig(ctx context.Context, jitRunnerSetting *scaleset.RunnerScaleSetJitRunnerSetting, scaleSetID int) (*scaleset.RunnerScaleSetJitRunnerConfig, error)
+
+	GetRunner(ctx context.Context, runnerID int) (*scaleset.RunnerReference, error)
+	GetRunnerByName(ctx context.Context, runnerName string) (*scaleset.RunnerReference, error)
+	RemoveRunner(ctx context.Context, runnerID int64) error
+
+	GetRunnerGroupByName(ctx context.Context, runnerGroup string) (*scaleset.RunnerGroup, error)
+
+	GetRunnerScaleSet(ctx context.Context, runnerGroupID int, runnerScaleSetName string) (*scaleset.RunnerScaleSet, error)
+	GetRunnerScaleSetByID(ctx context.Context, runnerScaleSetID int) (*scaleset.RunnerScaleSet, error)
+
+	CreateRunnerScaleSet(ctx context.Context, runnerScaleSet *scaleset.RunnerScaleSet) (*scaleset.RunnerScaleSet, error)
+	UpdateRunnerScaleSet(ctx context.Context, runnerScaleSetID int, runnerScaleSet *scaleset.RunnerScaleSet) (*scaleset.RunnerScaleSet, error)
+	DeleteRunnerScaleSet(ctx context.Context, runnerScaleSetID int) error
+}
+
+func (m *Scaleset) GetClientFor(ctx context.Context, opts *ClientForOptions) (Client, error) {
+	identifier, err := opts.identifier()
+	if err != nil {
+		return nil, fmt.Errorf("failed to generate client identifier: %w", err)
+	}
+
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
+	entry, ok := m.clients[identifier]
+	if ok && entry.rootCAs.Equal(opts.RootCAs) {
+		return entry.client, nil
+	}
+
+	client, err := opts.newClient()
+	if err != nil {
+		return nil, fmt.Errorf("failed to create new client: %w", err)
+	}
+
+	m.clients[identifier] = &multiClientEntry{
+		client:  client,
+		rootCAs: opts.RootCAs,
+	}
+
+	return client, nil
+}
+
+type ClientForOptions struct {
+	GithubConfigURL string
+	AppConfig       appconfig.AppConfig
+	Namespace       string
+	RootCAs         *x509.CertPool
+	ProxyFunc       func(*http.Request) (*url.URL, error)
+}
+
+func (o *ClientForOptions) identifier() (string, error) {
+	if err := o.AppConfig.Validate(); err != nil {
+		return "", fmt.Errorf("failed to validate app config: %w", err)
+	}
+	if _, err := actions.ParseGitHubConfigFromURL(o.GithubConfigURL); err != nil {
+		return "", fmt.Errorf("failed to parse GitHub config URL: %w", err)
+	}
+	if o.Namespace == "" {
+		return "", fmt.Errorf("namespace is required to generate client identifier")
+	}
+	identifier := fmt.Sprintf("configURL:%q,namespace:%q,proxy:%t", o.GithubConfigURL, o.Namespace, o.ProxyFunc != nil)
+
+	if o.AppConfig.Token != "" {
+		identifier += fmt.Sprintf(",token:%q,", o.AppConfig.Token)
+	} else {
+		identifier += fmt.Sprintf(
+			",appID:%q,installationID:%q,key:%q",
+			o.AppConfig.AppID,
+			strconv.FormatInt(o.AppConfig.AppInstallationID, 10),
+			o.AppConfig.AppPrivateKey,
+		)
+	}
+
+	if o.RootCAs != nil {
+		// ignoring because this cert pool is intended not to come from SystemCertPool
+		// nolint:staticcheck
+		identifier += fmt.Sprintf(",rootCAs:%q", o.RootCAs.Subjects())
+	}
+
+	return uuid.NewHash(sha256.New(), uuid.NameSpaceOID, []byte(identifier), 6).String(), nil
+}
+
+func (o *ClientForOptions) newClient() (*scaleset.Client, error) {
+	systemInfo := scaleset.SystemInfo{
+		System:     "actions-runner-controller",
+		Version:    build.Version,
+		CommitSHA:  build.CommitSHA,
+		ScaleSetID: 0, // by default, scale set is 0 (not created yet)
+		Subsystem:  "gha-scale-set-controller",
+	}
+
+	var options []scaleset.HTTPOption
+	if o.RootCAs != nil {
+		options = append(options, scaleset.WithRootCAs(o.RootCAs))
+	}
+	if o.ProxyFunc != nil {
+		options = append(options, scaleset.WithProxy(o.ProxyFunc))
+	}
+
+	if o.AppConfig.Token != "" {
+		c, err := scaleset.NewClientWithPersonalAccessToken(
+			scaleset.NewClientWithPersonalAccessTokenConfig{
+				GitHubConfigURL:     o.GithubConfigURL,
+				PersonalAccessToken: o.AppConfig.Token,
+				SystemInfo:          systemInfo,
+			},
+			options...,
+		)
+		if err != nil {
+			return nil, fmt.Errorf("failed to instantiate client with personal access token auth: %w", err)
+		}
+		return c, nil
+	}
+
+	c, err := scaleset.NewClientWithGitHubApp(
+		scaleset.ClientWithGitHubAppConfig{
+			GitHubConfigURL: o.GithubConfigURL,
+			GitHubAppAuth: scaleset.GitHubAppAuth{
+				ClientID:       o.AppConfig.AppID,
+				InstallationID: o.AppConfig.AppInstallationID,
+				PrivateKey:     o.AppConfig.AppPrivateKey,
+			},
+			SystemInfo: systemInfo,
+		},
+		options...,
+	)
+	if err != nil {
+		return nil, fmt.Errorf("failed to instantiate client with GitHub App auth: %w", err)
+	}
+	return c, nil
+}