From 82f050d7b2bbfc3d7e235b074a4283952000930c Mon Sep 17 00:00:00 2001
From: Piotr Dyraga <piotr.dyraga@akena.engineering>
Date: Mon, 30 Jul 2018 15:46:16 +0200
Subject: [PATCH] Simplified concurrency mechanisms for goroutines

It does not longer make sense to have a mutex and cancelled context check
since we have a wait group and channels with buffers with the size equal
the number of workers.

Without the wait group we had to make sure no goroutine tries to write
to the closed channel. Now, channel is closed after all goroutines
completed (either by finding prime or cancelling the work), so there is
no chance someone tries to write to closed channel.
---
 safe_prime_generator.go | 21 ++++-----------------
 1 file changed, 4 insertions(+), 17 deletions(-)

diff --git a/safe_prime_generator.go b/safe_prime_generator.go
index 3a3355f..351c602 100644
--- a/safe_prime_generator.go
+++ b/safe_prime_generator.go
@@ -74,7 +74,6 @@ func GenerateSafePrime(
 	defer close(primeChan)
 	defer close(errChan)
 
-	mutex := &sync.Mutex{}
 	waitGroup := &sync.WaitGroup{}
 
 	ctx, cancel := context.WithCancel(context.Background())
@@ -82,28 +81,22 @@ func GenerateSafePrime(
 	for i := 0; i < concurrencyLevel; i++ {
 		waitGroup.Add(1)
 		runGenPrimeRoutine(
-			ctx, primeChan, errChan, mutex, waitGroup, random, bitLen,
+			ctx, primeChan, errChan, waitGroup, random, bitLen,
 		)
 	}
 
 	// Cancel after the specified timeout.
 	go func() {
 		time.Sleep(timeout)
-		mutex.Lock()
 		cancel()
-		mutex.Unlock()
 	}()
 
 	select {
 	case result := <-primeChan:
-		mutex.Lock()
 		cancel()
-		mutex.Unlock()
 		p, q, err = result.p, result.q, nil
 	case err := <-errChan:
-		mutex.Lock()
 		cancel()
-		mutex.Unlock()
 		p, q, err = nil, nil, err
 	case <-ctx.Done():
 		p, q, err = nil, nil, fmt.Errorf("generator timed out after %v", timeout)
@@ -157,7 +150,6 @@ func runGenPrimeRoutine(
 	ctx context.Context,
 	primeChan chan safePrime,
 	errChan chan error,
-	mutex *sync.Mutex,
 	waitGroup *sync.WaitGroup,
 	rand io.Reader,
 	pBitLen int,
@@ -175,10 +167,11 @@ func runGenPrimeRoutine(
 	bigMod := new(big.Int)
 
 	go func() {
+		defer waitGroup.Done()
+
 		for {
 			select {
 			case <-ctx.Done():
-				waitGroup.Done()
 				return
 			default:
 				_, err := io.ReadFull(rand, bytes)
@@ -266,13 +259,7 @@ func runGenPrimeRoutine(
 					isPocklingtonCriterionSatisfied(p) &&
 					q.BitLen() == qBitLen {
 
-					mutex.Lock()
-					if ctx.Err() == nil {
-						primeChan <- safePrime{p, q}
-					}
-					mutex.Unlock()
-
-					waitGroup.Done()
+					primeChan <- safePrime{p, q}
 					return
 				}
 			}