Limit concurrency in phases (#431)

kke · web-flow · commit 0e8e8c566b36 · 2022-12-07T10:55:42.000+02:00
* Limit concurrency in phases

Signed-off-by: Kimmo Lehto &lt;klehto@mirantis.com&gt;
diff --git a/cmd/apply.go b/cmd/apply.go
@@ -18,6 +18,8 @@ var applyCommand = &cli.Command{
 	Usage: "Apply a k0sctl configuration",
 	Flags: []cli.Flag{
 		configFlag,
+		concurrencyFlag,
+		concurrentUploadsFlag,
 		&cli.BoolFlag{
 			Name:  "no-wait",
 			Usage: "Do not wait for worker nodes to join",
@@ -57,7 +59,7 @@ var applyCommand = &cli.Command{
 		start := time.Now()
 		phase.NoWait = ctx.Bool("no-wait")
 
-		manager := phase.Manager{Config: ctx.Context.Value(ctxConfigKey{}).(*v1beta1.Cluster)}
+		manager := phase.Manager{Config: ctx.Context.Value(ctxConfigKey{}).(*v1beta1.Cluster), Concurrency: ctx.Int("concurrency"), ConcurrentUploads: ctx.Int("concurrent-uploads")}
 		lockPhase := &phase.Lock{}
 
 		manager.AddPhase(
diff --git a/cmd/backup.go b/cmd/backup.go
@@ -16,6 +16,7 @@ var backupCommand = &cli.Command{
 	Usage: "Take backup of existing clusters state",
 	Flags: []cli.Flag{
 		configFlag,
+		concurrencyFlag,
 		debugFlag,
 		traceFlag,
 		redactFlag,
@@ -27,7 +28,7 @@ var backupCommand = &cli.Command{
 	Action: func(ctx *cli.Context) error {
 		start := time.Now()
 
-		manager := phase.Manager{Config: ctx.Context.Value(ctxConfigKey{}).(*v1beta1.Cluster)}
+		manager := phase.Manager{Config: ctx.Context.Value(ctxConfigKey{}).(*v1beta1.Cluster), Concurrency: ctx.Int("concurrency")}
 		lockPhase := &phase.Lock{}
 
 		manager.AddPhase(
diff --git a/cmd/flags.go b/cmd/flags.go
@@ -71,6 +71,18 @@ var (
 		EnvVars: []string{"DISABLE_UPGRADE_CHECK"},
 	}
 
+	concurrencyFlag = &cli.IntFlag{
+		Name:  "concurrency",
+		Usage: "Maximum number of hosts to configure in parallel, set to 0 for unlimited",
+		Value: 30,
+	}
+
+	concurrentUploadsFlag = &cli.IntFlag{
+		Name:  "concurrent-uploads",
+		Usage: "Maximum number of files to upload in parallel, set to 0 for unlimited",
+		Value: 5,
+	}
+
 	Colorize = aurora.NewAurora(false)
 )
 
diff --git a/cmd/kubeconfig.go b/cmd/kubeconfig.go
@@ -20,6 +20,7 @@ var kubeconfigCommand = &cli.Command{
 			Value: "",
 		},
 		configFlag,
+		concurrencyFlag,
 		debugFlag,
 		traceFlag,
 		redactFlag,
@@ -36,7 +37,7 @@ var kubeconfigCommand = &cli.Command{
 		// Change so that the internal config has only single controller host as we
 		// do not need to connect to all nodes
 		c.Spec.Hosts = cluster.Hosts{c.Spec.K0sLeader()}
-		manager := phase.Manager{Config: c}
+		manager := phase.Manager{Config: c, Concurrency: ctx.Int("concurrency")}
 
 		manager.AddPhase(
 			&phase.Connect{},
diff --git a/cmd/reset.go b/cmd/reset.go
@@ -20,6 +20,7 @@ var resetCommand = &cli.Command{
 	Usage: "Remove traces of k0s from all of the hosts",
 	Flags: []cli.Flag{
 		configFlag,
+		concurrencyFlag,
 		debugFlag,
 		traceFlag,
 		redactFlag,
@@ -50,7 +51,7 @@ var resetCommand = &cli.Command{
 
 		start := time.Now()
 
-		manager := phase.Manager{Config: ctx.Context.Value(ctxConfigKey{}).(*v1beta1.Cluster)}
+		manager := phase.Manager{Config: ctx.Context.Value(ctxConfigKey{}).(*v1beta1.Cluster), Concurrency: ctx.Int("concurrency")}
 		for _, h := range manager.Config.Spec.Hosts {
 			h.Reset = true
 		}
diff --git a/go.mod b/go.mod
@@ -28,7 +28,7 @@ require (
 	github.com/shiena/ansicolor v0.0.0-20200904210342-c7312218db18
 	github.com/sirupsen/logrus v1.8.1
 	github.com/stretchr/testify v1.8.0
-	github.com/urfave/cli/v2 v2.3.0
+	github.com/urfave/cli/v2 v2.23.6
 	github.com/xtgo/uuid v0.0.0-20140804021211-a0b114877d4c // indirect
 	golang.org/x/crypto v0.3.0 // indirect
 	golang.org/x/net v0.2.0 // indirect
@@ -82,6 +82,7 @@ require (
 	github.com/pmezard/go-difflib v1.0.0 // indirect
 	github.com/russross/blackfriday/v2 v2.1.0 // indirect
 	github.com/spf13/pflag v1.0.5 // indirect
+	github.com/xrash/smetrics v0.0.0-20201216005158-039620a65673 // indirect
 	golang.org/x/mod v0.7.0 // indirect
 	golang.org/x/oauth2 v0.0.0-20210819190943-2bc19b11175f // indirect
 	golang.org/x/time v0.0.0-20210723032227-1f47c861a9ac // indirect
diff --git a/go.sum b/go.sum
@@ -343,6 +343,10 @@ github.com/stretchr/testify v1.8.0 h1:pSgiaMZlXftHpm5L7V1+rVB+AZJydKsMxsQBIJw4PK
 github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU=
 github.com/urfave/cli/v2 v2.3.0 h1:qph92Y649prgesehzOrQjdWyxFOp/QVM+6imKHad91M=
 github.com/urfave/cli/v2 v2.3.0/go.mod h1:LJmUH05zAU44vOAcrfzZQKsZbVcdbOG8rtL3/XcUArI=
+github.com/urfave/cli/v2 v2.23.6 h1:iWmtKD+prGo1nKUtLO0Wg4z9esfBM4rAV4QRLQiEmJ4=
+github.com/urfave/cli/v2 v2.23.6/go.mod h1:GHupkWPMM0M/sj1a2b4wUrWBPzazNrIjouW6fmdJLxc=
+github.com/xrash/smetrics v0.0.0-20201216005158-039620a65673 h1:bAn7/zixMGCfxrRTfdpNzjtPYqr8smhKouy9mxVdGPU=
+github.com/xrash/smetrics v0.0.0-20201216005158-039620a65673/go.mod h1:N3UwUGtsrSj3ccvlPHLoLsHnpR27oXr4ZE984MbSER8=
 github.com/xtgo/uuid v0.0.0-20140804021211-a0b114877d4c h1:3lbZUMbMiGUW/LMkfsEABsc5zNT9+b1CvsJx47JzJ8g=
 github.com/xtgo/uuid v0.0.0-20140804021211-a0b114877d4c/go.mod h1:UrdRz5enIKZ63MEE3IF9l2/ebyx59GyGgPi+tICQdmM=
 github.com/yuin/goldmark v1.1.25/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=
diff --git a/phase/arm_prepare.go b/phase/arm_prepare.go
@@ -40,7 +40,7 @@ func (p *PrepareArm) ShouldRun() bool {
 
 // Run the phase
 func (p *PrepareArm) Run() error {
-	return p.hosts.ParallelEach(p.etcdUnsupportedArch)
+	return p.parallelDo(p.hosts, p.etcdUnsupportedArch)
 }
 
 func (p *PrepareArm) etcdUnsupportedArch(h *cluster.Host) error {
diff --git a/phase/configure_k0s.go b/phase/configure_k0s.go
@@ -59,7 +59,7 @@ func (p *ConfigureK0s) Run() error {
 	}
 
 	controllers := p.Config.Spec.Hosts.Controllers()
-	return controllers.ParallelEach(p.configureK0s)
+	return p.parallelDo(controllers, p.configureK0s)
 }
 
 func (p *ConfigureK0s) validateConfig(h *cluster.Host) error {
diff --git a/phase/connect.go b/phase/connect.go
@@ -24,7 +24,7 @@ var retries = uint(60)
 
 // Run the phase
 func (p *Connect) Run() error {
-	return p.Config.Spec.Hosts.ParallelEach(func(h *cluster.Host) error {
+	return p.parallelDo(p.Config.Spec.Hosts, func(h *cluster.Host) error {
 		err := retry.Do(
 			func() error {
 				return h.Connect()
diff --git a/phase/detect_os.go b/phase/detect_os.go
@@ -25,7 +25,7 @@ func (p *DetectOS) Title() string {
 
 // Run the phase
 func (p *DetectOS) Run() error {
-	return p.Config.Spec.Hosts.ParallelEach(func(h *cluster.Host) error {
+	return p.parallelDo(p.Config.Spec.Hosts, func(h *cluster.Host) error {
 		if h.OSIDOverride != "" {
 			log.Infof("%s: OS ID has been manually set to %s", h, h.OSIDOverride)
 		}
diff --git a/phase/download_k0s.go b/phase/download_k0s.go
@@ -56,7 +56,7 @@ func (p *DownloadK0s) ShouldRun() bool {
 
 // Run the phase
 func (p *DownloadK0s) Run() error {
-	return p.hosts.ParallelEach(p.downloadK0s)
+	return p.parallelDo(p.hosts, p.downloadK0s)
 }
 
 func (p *DownloadK0s) downloadK0s(h *cluster.Host) error {
diff --git a/phase/gather_facts.go b/phase/gather_facts.go
@@ -21,7 +21,7 @@ func (p *GatherFacts) Title() string {
 
 // Run the phase
 func (p *GatherFacts) Run() error {
-	return p.Config.Spec.Hosts.ParallelEach(p.investigateHost)
+	return p.parallelDo(p.Config.Spec.Hosts, p.investigateHost)
 }
 
 func (p *GatherFacts) investigateHost(h *cluster.Host) error {
diff --git a/phase/gather_k0s_facts.go b/phase/gather_k0s_facts.go
@@ -50,7 +50,7 @@ func (p *GatherK0sFacts) Title() string {
 // Run the phase
 func (p *GatherK0sFacts) Run() error {
 	var controllers cluster.Hosts = p.Config.Spec.Hosts.Controllers()
-	if err := controllers.ParallelEach(p.investigateK0s); err != nil {
+	if err := p.parallelDo(controllers, p.investigateK0s); err != nil {
 		return err
 	}
 	p.leader = p.Config.Spec.K0sLeader()
@@ -61,7 +61,7 @@ func (p *GatherK0sFacts) Run() error {
 	}
 
 	var workers cluster.Hosts = p.Config.Spec.Hosts.Workers()
-	if err := workers.ParallelEach(p.investigateK0s); err != nil {
+	if err := p.parallelDo(workers, p.investigateK0s); err != nil {
 		return err
 	}
 
diff --git a/phase/generic_phase.go b/phase/generic_phase.go
@@ -3,12 +3,15 @@ package phase
 import (
 	"github.com/k0sproject/k0sctl/analytics"
 	"github.com/k0sproject/k0sctl/pkg/apis/k0sctl.k0sproject.io/v1beta1"
+	"github.com/k0sproject/k0sctl/pkg/apis/k0sctl.k0sproject.io/v1beta1/cluster"
 )
 
 // GenericPhase is a basic phase which gets a config via prepare, sets it into p.Config
 type GenericPhase struct {
 	analytics.Phase
 	Config *v1beta1.Cluster
+
+	manager *Manager
 }
 
 // GetConfig is an accessor to phase Config
@@ -21,3 +24,22 @@ func (p *GenericPhase) Prepare(c *v1beta1.Cluster) error {
 	p.Config = c
 	return nil
 }
+
+// SetManager adds a reference to the phase manager
+func (p *GenericPhase) SetManager(m *Manager) {
+	p.manager = m
+}
+
+func (p *GenericPhase) parallelDo(hosts cluster.Hosts, funcs ...func(h *cluster.Host) error) error {
+	if p.manager.Concurrency == 0 {
+		return hosts.ParallelEach(funcs...)
+	}
+	return hosts.BatchedParallelEach(p.manager.Concurrency, funcs...)
+}
+
+func (p *GenericPhase) parallelDoUpload(hosts cluster.Hosts, funcs ...func(h *cluster.Host) error) error {
+	if p.manager.Concurrency == 0 {
+		return hosts.ParallelEach(funcs...)
+	}
+	return hosts.BatchedParallelEach(p.manager.ConcurrentUploads, funcs...)
+}
diff --git a/phase/install_workers.go b/phase/install_workers.go
@@ -55,7 +55,7 @@ func (p *InstallWorkers) Run() error {
 	url := p.Config.Spec.KubeAPIURL()
 	healthz := fmt.Sprintf("%s/healthz", url)
 
-	err := p.hosts.ParallelEach(func(h *cluster.Host) error {
+	err := p.parallelDo(p.hosts, func(h *cluster.Host) error {
 		log.Infof("%s: validating api connection to %s", h, url)
 		if err := h.WaitHTTPStatus(healthz, 200, 401); err != nil {
 			return fmt.Errorf("failed to connect from worker to kubernetes api at %s - check networking", url)
@@ -91,7 +91,7 @@ func (p *InstallWorkers) Run() error {
 		}()
 	}
 
-	return p.hosts.ParallelEach(func(h *cluster.Host) error {
+	return p.parallelDo(p.hosts, func(h *cluster.Host) error {
 		log.Infof("%s: writing join token", h)
 		if err := h.Configurer.WriteFile(h, h.K0sJoinTokenPath(), token, "0640"); err != nil {
 			return err
diff --git a/phase/lock.go b/phase/lock.go
@@ -48,7 +48,7 @@ func (p *Lock) Cancel() {
 
 // Run the phase
 func (p *Lock) Run() error {
-	if err := p.Config.Spec.Hosts.ParallelEach(p.startLock); err != nil {
+	if err := p.parallelDo(p.Config.Spec.Hosts, p.startLock); err != nil {
 		return err
 	}
 	return p.Config.Spec.Hosts.ParallelEach(p.startTicker)
diff --git a/phase/manager.go b/phase/manager.go
@@ -42,10 +42,16 @@ type withcleanup interface {
 	CleanUp()
 }
 
+type withmanager interface {
+	SetManager(*Manager)
+}
+
 // Manager executes phases to construct the cluster
 type Manager struct {
-	phases []phase
-	Config *v1beta1.Cluster
+	phases            []phase
+	Config            *v1beta1.Cluster
+	Concurrency       int
+	ConcurrentUploads int
 }
 
 // AddPhase adds a Phase to Manager
@@ -72,6 +78,10 @@ func (m *Manager) Run() error {
 	for _, p := range m.phases {
 		title := p.Title()
 
+		if p, ok := p.(withmanager); ok {
+			p.SetManager(m)
+		}
+
 		if p, ok := p.(withconfig); ok {
 			log.Debugf("Preparing phase '%s'", p.Title())
 			if err := p.Prepare(m.Config); err != nil {
diff --git a/phase/prepare_hosts.go b/phase/prepare_hosts.go
@@ -20,7 +20,7 @@ func (p *PrepareHosts) Title() string {
 
 // Run the phase
 func (p *PrepareHosts) Run() error {
-	return p.Config.Spec.Hosts.ParallelEach(p.prepareHost)
+	return p.parallelDo(p.Config.Spec.Hosts, p.prepareHost)
 }
 
 type prepare interface {
diff --git a/phase/reset_workers.go b/phase/reset_workers.go
@@ -59,7 +59,7 @@ func (p *ResetWorkers) CleanUp() {
 
 // Run the phase
 func (p *ResetWorkers) Run() error {
-	return p.hosts.ParallelEach(func(h *cluster.Host) error {
+	return p.parallelDo(p.hosts, func(h *cluster.Host) error {
 		log.Debugf("%s: draining node", h)
 		if !p.NoDrain {
 			if err := p.leader.DrainNode(&cluster.Host{
diff --git a/phase/upload_binaries.go b/phase/upload_binaries.go
@@ -59,7 +59,7 @@ func (p *UploadBinaries) ShouldRun() bool {
 
 // Run the phase
 func (p *UploadBinaries) Run() error {
-	return p.hosts.ParallelEach(p.uploadBinary)
+	return p.parallelDoUpload(p.hosts, p.uploadBinary)
 }
 
 func (p *UploadBinaries) ensureBinPath(h *cluster.Host) error {
diff --git a/phase/uploadfiles.go b/phase/uploadfiles.go
@@ -42,7 +42,7 @@ func (p *UploadFiles) ShouldRun() bool {
 
 // Run the phase
 func (p *UploadFiles) Run() error {
-	return p.Config.Spec.Hosts.ParallelEach(p.uploadFiles)
+	return p.parallelDoUpload(p.Config.Spec.Hosts, p.uploadFiles)
 }
 
 func (p *UploadFiles) uploadFiles(h *cluster.Host) error {
diff --git a/phase/validate_hosts.go b/phase/validate_hosts.go
@@ -24,7 +24,7 @@ func (p *ValidateHosts) Run() error {
 		p.hncount[h.Metadata.Hostname]++
 	}
 
-	return p.Config.Spec.Hosts.ParallelEach(p.validateUniqueHostname, p.validateSudo)
+	return p.parallelDo(p.Config.Spec.Hosts, p.validateUniqueHostname, p.validateSudo)
 }
 
 func (p *ValidateHosts) validateUniqueHostname(h *cluster.Host) error {
diff --git a/pkg/apis/k0sctl.k0sproject.io/v1beta1/cluster/hosts.go b/pkg/apis/k0sctl.k0sproject.io/v1beta1/cluster/hosts.go
@@ -134,3 +134,18 @@ func (hosts Hosts) ParallelEach(filter ...func(h *Host) error) error {
 
 	return nil
 }
+
+// BatchedParallelEach runs a function (or multiple functions chained) on every Host parallelly in groups of batchSize hosts.
+func (hosts Hosts) BatchedParallelEach(batchSize int, filter ...func(h *Host) error) error {
+	for i := 0; i < len(hosts); i += batchSize {
+		end := i + batchSize
+		if end > len(hosts) {
+			end = len(hosts)
+		}
+		if err := hosts[i:end].ParallelEach(filter...); err != nil {
+			return err
+		}
+	}
+
+	return nil
+}

Original file line number	Diff line number	Diff line change
`@@ -40,7 +40,7 @@ func (p *PrepareArm) ShouldRun() bool {`
`40`	`40`
`41`	`41`	`// Run the phase`
`42`	`42`	`func (p *PrepareArm) Run() error {`
`43`		`- return p.hosts.ParallelEach(p.etcdUnsupportedArch)`
	`43`	`+ return p.parallelDo(p.hosts, p.etcdUnsupportedArch)`
`44`	`44`	`}`
`45`	`45`
`46`	`46`	`func (p PrepareArm) etcdUnsupportedArch(h cluster.Host) error {`
Original file line number	Diff line number	Diff line change
`@@ -59,7 +59,7 @@ func (p *ConfigureK0s) Run() error {`
`59`	`59`	`}`
`60`	`60`
`61`	`61`	`controllers := p.Config.Spec.Hosts.Controllers()`
`62`		`- return controllers.ParallelEach(p.configureK0s)`
	`62`	`+ return p.parallelDo(controllers, p.configureK0s)`
`63`	`63`	`}`
`64`	`64`
`65`	`65`	`func (p ConfigureK0s) validateConfig(h cluster.Host) error {`