Merge pull request #54 from josedonizetti/refactor-waitgroup-usage

Refactor: clean up WaitGroup usage
Shopify · Oct 31, 2017 · cb1d58a · cb1d58a
2 parents 9dc109d + 70d6341
commit cb1d58a
Show file tree

Hide file tree

Showing 13 changed files with 59 additions and 33 deletions.
diff --git a/cmd/automountServiceAccountToken.go b/cmd/automountServiceAccountToken.go
@@ -1,6 +1,8 @@
 package cmd
 
 import (
+	"sync"
+
 	log "github.com/sirupsen/logrus"
 	"github.com/spf13/cobra"
 )
@@ -40,7 +42,6 @@ func auditAutomountServiceAccountToken(items Items) (results []Result) {
 	for _, result := range results {
 		result.Print()
 	}
-	defer wg.Done()
 	return
 }
 
@@ -80,11 +81,16 @@ kubeaudit rbac sat`,
 			resources = getKubeResources(kube)
 		}
 
-		count := len(resources)
-		wg.Add(count)
+		var wg sync.WaitGroup
+		wg.Add(len(resources))
+
 		for _, resource := range resources {
-			go auditAutomountServiceAccountToken(resource)
+			go func() {
+				auditAutomountServiceAccountToken(resource)
+				wg.Done()
+			}()
 		}
+
 		wg.Wait()
 	},
 }

diff --git a/cmd/automountServiceAccountToken_test.go b/cmd/automountServiceAccountToken_test.go
@@ -13,7 +13,6 @@ func init() {
 
 func TestReplicationControllerASAT(t *testing.T) {
 	fakeReplicationControllers := fakeaudit.GetReplicationControllers("fakeReplicationControllerASAT")
-	wg.Add(1)
 	results := auditAutomountServiceAccountToken(kubeAuditReplicationControllers{list: fakeReplicationControllers})
 
 	if len(results) != 2 {

diff --git a/cmd/image.go b/cmd/image.go
@@ -2,6 +2,7 @@ package cmd
 
 import (
 	"strings"
+	"sync"
 
 	log "github.com/sirupsen/logrus"
 	"github.com/spf13/cobra"
@@ -65,7 +66,6 @@ func auditImages(image imgFlags, items Items) (results []Result) {
 	for _, result := range results {
 		result.Print()
 	}
-	defer wg.Done()
 	return
 }
 
@@ -109,11 +109,16 @@ kubeaudit image -i gcr.io/google_containers/echoserver:1.7`,
 			resources = getKubeResources(kube)
 		}
 
-		count := len(resources)
-		wg.Add(count)
+		var wg sync.WaitGroup
+		wg.Add(len(resources))
+
 		for _, resource := range resources {
-			go auditImages(imgConfig, resource)
+			go func() {
+				auditImages(imgConfig, resource)
+				wg.Done()
+			}()
 		}
+
 		wg.Wait()
 	},
 }

diff --git a/cmd/image_test.go b/cmd/image_test.go
@@ -13,7 +13,6 @@ func init() {
 
 func TestPodImg(t *testing.T) {
 	fakePods := fakeaudit.GetPods("fakePodImg")
-	wg.Add(2)
 	image := imgFlags{img: "fakeContainerImg:1.5"}
 	image.splitImageString()
 	results := auditImages(image, kubeAuditPods{list: fakePods})

diff --git a/cmd/privileged.go b/cmd/privileged.go
@@ -1,6 +1,8 @@
 package cmd
 
 import (
+	"sync"
+
 	log "github.com/sirupsen/logrus"
 	"github.com/spf13/cobra"
 )
@@ -38,7 +40,6 @@ func auditPrivileged(items Items) (results []Result) {
 	for _, result := range results {
 		result.Print()
 	}
-	defer wg.Done()
 	return
 }
 
@@ -75,11 +76,16 @@ kubeaudit privileged`,
 			resources = getKubeResources(kube)
 		}
 
-		count := len(resources)
-		wg.Add(count)
+		var wg sync.WaitGroup
+		wg.Add(len(resources))
+
 		for _, resource := range resources {
-			go auditPrivileged(resource)
+			go func() {
+				auditPrivileged(resource)
+				wg.Done()
+			}()
 		}
+
 		wg.Wait()
 	},
 }

diff --git a/cmd/privileged_test.go b/cmd/privileged_test.go
@@ -13,7 +13,6 @@ func init() {
 
 func TestDaemonSetPrivileged(t *testing.T) {
 	fakeDaemonSets := fakeaudit.GetDaemonSets("fakeDaemonSetPrivileged")
-	wg.Add(1)
 	results := auditPrivileged(kubeAuditDaemonSets{list: fakeDaemonSets})
 
 	if len(results) != 3 {

diff --git a/cmd/readOnlyRootFilesystem.go b/cmd/readOnlyRootFilesystem.go
@@ -1,6 +1,8 @@
 package cmd
 
 import (
+	"sync"
+
 	log "github.com/sirupsen/logrus"
 	"github.com/spf13/cobra"
 )
@@ -36,7 +38,6 @@ func auditReadOnlyRootFS(items Items) (results []Result) {
 	for _, result := range results {
 		result.Print()
 	}
-	defer wg.Done()
 	return
 }
 
@@ -71,11 +72,16 @@ kubeaudit runAsNonRoot`,
 			resources = getKubeResources(kube)
 		}
 
-		count := len(resources)
-		wg.Add(count)
+		var wg sync.WaitGroup
+		wg.Add(len(resources))
+
 		for _, resource := range resources {
-			go auditReadOnlyRootFS(resource)
+			go func() {
+				auditReadOnlyRootFS(resource)
+				wg.Done()
+			}()
 		}
+
 		wg.Wait()
 	},
 }

diff --git a/cmd/readOnlyRootFilesystem_test.go b/cmd/readOnlyRootFilesystem_test.go
@@ -13,7 +13,6 @@ func init() {
 
 func TestStatefulSetRORF(t *testing.T) {
 	fakeStatefulSets := fakeaudit.GetStatefulSets("fakeStatefulSetRORF")
-	wg.Add(1)
 	results := auditReadOnlyRootFS(kubeAuditStatefulSets{list: fakeStatefulSets})
 
 	if len(results) != 3 {

diff --git a/cmd/runAsNonRoot.go b/cmd/runAsNonRoot.go
@@ -1,6 +1,8 @@
 package cmd
 
 import (
+	"sync"
+
 	log "github.com/sirupsen/logrus"
 	"github.com/spf13/cobra"
 )
@@ -36,7 +38,6 @@ func auditRunAsNonRoot(items Items) (results []Result) {
 	for _, result := range results {
 		result.Print()
 	}
-	defer wg.Done()
 	return
 }
 
@@ -73,11 +74,17 @@ kubeaudit runAsNonRoot`,
 			resources = getKubeResources(kube)
 		}
 
-		count := len(resources)
-		wg.Add(count)
+		var wg sync.WaitGroup
+		wg.Add(len(resources))
+
 		for _, resource := range resources {
-			go auditRunAsNonRoot(resource)
+			go func() {
+				auditRunAsNonRoot(resource)
+				wg.Done()
+			}()
+
 		}
+
 		wg.Wait()
 	},
 }

diff --git a/cmd/runAsNonRoot_test.go b/cmd/runAsNonRoot_test.go
@@ -13,7 +13,6 @@ func init() {
 
 func TestDeploymentRANR(t *testing.T) {
 	fakeDeployments := fakeaudit.GetDeployments("fakeDeploymentRANR")
-	wg.Add(1)
 	results := auditRunAsNonRoot(kubeAuditDeployments{list: fakeDeployments})
 
 	if len(results) != 3 {

diff --git a/cmd/securityContext.go b/cmd/securityContext.go
@@ -2,6 +2,7 @@ package cmd
 
 import (
 	"io/ioutil"
+	"sync"
 
 	log "github.com/sirupsen/logrus"
 	"github.com/spf13/cobra"
@@ -100,7 +101,6 @@ func auditSecurityContext(items Items) (results []Result) {
 	for _, result := range results {
 		result.Print()
 	}
-	defer wg.Done()
 	return
 }
 
@@ -137,11 +137,16 @@ kubeaudit sc rootfs`,
 			resources = getKubeResources(kube)
 		}
 
-		count := len(resources)
-		wg.Add(count)
+		var wg sync.WaitGroup
+		wg.Add(len(resources))
+
 		for _, resource := range resources {
-			go auditSecurityContext(resource)
+			go func() {
+				auditSecurityContext(resource)
+				wg.Done()
+			}()
 		}
+
 		wg.Wait()
 	},
 }

diff --git a/cmd/securityContext_test.go b/cmd/securityContext_test.go
@@ -29,7 +29,6 @@ func TestCapsNotDropped(t *testing.T) {
 
 func TestDeploymentSC(t *testing.T) {
 	fakeDeployments := fakeaudit.GetDeployments("fakeDeploymentSC")
-	wg.Add(1)
 	results := auditSecurityContext(kubeAuditDeployments{list: fakeDeployments})
 
 	if len(results) != 4 {

diff --git a/cmd/util.go b/cmd/util.go
@@ -3,15 +3,12 @@ package cmd
 import (
 	"reflect"
 	"runtime"
-	"sync"
 
 	fakeaudit "github.com/Shopify/kubeaudit/fakeaudit"
 	log "github.com/sirupsen/logrus"
 	"k8s.io/client-go/kubernetes"
 )
 
-var wg sync.WaitGroup
-
 func debugPrint() {
 	if rootConfig.verbose == "DEBUG" {
 		buf := make([]byte, 1<<16)