Merge pull request #19 from carlmontanari/uhoh-fix-concurrency

fix: fix borked concurrency bits
hellt · Dec 14, 2022 · b478396 · b478396
2 parents a676023 + 61c39b8
commit b478396
Show file tree

Hide file tree

Showing 2 changed files with 53 additions and 21 deletions.
diff --git a/commando/cmdo.go b/commando/cmdo.go
@@ -110,6 +110,11 @@ type appCfg struct {
 	commands    string                  // commands to send
 }
 
+type respTuple struct {
+	name string
+	resp []interface{}
+}
+
 // run runs the commando.
 func (app *appCfg) run() error {
 	i := &inventory{}
@@ -125,7 +130,10 @@ func (app *appCfg) run() error {
 	}
 
 	rw := app.newResponseWriter(app.output)
-	rCh := make(chan []interface{})
+
+	respCh := make(chan respTuple)
+
+	doneCh := make(chan interface{})
 
 	if app.output == fileOutput {
 		log.SetOutput(os.Stderr)
@@ -136,14 +144,15 @@ func (app *appCfg) run() error {
 	wg.Add(len(i.Devices))
 
 	for n, d := range i.Devices {
-		go app.runOperations(n, d, rCh)
-
-		resps := <-rCh
-		go app.outputResult(wg, rw, n, resps)
+		go app.runOperations(n, d, respCh)
 	}
 
+	go app.outputResult(wg, rw, respCh, doneCh)
+
 	wg.Wait()
 
+	doneCh <- nil
+
 	if app.output == fileOutput {
 		log.Infof("outputs have been saved to '%s' directory", app.outDir)
 	}
@@ -321,10 +330,13 @@ func runCommands(name string, d *device, driver *network.Driver) ([]interface{},
 func (app *appCfg) runOperations(
 	name string,
 	d *device,
-	rCh chan<- []interface{}) {
+	rCh chan<- respTuple) {
 	driver, err := app.openCoreConn(name, d)
 	if err != nil {
-		rCh <- nil
+		rCh <- respTuple{
+			name: name,
+			resp: nil,
+		}
 
 		return
 	}
@@ -333,7 +345,10 @@ func (app *appCfg) runOperations(
 
 	cfgResponses, err := runCfg(name, d, driver)
 	if err != nil {
-		rCh <- nil
+		rCh <- respTuple{
+			name: name,
+			resp: nil,
+		}
 
 		return
 	}
@@ -342,31 +357,52 @@ func (app *appCfg) runOperations(
 
 	err = runConfigs(name, d, driver)
 	if err != nil {
-		rCh <- nil
+		rCh <- respTuple{
+			name: name,
+			resp: nil,
+		}
 
 		return
 	}
 
 	cmdResponses, err := runCommands(name, d, driver)
 	if err != nil {
-		rCh <- nil
+		rCh <- respTuple{
+			name: name,
+			resp: nil,
+		}
 
 		return
 	}
 
 	responses = append(responses, cmdResponses...)
 
-	rCh <- responses
+	rCh <- respTuple{
+		name: name,
+		resp: responses,
+	}
 }
 
 func (app *appCfg) outputResult(
 	wg *sync.WaitGroup,
 	rw responseWriter,
-	name string,
-	r []interface{}) {
-	defer wg.Done()
-
-	if err := rw.WriteResponse(r, name); err != nil {
-		log.Errorf("error while writing the response: %v", err)
+	rCh chan respTuple,
+	doneCh chan interface{},
+) {
+	for {
+		select {
+		case <-doneCh:
+			return
+		case r := <-rCh:
+			if err := rw.WriteResponse(r.resp, r.name); err != nil {
+				log.Errorf("error while writing the response: %v", err)
+
+				// don't defer the wg.Done because it needs to always be decremented at each
+				// iteration!
+				wg.Done()
+			} else {
+				wg.Done()
+			}
+		}
 	}
 }
diff --git a/go.sum b/go.sum
@@ -31,10 +31,6 @@ github.com/rogpeppe/go-internal v1.6.1/go.mod h1:xXDCJY+GAPziupqXw64V24skbSoqbTE
 github.com/russross/blackfriday/v2 v2.1.0 h1:JIOH55/0cWyOuilr9/qlrm0BSXldqnqwMsf35Ld67mk=
 github.com/russross/blackfriday/v2 v2.1.0/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQDYRxCVz55jmeOWTM=
 github.com/scrapli/scrapligo v1.0.0/go.mod h1:jvRMdb90MNnswMiku8UNXj8JZaOIPhwhcqqFwr9qeoY=
-github.com/scrapli/scrapligo v1.0.2 h1:IzZPtfuuINvewd72Jjt91fr43s3y1shI4rXWnKmpC7A=
-github.com/scrapli/scrapligo v1.0.2/go.mod h1:jvRMdb90MNnswMiku8UNXj8JZaOIPhwhcqqFwr9qeoY=
-github.com/scrapli/scrapligo v1.0.3 h1:nLqW1FquCG//l8lzPS9rGIMX/9fe8v9dJoupnem3Hdc=
-github.com/scrapli/scrapligo v1.0.3/go.mod h1:jvRMdb90MNnswMiku8UNXj8JZaOIPhwhcqqFwr9qeoY=
 github.com/scrapli/scrapligo v1.1.0 h1:KjCam57kIV2rlxAQg/J1G7v/xgRHvpJF+Gjz+LXhQaI=
 github.com/scrapli/scrapligo v1.1.0/go.mod h1:jvRMdb90MNnswMiku8UNXj8JZaOIPhwhcqqFwr9qeoY=
 github.com/scrapli/scrapligocfg v1.0.0 h1:540SuGqqM6rKN87SLCfR54IageQ6s3a/ZOycGRgbbak=