Add custom WaitGroup

Author: cschleiden

internal/sync/waitgroup.go

@@ -0,0 +1,51 @@
+package sync
+
+type WaitGroup interface {
+	Add(delta int)
+	Done()
+	Wait(ctx Context)
+}
+
+type waitGroup struct {
+	n       int
+	f       Future
+	waiting bool
+}
+
+func NewWaitGroup() WaitGroup {
+	return &waitGroup{
+		f: NewFuture(),
+	}
+}
+
+func (wg *waitGroup) Wait(ctx Context) {
+	wg.waiting = true
+
+	if err := wg.f.Get(ctx, nil); err != nil {
+		panic(err)
+	}
+}
+
+func (wg *waitGroup) Add(delta int) {
+	wg.n += delta
+
+	if wg.n < 0 {
+		panic("negative WaitGroup counter")
+	}
+
+	if wg.waiting && delta > 0 && wg.n == delta {
+		panic("WaitGroup misuse: Add called concurrently with Wait")
+	}
+
+	if wg.n > 0 || !wg.waiting {
+		return
+	}
+
+	if wg.n == 0 {
+		wg.f.Set(nil, nil)
+	}
+}
+
+func (wg *waitGroup) Done() {
+	wg.Add(-1)
+}

internal/sync/waitgroup_test.go

@@ -0,0 +1,50 @@
+package sync
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func Test_WaitGroup_PanicsForInvalidCounters(t *testing.T) {
+	wg := NewWaitGroup()
+
+	require.PanicsWithValue(t, "negative WaitGroup counter", func() {
+		wg.Add(-2)
+	})
+}
+
+func Test_WaitGroup_Blocks(t *testing.T) {
+	s := NewScheduler()
+	ctx := Background()
+
+	wg := NewWaitGroup()
+	wg.Add(2)
+
+	s.NewCoroutine(ctx, func(ctx Context) error {
+		wg.Wait(ctx)
+
+		return nil
+	})
+
+	s.Execute(ctx)
+	require.Equal(t, 1, s.RunningCoroutines())
+
+	s.NewCoroutine(ctx, func(ctx Context) error {
+		wg.Done()
+
+		return nil
+	})
+
+	s.Execute(ctx)
+	require.Equal(t, 1, s.RunningCoroutines())
+
+	s.NewCoroutine(ctx, func(ctx Context) error {
+		wg.Done()
+
+		return nil
+	})
+
+	s.Execute(ctx)
+	require.Equal(t, 0, s.RunningCoroutines())
+}

workflow/sync.go

@@ -7,5 +7,6 @@ import (
 type Future = sync.Future
 type Channel = sync.Channel
 type Context = sync.Context
+type WaitGroup = sync.WaitGroup
 
 var Canceled = sync.Canceled

workflow/workflow.go

@@ -90,3 +90,7 @@ func NewChannel() Channel {
 func NewBufferedChannel(size int) Channel {
 	return sync.NewBufferedChannel(size)
 }
+
+func NewWaitGroup() WaitGroup {
+	return sync.NewWaitGroup()
+}