From 632e4f5ffc2d281f824c3ded3dd1cdcbb1db6d58 Mon Sep 17 00:00:00 2001
From: Mikhail Klementev <blame@dumpstack.io>
Date: Sun, 21 May 2023 21:43:18 +0000
Subject: [PATCH] feat: parallel kernel installation

---
 kernel.go        | 20 +++++++++++--------
 kernel/kernel.go | 51 ++++++++++++++++++++++++++++--------------------
 pack.go          |  2 +-
 3 files changed, 43 insertions(+), 30 deletions(-)

diff --git a/kernel.go b/kernel.go
index f098a08..611f321 100644
--- a/kernel.go
+++ b/kernel.go
@@ -18,13 +18,14 @@ import (
 )
 
 type KernelCmd struct {
-	NoDownload bool  `help:"do not download qemu image while kernel generation"`
-	UseHost    bool  `help:"also use host kernels"`
-	Force      bool  `help:"force reinstall kernel"`
-	NoHeaders  bool  `help:"do not install kernel headers"`
-	Shuffle    bool  `help:"randomize kernels installation order"`
-	Retries    int64 `help:"amount of tries for each kernel" default:"10"`
-	Update     bool  `help:"update container"`
+	NoDownload bool `help:"do not download qemu image while kernel generation"`
+	UseHost    bool `help:"also use host kernels"`
+	Force      bool `help:"force reinstall kernel"`
+	NoHeaders  bool `help:"do not install kernel headers"`
+	Shuffle    bool `help:"randomize kernels installation order"`
+	Retries    int  `help:"amount of tries for each kernel" default:"10"`
+	Threads    int  `help:"threads for parallel installation" default:"1"`
+	Update     bool `help:"update container"`
 
 	List        KernelListCmd        `cmd:"" help:"list kernels"`
 	ListRemote  KernelListRemoteCmd  `cmd:"" help:"list remote kernels"`
@@ -95,7 +96,7 @@ func (cmd *KernelListRemoteCmd) Run(kernelCmd *KernelCmd, g *Globals) (err error
 }
 
 type KernelAutogenCmd struct {
-	Max int64 `help:"download kernels from set defined by regex in release_mask, but no more than X for each of release_mask" default:"100500"`
+	Max int `help:"download kernels from set defined by regex in release_mask, but no more than X for each of release_mask" default:"100500"`
 }
 
 func (cmd KernelAutogenCmd) Run(kernelCmd *KernelCmd, g *Globals) (err error) {
@@ -117,6 +118,7 @@ func (cmd KernelAutogenCmd) Run(kernelCmd *KernelCmd, g *Globals) (err error) {
 			g.Config.Docker.Registry,
 			g.Config.Docker.Commands,
 			cmd.Max, kernelCmd.Retries,
+			kernelCmd.Threads,
 			!kernelCmd.NoDownload,
 			kernelCmd.Force,
 			!kernelCmd.NoHeaders,
@@ -157,6 +159,7 @@ func (cmd *KernelGenallCmd) Run(kernelCmd *KernelCmd, g *Globals) (err error) {
 		g.Config.Docker.Registry,
 		g.Config.Docker.Commands,
 		math.MaxUint32, kernelCmd.Retries,
+		kernelCmd.Threads,
 		!kernelCmd.NoDownload,
 		kernelCmd.Force,
 		!kernelCmd.NoHeaders,
@@ -194,6 +197,7 @@ func (cmd *KernelInstallCmd) Run(kernelCmd *KernelCmd, g *Globals) (err error) {
 		g.Config.Docker.Registry,
 		g.Config.Docker.Commands,
 		math.MaxUint32, kernelCmd.Retries,
+		kernelCmd.Threads,
 		!kernelCmd.NoDownload,
 		kernelCmd.Force,
 		!kernelCmd.NoHeaders,
diff --git a/kernel/kernel.go b/kernel/kernel.go
index c068e48..1877601 100644
--- a/kernel/kernel.go
+++ b/kernel/kernel.go
@@ -19,6 +19,7 @@ import (
 	"time"
 
 	"github.com/naoina/toml"
+	"github.com/remeh/sizedwaitgroup"
 	"github.com/rs/zerolog/log"
 
 	"code.dumpstack.io/tools/out-of-tree/cache"
@@ -505,7 +506,7 @@ func SetSigintHandler(variable *bool) {
 
 // FIXME too many parameters
 func GenerateKernels(km config.Target, registry string,
-	commands []config.DockerCommand, max, retries int64,
+	commands []config.DockerCommand, max, retries, threads int,
 	download, force, headers, shuffle, update bool,
 	shutdown *bool) (err error) {
 
@@ -530,6 +531,9 @@ func GenerateKernels(km config.Target, registry string,
 	if shuffle {
 		pkgs = shuffleStrings(pkgs)
 	}
+
+	swg := sizedwaitgroup.New(threads)
+
 	for i, pkg := range pkgs {
 		if max <= 0 {
 			log.Print("Max is reached")
@@ -542,30 +546,35 @@ func GenerateKernels(km config.Target, registry string,
 		}
 		log.Info().Msgf("%d/%d %s", i+1, len(pkgs), pkg)
 
-		var attempt int64
-		for {
-			attempt++
+		swg.Add()
+		go func() {
+			defer swg.Done()
+			var attempt int
+			for {
+				attempt++
 
-			if *shutdown {
-				err = nil
-				return
-			}
+				if *shutdown {
+					err = nil
+					return
+				}
 
-			err = installKernel(km, pkg, force, headers)
-			if err == nil {
-				max--
-				break
-			} else if attempt >= retries {
-				log.Error().Err(err).Msg("install kernel")
-				log.Debug().Msg("skip")
-				break
-			} else {
-				log.Warn().Err(err).Msg("install kernel")
-				time.Sleep(time.Second)
-				log.Info().Msg("retry")
+				err = installKernel(km, pkg, force, headers)
+				if err == nil {
+					max--
+					break
+				} else if attempt >= retries {
+					log.Error().Err(err).Msg("install kernel")
+					log.Debug().Msg("skip")
+					break
+				} else {
+					log.Warn().Err(err).Msg("install kernel")
+					time.Sleep(time.Second)
+					log.Info().Msg("retry")
+				}
 			}
-		}
+		}()
 	}
+	swg.Wait()
 
 	return
 }
diff --git a/pack.go b/pack.go
index 9ee8375..ea86eea 100644
--- a/pack.go
+++ b/pack.go
@@ -20,7 +20,7 @@ type PackCmd struct {
 	NoDownload  bool  `help:"do not download qemu image while kernel generation"`
 	ExploitRuns int64 `default:"4" help:"amount of runs of each exploit"`
 	KernelRuns  int64 `default:"1" help:"amount of runs of each kernel"`
-	Max         int64 `help:"download random kernels from set defined by regex in release_mask, but no more than X for each of release_mask" default:"1"`
+	Max         int   `help:"download random kernels from set defined by regex in release_mask, but no more than X for each of release_mask" default:"1"`
 
 	Threads int `help:"threads" default:"4"`