From 2c0d4e093d00514835245fc45995fcb0ed341093 Mon Sep 17 00:00:00 2001
From: Piotr Dyraga <piotr.dyraga@akena.engineering>
Date: Thu, 26 Jul 2018 11:12:23 +0200
Subject: [PATCH] Wait for all goroutines to finish

Initially, we did not wait for all goroutines to finish which was fine
as long as we did not run all tests with `go test` and observe deadlock
in 50% of cases. This is interesting, since deadlock never occurred when
running safe_prime_generator_test.go. Anyway, waiting for all go
routines to finish shouldn't harm the performance a lot and it gives a
guarantee we did not miss anything.
---
 safe_prime_generator.go | 20 +++++++++++++++-----
 1 file changed, 15 insertions(+), 5 deletions(-)

diff --git a/safe_prime_generator.go b/safe_prime_generator.go
index c947041..8fa65a2 100644
--- a/safe_prime_generator.go
+++ b/safe_prime_generator.go
@@ -63,7 +63,7 @@ func GenerateSafePrime(
 	bitLen int,
 	concurrencyLevel int,
 	timeout time.Duration,
-) (*big.Int, *big.Int, error) {
+) (p *big.Int, q *big.Int, err error) {
 	if bitLen < 6 {
 		return nil, nil, errors.New("safe prime size must be at least 6 bits")
 	}
@@ -75,11 +75,15 @@ func GenerateSafePrime(
 	defer close(errChan)
 
 	mutex := &sync.Mutex{}
+	waitGroup := &sync.WaitGroup{}
+	waitGroup.Add(concurrencyLevel)
 
 	ctx, cancel := context.WithCancel(context.Background())
 
 	for i := 0; i < concurrencyLevel; i++ {
-		runGenPrimeRoutine(ctx, primeChan, errChan, mutex, rand.Reader, bitLen)
+		runGenPrimeRoutine(
+			ctx, primeChan, errChan, mutex, waitGroup, rand.Reader, bitLen,
+		)
 	}
 
 	// Cancel after the specified timeout.
@@ -95,15 +99,18 @@ func GenerateSafePrime(
 		mutex.Lock()
 		cancel()
 		mutex.Unlock()
-		return result.p, result.q, nil
+		p, q, err = result.p, result.q, nil
 	case err := <-errChan:
 		mutex.Lock()
 		cancel()
 		mutex.Unlock()
-		return nil, nil, err
+		p, q, err = nil, nil, err
 	case <-ctx.Done():
-		return nil, nil, fmt.Errorf("generator timed out after %v", timeout)
+		p, q, err = nil, nil, fmt.Errorf("generator timed out after %v", timeout)
 	}
+
+	waitGroup.Wait()
+	return
 }
 
 type safePrime struct {
@@ -151,6 +158,7 @@ func runGenPrimeRoutine(
 	primeChan chan safePrime,
 	errChan chan error,
 	mutex *sync.Mutex,
+	waitGroup *sync.WaitGroup,
 	rand io.Reader,
 	pBitLen int,
 ) {
@@ -170,6 +178,7 @@ func runGenPrimeRoutine(
 		for {
 			select {
 			case <-ctx.Done():
+				waitGroup.Done()
 				return
 			default:
 				_, err := io.ReadFull(rand, bytes)
@@ -253,6 +262,7 @@ func runGenPrimeRoutine(
 					}
 					mutex.Unlock()
 
+					waitGroup.Done()
 					return
 				}
 			}