From e1d46e6973b9d4e4bf3acd7fbd4136edbadc46fc Mon Sep 17 00:00:00 2001
From: adisbladis <adisbladis@gmail.com>
Date: Tue, 31 May 2022 00:08:36 +0800
Subject: [PATCH] Reintroduce the max workers flag

It turns out that this is actually useful to limit the number of open files.

Fixes https://github.com/tweag/gomod2nix/issues/50.
---
 generate/generate.go | 12 ++++++++++--
 lib/executor.go      |  9 ++++++++-
 main.go              |  3 ++-
 3 files changed, 20 insertions(+), 4 deletions(-)

diff --git a/generate/generate.go b/generate/generate.go
index 3ca63cb..41365f8 100644
--- a/generate/generate.go
+++ b/generate/generate.go
@@ -30,7 +30,7 @@ type goModDownload struct {
 	GoModSum string
 }
 
-func GeneratePkgs(directory string, goMod2NixPath string) ([]*schema.Package, error) {
+func GeneratePkgs(directory string, goMod2NixPath string, numWorkers int) ([]*schema.Package, error) {
 	goModPath := filepath.Join(directory, "go.mod")
 
 	log.WithFields(log.Fields{
@@ -81,7 +81,7 @@ func GeneratePkgs(directory string, goMod2NixPath string) ([]*schema.Package, er
 		log.Info("Done downloading dependencies")
 	}
 
-	executor := lib.NewParallellExecutor()
+	executor := lib.NewParallellExecutor(numWorkers)
 	var mux sync.Mutex
 
 	cache := schema.ReadCache(goMod2NixPath)
@@ -108,6 +108,10 @@ func GeneratePkgs(directory string, goMod2NixPath string) ([]*schema.Package, er
 		}
 
 		executor.Add(func() error {
+			log.WithFields(log.Fields{
+				"goPackagePath": goPackagePath,
+			}).Info("Calculating NAR hash")
+
 			h := sha256.New()
 			err := nar.DumpPath(h, dl.Dir)
 			if err != nil {
@@ -126,6 +130,10 @@ func GeneratePkgs(directory string, goMod2NixPath string) ([]*schema.Package, er
 
 			addPkg(pkg)
 
+			log.WithFields(log.Fields{
+				"goPackagePath": goPackagePath,
+			}).Info("Done calculating NAR hash")
+
 			return nil
 		})
 	}
diff --git a/lib/executor.go b/lib/executor.go
index 51f3680..7b605d3 100644
--- a/lib/executor.go
+++ b/lib/executor.go
@@ -9,17 +9,19 @@ type ParallellExecutor struct {
 	errChan chan error
 	wg      *sync.WaitGroup
 	mux     *sync.Mutex
+	guard   chan struct{}
 
 	// Error returned by Wait(), cached for other Wait() invocations
 	err  error
 	done bool
 }
 
-func NewParallellExecutor() *ParallellExecutor {
+func NewParallellExecutor(maxWorkers int) *ParallellExecutor {
 	return &ParallellExecutor{
 		errChan: make(chan error),
 		mux:     new(sync.Mutex),
 		wg:      new(sync.WaitGroup),
+		guard:   make(chan struct{}, maxWorkers),
 
 		err:  nil,
 		done: false,
@@ -29,8 +31,13 @@ func NewParallellExecutor() *ParallellExecutor {
 func (e *ParallellExecutor) Add(fn func() error) {
 	e.wg.Add(1)
 
+	e.guard <- struct{}{} // Block until a worker is available
+
 	go func() {
 		defer e.wg.Done()
+		defer func() {
+			<-e.guard
+		}()
 
 		err := fn()
 		if err != nil {
diff --git a/main.go b/main.go
index a0f3726..6e7be0a 100644
--- a/main.go
+++ b/main.go
@@ -13,6 +13,7 @@ import (
 func main() {
 
 	var directory = flag.String("dir", "./", "Go project directory")
+	var maxJobs = flag.Int("jobs", 10, "Number of max parallel jobs")
 	var outDirFlag = flag.String("outdir", "", "output directory (if different from project directory)")
 	flag.Parse()
 
@@ -23,7 +24,7 @@ func main() {
 
 	goMod2NixPath := filepath.Join(outDir, "gomod2nix.toml")
 	outFile := goMod2NixPath
-	pkgs, err := generate.GeneratePkgs(*directory, goMod2NixPath)
+	pkgs, err := generate.GeneratePkgs(*directory, goMod2NixPath, *maxJobs)
 	if err != nil {
 		panic(err)
 	}