From c3d87289bbb4b50b73994161417109564a97a6e6 Mon Sep 17 00:00:00 2001
From: Christian Schwarz <me@cschwarz.com>
Date: Sun, 22 Nov 2020 15:04:48 +0100
Subject: [PATCH] [#388] util/semaphore: fix TestSemaphore test

fixes #388
---
 util/semaphore/semaphore_test.go | 24 ++++++++++--------------
 1 file changed, 10 insertions(+), 14 deletions(-)

diff --git a/util/semaphore/semaphore_test.go b/util/semaphore/semaphore_test.go
index b2fb592..7c3f9c6 100644
--- a/util/semaphore/semaphore_test.go
+++ b/util/semaphore/semaphore_test.go
@@ -2,7 +2,6 @@ package semaphore
 
 import (
 	"context"
-	"sync"
 	"sync/atomic"
 	"testing"
 	"time"
@@ -15,7 +14,7 @@ import (
 
 func TestSemaphore(t *testing.T) {
 	const numGoroutines = 10
-	const concurrentSemaphore = 5
+	const concurrentSemaphore = 6
 	const sleepTime = 1 * time.Second
 
 	begin := time.Now()
@@ -26,29 +25,26 @@ func TestSemaphore(t *testing.T) {
 		beforeT, afterT uint32
 	}
 
-	ctx := context.Background()
-	defer trace.WithTaskFromStackUpdateCtx(&ctx)()
+	rootCtx, endRoot := trace.WithTaskFromStack(context.Background())
+	defer endRoot()
+	_, add, waitEnd := trace.WithTaskGroup(rootCtx, "TestSemaphore")
 
-	var wg sync.WaitGroup
-	wg.Add(numGoroutines)
 	for i := 0; i < numGoroutines; i++ {
-		go func() {
-			ctx, end := trace.WithTaskFromStack(ctx)
-			defer end()
-			defer wg.Done()
+		// not capturing i so no need for local copy
+		add(func(ctx context.Context) {
 			res, err := sem.Acquire(ctx)
 			require.NoError(t, err)
 			defer res.Release()
 			if time.Since(begin) > sleepTime {
-				atomic.AddUint32(&acquisitions.beforeT, 1)
-			} else {
 				atomic.AddUint32(&acquisitions.afterT, 1)
+			} else {
+				atomic.AddUint32(&acquisitions.beforeT, 1)
 			}
 			time.Sleep(sleepTime)
-		}()
+		})
 	}
 
-	wg.Wait()
+	waitEnd()
 
 	assert.True(t, acquisitions.beforeT == concurrentSemaphore)
 	assert.True(t, acquisitions.afterT == numGoroutines-concurrentSemaphore)