Merge pull request #233 from gabriel-samfira/add-balancing-strategy

Add pool balancing strategy
cloudbase · Mar 15, 2024 · 6bfcddc · 6bfcddc
2 parents 569d512 + 206fe42
commit 6bfcddc
Show file tree

Hide file tree

Showing 35 changed files with 664 additions and 184 deletions.
diff --git a/cmd/garm-cli/cmd/enterprise.go b/cmd/garm-cli/cmd/enterprise.go
@@ -57,9 +57,10 @@ var enterpriseAddCmd = &cobra.Command{
 
 		newEnterpriseReq := apiClientEnterprises.NewCreateEnterpriseParams()
 		newEnterpriseReq.Body = params.CreateEnterpriseParams{
-			Name:            enterpriseName,
-			WebhookSecret:   enterpriseWebhookSecret,
-			CredentialsName: enterpriseCreds,
+			Name:             enterpriseName,
+			WebhookSecret:    enterpriseWebhookSecret,
+			CredentialsName:  enterpriseCreds,
+			PoolBalancerType: params.PoolBalancerType(poolBalancerType),
 		}
 		response, err := apiCli.Enterprises.CreateEnterprise(newEnterpriseReq, authToken)
 		if err != nil {
@@ -161,8 +162,9 @@ var enterpriseUpdateCmd = &cobra.Command{
 		}
 		updateEnterpriseReq := apiClientEnterprises.NewUpdateEnterpriseParams()
 		updateEnterpriseReq.Body = params.UpdateEntityParams{
-			WebhookSecret:   repoWebhookSecret,
-			CredentialsName: repoCreds,
+			WebhookSecret:    repoWebhookSecret,
+			CredentialsName:  repoCreds,
+			PoolBalancerType: params.PoolBalancerType(poolBalancerType),
 		}
 		updateEnterpriseReq.EnterpriseID = args[0]
 		response, err := apiCli.Enterprises.UpdateEnterprise(updateEnterpriseReq, authToken)
@@ -178,11 +180,13 @@ func init() {
 	enterpriseAddCmd.Flags().StringVar(&enterpriseName, "name", "", "The name of the enterprise")
 	enterpriseAddCmd.Flags().StringVar(&enterpriseWebhookSecret, "webhook-secret", "", "The webhook secret for this enterprise")
 	enterpriseAddCmd.Flags().StringVar(&enterpriseCreds, "credentials", "", "Credentials name. See credentials list.")
+	enterpriseAddCmd.Flags().StringVar(&poolBalancerType, "pool-balancer-type", string(params.PoolBalancerTypeRoundRobin), "The balancing strategy to use when creating runners in pools matching requested labels.")
 
 	enterpriseAddCmd.MarkFlagRequired("credentials") //nolint
 	enterpriseAddCmd.MarkFlagRequired("name")        //nolint
 	enterpriseUpdateCmd.Flags().StringVar(&enterpriseWebhookSecret, "webhook-secret", "", "The webhook secret for this enterprise")
 	enterpriseUpdateCmd.Flags().StringVar(&enterpriseCreds, "credentials", "", "Credentials name. See credentials list.")
+	enterpriseUpdateCmd.Flags().StringVar(&poolBalancerType, "pool-balancer-type", "", "The balancing strategy to use when creating runners in pools matching requested labels.")
 
 	enterpriseCmd.AddCommand(
 		enterpriseListCmd,
@@ -197,10 +201,10 @@ func init() {
 
 func formatEnterprises(enterprises []params.Enterprise) {
 	t := table.NewWriter()
-	header := table.Row{"ID", "Name", "Credentials name", "Pool mgr running"}
+	header := table.Row{"ID", "Name", "Credentials name", "Pool Balancer Type", "Pool mgr running"}
 	t.AppendHeader(header)
 	for _, val := range enterprises {
-		t.AppendRow(table.Row{val.ID, val.Name, val.CredentialsName, val.PoolManagerStatus.IsRunning})
+		t.AppendRow(table.Row{val.ID, val.Name, val.CredentialsName, val.GetBalancerType(), val.PoolManagerStatus.IsRunning})
 		t.AppendSeparator()
 	}
 	fmt.Println(t.Render())
@@ -213,6 +217,7 @@ func formatOneEnterprise(enterprise params.Enterprise) {
 	t.AppendHeader(header)
 	t.AppendRow(table.Row{"ID", enterprise.ID})
 	t.AppendRow(table.Row{"Name", enterprise.Name})
+	t.AppendRow(table.Row{"Pool balancer type", enterprise.GetBalancerType()})
 	t.AppendRow(table.Row{"Credentials", enterprise.CredentialsName})
 	t.AppendRow(table.Row{"Pool manager running", enterprise.PoolManagerStatus.IsRunning})
 	if !enterprise.PoolManagerStatus.IsRunning {

diff --git a/cmd/garm-cli/cmd/organization.go b/cmd/garm-cli/cmd/organization.go
@@ -163,9 +163,10 @@ var orgAddCmd = &cobra.Command{
 
 		newOrgReq := apiClientOrgs.NewCreateOrgParams()
 		newOrgReq.Body = params.CreateOrgParams{
-			Name:            orgName,
-			WebhookSecret:   orgWebhookSecret,
-			CredentialsName: orgCreds,
+			Name:             orgName,
+			WebhookSecret:    orgWebhookSecret,
+			CredentialsName:  orgCreds,
+			PoolBalancerType: params.PoolBalancerType(poolBalancerType),
 		}
 		response, err := apiCli.Organizations.CreateOrg(newOrgReq, authToken)
 		if err != nil {
@@ -213,8 +214,9 @@ var orgUpdateCmd = &cobra.Command{
 		}
 		updateOrgReq := apiClientOrgs.NewUpdateOrgParams()
 		updateOrgReq.Body = params.UpdateEntityParams{
-			WebhookSecret:   orgWebhookSecret,
-			CredentialsName: orgCreds,
+			WebhookSecret:    orgWebhookSecret,
+			CredentialsName:  orgCreds,
+			PoolBalancerType: params.PoolBalancerType(poolBalancerType),
 		}
 		updateOrgReq.OrgID = args[0]
 		response, err := apiCli.Organizations.UpdateOrg(updateOrgReq, authToken)
@@ -301,6 +303,7 @@ var orgDeleteCmd = &cobra.Command{
 
 func init() {
 	orgAddCmd.Flags().StringVar(&orgName, "name", "", "The name of the organization")
+	orgAddCmd.Flags().StringVar(&poolBalancerType, "pool-balancer-type", string(params.PoolBalancerTypeRoundRobin), "The balancing strategy to use when creating runners in pools matching requested labels.")
 	orgAddCmd.Flags().StringVar(&orgWebhookSecret, "webhook-secret", "", "The webhook secret for this organization")
 	orgAddCmd.Flags().StringVar(&orgCreds, "credentials", "", "Credentials name. See credentials list.")
 	orgAddCmd.Flags().BoolVar(&orgRandomWebhookSecret, "random-webhook-secret", false, "Generate a random webhook secret for this organization.")
@@ -315,6 +318,7 @@ func init() {
 
 	orgUpdateCmd.Flags().StringVar(&orgWebhookSecret, "webhook-secret", "", "The webhook secret for this organization")
 	orgUpdateCmd.Flags().StringVar(&orgCreds, "credentials", "", "Credentials name. See credentials list.")
+	orgUpdateCmd.Flags().StringVar(&poolBalancerType, "pool-balancer-type", "", "The balancing strategy to use when creating runners in pools matching requested labels.")
 
 	orgWebhookInstallCmd.Flags().BoolVar(&insecureOrgWebhook, "insecure", false, "Ignore self signed certificate errors.")
 	orgWebhookCmd.AddCommand(
@@ -337,10 +341,10 @@ func init() {
 
 func formatOrganizations(orgs []params.Organization) {
 	t := table.NewWriter()
-	header := table.Row{"ID", "Name", "Credentials name", "Pool mgr running"}
+	header := table.Row{"ID", "Name", "Credentials name", "Pool Balancer Type", "Pool mgr running"}
 	t.AppendHeader(header)
 	for _, val := range orgs {
-		t.AppendRow(table.Row{val.ID, val.Name, val.CredentialsName, val.PoolManagerStatus.IsRunning})
+		t.AppendRow(table.Row{val.ID, val.Name, val.CredentialsName, val.GetBalancerType(), val.PoolManagerStatus.IsRunning})
 		t.AppendSeparator()
 	}
 	fmt.Println(t.Render())
@@ -353,6 +357,7 @@ func formatOneOrganization(org params.Organization) {
 	t.AppendHeader(header)
 	t.AppendRow(table.Row{"ID", org.ID})
 	t.AppendRow(table.Row{"Name", org.Name})
+	t.AppendRow(table.Row{"Pool balancer type", org.GetBalancerType()})
 	t.AppendRow(table.Row{"Credentials", org.CredentialsName})
 	t.AppendRow(table.Row{"Pool manager running", org.PoolManagerStatus.IsRunning})
 	if !org.PoolManagerStatus.IsRunning {

diff --git a/cmd/garm-cli/cmd/pool.go b/cmd/garm-cli/cmd/pool.go
@@ -51,6 +51,7 @@ var (
 	poolExtraSpecs             string
 	poolAll                    bool
 	poolGitHubRunnerGroup      string
+	priority                   uint
 )
 
 type poolPayloadGetter interface {
@@ -218,6 +219,7 @@ var poolAddCmd = &cobra.Command{
 			Enabled:                poolEnabled,
 			RunnerBootstrapTimeout: poolRunnerBootstrapTimeout,
 			GitHubRunnerGroup:      poolGitHubRunnerGroup,
+			Priority:               priority,
 		}
 
 		if cmd.Flags().Changed("extra-specs") {
@@ -326,6 +328,9 @@ explicitly remove them using the runner delete command.
 		if cmd.Flags().Changed("max-runners") {
 			poolUpdateParams.MaxRunners = &poolMaxRunners
 		}
+		if cmd.Flags().Changed("priority") {
+			poolUpdateParams.Priority = &priority
+		}
 
 		if cmd.Flags().Changed("min-idle-runners") {
 			poolUpdateParams.MinIdleRunners = &poolMinIdleRunners
@@ -385,6 +390,7 @@ func init() {
 	poolListCmd.MarkFlagsMutuallyExclusive("repo", "org", "all", "enterprise")
 
 	poolUpdateCmd.Flags().StringVar(&poolImage, "image", "", "The provider-specific image name to use for runners in this pool.")
+	poolUpdateCmd.Flags().UintVar(&priority, "priority", 0, "When multiple pools match the same labels, priority dictates the order by which they are returned, in descending order.")
 	poolUpdateCmd.Flags().StringVar(&poolFlavor, "flavor", "", "The flavor to use for this runner.")
 	poolUpdateCmd.Flags().StringVar(&poolTags, "tags", "", "A comma separated list of tags to assign to this runner.")
 	poolUpdateCmd.Flags().StringVar(&poolOSType, "os-type", "linux", "Operating system type (windows, linux, etc).")
@@ -400,6 +406,7 @@ func init() {
 	poolUpdateCmd.MarkFlagsMutuallyExclusive("extra-specs-file", "extra-specs")
 
 	poolAddCmd.Flags().StringVar(&poolProvider, "provider-name", "", "The name of the provider where runners will be created.")
+	poolAddCmd.Flags().UintVar(&priority, "priority", 0, "When multiple pools match the same labels, priority dictates the order by which they are returned, in descending order.")
 	poolAddCmd.Flags().StringVar(&poolImage, "image", "", "The provider-specific image name to use for runners in this pool.")
 	poolAddCmd.Flags().StringVar(&poolFlavor, "flavor", "", "The flavor to use for this runner.")
 	poolAddCmd.Flags().StringVar(&poolRunnerPrefix, "runner-prefix", "", "The name prefix to use for runners in this pool.")
@@ -462,7 +469,7 @@ func asRawMessage(data []byte) (json.RawMessage, error) {
 
 func formatPools(pools []params.Pool) {
 	t := table.NewWriter()
-	header := table.Row{"ID", "Image", "Flavor", "Tags", "Belongs to", "Level", "Enabled", "Runner Prefix"}
+	header := table.Row{"ID", "Image", "Flavor", "Tags", "Belongs to", "Level", "Enabled", "Runner Prefix", "Priority"}
 	t.AppendHeader(header)
 
 	for _, pool := range pools {
@@ -484,7 +491,7 @@ func formatPools(pools []params.Pool) {
 			belongsTo = pool.EnterpriseName
 			level = "enterprise"
 		}
-		t.AppendRow(table.Row{pool.ID, pool.Image, pool.Flavor, strings.Join(tags, " "), belongsTo, level, pool.Enabled, pool.GetRunnerPrefix()})
+		t.AppendRow(table.Row{pool.ID, pool.Image, pool.Flavor, strings.Join(tags, " "), belongsTo, level, pool.Enabled, pool.GetRunnerPrefix(), pool.Priority})
 		t.AppendSeparator()
 	}
 	fmt.Println(t.Render())
@@ -519,6 +526,7 @@ func formatOnePool(pool params.Pool) {
 	t.AppendHeader(header)
 	t.AppendRow(table.Row{"ID", pool.ID})
 	t.AppendRow(table.Row{"Provider Name", pool.ProviderName})
+	t.AppendRow(table.Row{"Priority", pool.Priority})
 	t.AppendRow(table.Row{"Image", pool.Image})
 	t.AppendRow(table.Row{"Flavor", pool.Flavor})
 	t.AppendRow(table.Row{"OS Type", pool.OSType})

diff --git a/cmd/garm-cli/cmd/repository.go b/cmd/garm-cli/cmd/repository.go
@@ -164,10 +164,11 @@ var repoAddCmd = &cobra.Command{
 
 		newRepoReq := apiClientRepos.NewCreateRepoParams()
 		newRepoReq.Body = params.CreateRepoParams{
-			Owner:           repoOwner,
-			Name:            repoName,
-			WebhookSecret:   repoWebhookSecret,
-			CredentialsName: repoCreds,
+			Owner:            repoOwner,
+			Name:             repoName,
+			WebhookSecret:    repoWebhookSecret,
+			CredentialsName:  repoCreds,
+			PoolBalancerType: params.PoolBalancerType(poolBalancerType),
 		}
 		response, err := apiCli.Repositories.CreateRepo(newRepoReq, authToken)
 		if err != nil {
@@ -236,8 +237,9 @@ var repoUpdateCmd = &cobra.Command{
 		}
 		updateReposReq := apiClientRepos.NewUpdateRepoParams()
 		updateReposReq.Body = params.UpdateEntityParams{
-			WebhookSecret:   repoWebhookSecret,
-			CredentialsName: repoCreds,
+			WebhookSecret:    repoWebhookSecret,
+			CredentialsName:  repoCreds,
+			PoolBalancerType: params.PoolBalancerType(poolBalancerType),
 		}
 		updateReposReq.RepoID = args[0]
 
@@ -304,6 +306,7 @@ var repoDeleteCmd = &cobra.Command{
 
 func init() {
 	repoAddCmd.Flags().StringVar(&repoOwner, "owner", "", "The owner of this repository")
+	repoAddCmd.Flags().StringVar(&poolBalancerType, "pool-balancer-type", string(params.PoolBalancerTypeRoundRobin), "The balancing strategy to use when creating runners in pools matching requested labels.")
 	repoAddCmd.Flags().StringVar(&repoName, "name", "", "The name of the repository")
 	repoAddCmd.Flags().StringVar(&repoWebhookSecret, "webhook-secret", "", "The webhook secret for this repository")
 	repoAddCmd.Flags().StringVar(&repoCreds, "credentials", "", "Credentials name. See credentials list.")
@@ -320,6 +323,7 @@ func init() {
 
 	repoUpdateCmd.Flags().StringVar(&repoWebhookSecret, "webhook-secret", "", "The webhook secret for this repository. If you update this secret, you will have to manually update the secret in GitHub as well.")
 	repoUpdateCmd.Flags().StringVar(&repoCreds, "credentials", "", "Credentials name. See credentials list.")
+	repoUpdateCmd.Flags().StringVar(&poolBalancerType, "pool-balancer-type", "", "The balancing strategy to use when creating runners in pools matching requested labels.")
 
 	repoWebhookInstallCmd.Flags().BoolVar(&insecureRepoWebhook, "insecure", false, "Ignore self signed certificate errors.")
 
@@ -343,10 +347,10 @@ func init() {
 
 func formatRepositories(repos []params.Repository) {
 	t := table.NewWriter()
-	header := table.Row{"ID", "Owner", "Name", "Credentials name", "Pool mgr running"}
+	header := table.Row{"ID", "Owner", "Name", "Credentials name", "Pool Balancer Type", "Pool mgr running"}
 	t.AppendHeader(header)
 	for _, val := range repos {
-		t.AppendRow(table.Row{val.ID, val.Owner, val.Name, val.CredentialsName, val.PoolManagerStatus.IsRunning})
+		t.AppendRow(table.Row{val.ID, val.Owner, val.Name, val.CredentialsName, val.GetBalancerType(), val.PoolManagerStatus.IsRunning})
 		t.AppendSeparator()
 	}
 	fmt.Println(t.Render())
@@ -360,6 +364,7 @@ func formatOneRepository(repo params.Repository) {
 	t.AppendRow(table.Row{"ID", repo.ID})
 	t.AppendRow(table.Row{"Owner", repo.Owner})
 	t.AppendRow(table.Row{"Name", repo.Name})
+	t.AppendRow(table.Row{"Pool balancer type", repo.GetBalancerType()})
 	t.AppendRow(table.Row{"Credentials", repo.CredentialsName})
 	t.AppendRow(table.Row{"Pool manager running", repo.PoolManagerStatus.IsRunning})
 	if !repo.PoolManagerStatus.IsRunning {

diff --git a/cmd/garm-cli/cmd/root.go b/cmd/garm-cli/cmd/root.go
@@ -38,6 +38,7 @@ var (
 	authToken         runtime.ClientAuthInfoWriter
 	needsInit         bool
 	debug             bool
+	poolBalancerType  string
 	errNeedsInitError = fmt.Errorf("please log into a garm installation first")
 )
 

diff --git a/database/common/common.go b/database/common/common.go
@@ -21,7 +21,7 @@ import (
 )
 
 type RepoStore interface {
-	CreateRepository(ctx context.Context, owner, name, credentialsName, webhookSecret string) (params.Repository, error)
+	CreateRepository(ctx context.Context, owner, name, credentialsName, webhookSecret string, poolBalancerType params.PoolBalancerType) (params.Repository, error)
 	GetRepository(ctx context.Context, owner, name string) (params.Repository, error)
 	GetRepositoryByID(ctx context.Context, repoID string) (params.Repository, error)
 	ListRepositories(ctx context.Context) ([]params.Repository, error)
@@ -40,7 +40,7 @@ type RepoStore interface {
 }
 
 type OrgStore interface {
-	CreateOrganization(ctx context.Context, name, credentialsName, webhookSecret string) (params.Organization, error)
+	CreateOrganization(ctx context.Context, name, credentialsName, webhookSecret string, poolBalancerType params.PoolBalancerType) (params.Organization, error)
 	GetOrganization(ctx context.Context, name string) (params.Organization, error)
 	GetOrganizationByID(ctx context.Context, orgID string) (params.Organization, error)
 	ListOrganizations(ctx context.Context) ([]params.Organization, error)
@@ -58,7 +58,7 @@ type OrgStore interface {
 }
 
 type EnterpriseStore interface {
-	CreateEnterprise(ctx context.Context, name, credentialsName, webhookSecret string) (params.Enterprise, error)
+	CreateEnterprise(ctx context.Context, name, credentialsName, webhookSecret string, poolBalancerType params.PoolBalancerType) (params.Enterprise, error)
 	GetEnterprise(ctx context.Context, name string) (params.Enterprise, error)
 	GetEnterpriseByID(ctx context.Context, enterpriseID string) (params.Enterprise, error)
 	ListEnterprises(ctx context.Context) ([]params.Enterprise, error)