fix: make WaitGroupContext reusable (#28)

Author: reugn

wait_group_context.go

@@ -11,28 +11,28 @@ import (
 // time, Wait can be used to block until all goroutines have finished or the
 // given context is done.
 type WaitGroupContext struct {
-	ctx     context.Context
-	done    chan struct{}
-	counter atomic.Int32
-	state   atomic.Int32
+	ctx   context.Context
+	sem   chan struct{}
+	state atomic.Uint64 // high 32 bits are counter, low 32 bits are waiter count.
 }
 
 // NewWaitGroupContext returns a new WaitGroupContext with Context ctx.
 func NewWaitGroupContext(ctx context.Context) *WaitGroupContext {
 	return &WaitGroupContext{
-		ctx:  ctx,
-		done: make(chan struct{}),
+		ctx: ctx,
+		sem: make(chan struct{}),
 	}
 }
 
 // Add adds delta, which may be negative, to the WaitGroupContext counter.
 // If the counter becomes zero, all goroutines blocked on Wait are released.
 // If the counter goes negative, Add panics.
 func (wgc *WaitGroupContext) Add(delta int) {
-	counter := wgc.counter.Add(int32(delta))
-	if counter == 0 && wgc.state.CompareAndSwap(0, 1) {
-		wgc.release()
-	} else if counter < 0 && wgc.state.Load() == 0 {
+	state := wgc.state.Add(uint64(delta) << 32)
+	counter := int32(state >> 32)
+	if counter == 0 {
+		wgc.notifyAll()
+	} else if counter < 0 {
 		panic("async: negative WaitGroupContext counter")
 	}
 }
@@ -44,12 +44,36 @@ func (wgc *WaitGroupContext) Done() {
 
 // Wait blocks until the wait group counter is zero or ctx is done.
 func (wgc *WaitGroupContext) Wait() {
-	select {
-	case <-wgc.ctx.Done():
-	case <-wgc.done:
+	for {
+		state := wgc.state.Load()
+		counter := int32(state >> 32)
+		if counter == 0 {
+			return
+		}
+		if wgc.state.CompareAndSwap(state, state+1) {
+			select {
+			case <-wgc.sem:
+				if wgc.state.Load() != 0 {
+					panic("async: WaitGroupContext is reused before " +
+						"previous Wait has returned")
+				}
+			case <-wgc.ctx.Done():
+			}
+			return
+		}
 	}
 }
 
-func (wgc *WaitGroupContext) release() {
-	close(wgc.done)
+// notifyAll releases all goroutines blocked in Wait and resets
+// the wait group state.
+func (wgc *WaitGroupContext) notifyAll() {
+	state := wgc.state.Load()
+	waiting := uint32(state)
+	wgc.state.Store(0)
+	for ; waiting != 0; waiting-- {
+		select {
+		case wgc.sem <- struct{}{}:
+		case <-wgc.ctx.Done():
+		}
+	}
 }

wait_group_context_test.go

@@ -67,10 +67,56 @@ func TestWaitGroupContextCanceled(t *testing.T) {
 	assert.Equal(t, int(result.Load()), 111)
 }
 
-func TestWaitGroupContextPanic(t *testing.T) {
+func TestWaitGroupContextPanicNegativeCounter(t *testing.T) {
 	negativeCounter := func() {
 		wgc := NewWaitGroupContext(context.Background())
 		wgc.Add(-2)
 	}
 	assert.Panic(t, negativeCounter)
 }
+
+func TestWaitGroupContextPanicReused(t *testing.T) {
+	reusedBeforeWaitReturned := func() {
+		var result atomic.Int32
+		wgc := NewWaitGroupContext(context.Background())
+
+		n := 10
+		for i := 0; i < n; i++ {
+			wgc.Add(1)
+			go func() {
+				defer wgc.Add(1)
+				defer wgc.Done()
+				result.Add(1)
+			}()
+			wgc.Wait()
+		}
+	}
+	assert.Panic(t, reusedBeforeWaitReturned)
+}
+
+func TestWaitGroupContextReused(t *testing.T) {
+	var result atomic.Int32
+	wgc := NewWaitGroupContext(context.Background())
+
+	n := 1000
+	for i := 0; i < n; i++ {
+		assert.Equal(t, int(result.Load()), i*3)
+		wgc.Add(2)
+		go func() {
+			defer wgc.Done()
+			result.Add(1)
+		}()
+		go func() {
+			defer wgc.Done()
+			result.Add(1)
+		}()
+		go func() {
+			wgc.Wait()
+			result.Add(1)
+		}()
+		wgc.Wait()
+		time.Sleep(time.Millisecond)
+	}
+
+	assert.Equal(t, int(result.Load()), n*3)
+}