Merge pull request #5 from hstern/context

Add context awareness

sizedwaitgroup.go

@@ -10,6 +10,7 @@
 package sizedwaitgroup
 
 import (
+	"context"
 	"math"
 	"sync"
 )
@@ -47,8 +48,26 @@ func New(limit int) SizedWaitGroup {
 //
 // See sync.WaitGroup documentation for more information.
 func (s *SizedWaitGroup) Add() {
-	s.current <- struct{}{}
+	s.AddWithContext(context.Background())
+}
+
+// AddWithContext increments the internal WaitGroup counter.
+// It can be blocking if the limit of spawned goroutines
+// has been reached. It will stop blocking when Done is
+// been called, or when the context is canceled. Returns nil on
+// success or an error if the context is canceled before the lock
+// is acquired.
+//
+// See sync.WaitGroup documentation for more information.
+func (s *SizedWaitGroup) AddWithContext(ctx context.Context) error {
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	case s.current <- struct{}{}:
+		break
+	}
 	s.wg.Add(1)
+	return nil
 }
 
 // Done decrements the SizedWaitGroup counter.

sizedwaitgroup_test.go

@@ -1,6 +1,7 @@
 package sizedwaitgroup
 
 import (
+	"context"
 	"sync/atomic"
 	"testing"
 )
@@ -66,3 +67,19 @@ func TestNoThrottling(t *testing.T) {
 		t.Fatalf("%d, not all routines have been executed.", c)
 	}
 }
+
+func TestAddWithContext(t *testing.T) {
+	ctx, cancelFunc := context.WithCancel(context.TODO())
+
+	swg := New(1)
+
+	if err := swg.AddWithContext(ctx); err != nil {
+		t.Fatalf("AddContext returned error: %v", err)
+	}
+
+	cancelFunc()
+	if err := swg.AddWithContext(ctx); err != context.Canceled {
+		t.Fatalf("AddContext returned non-context.Canceled error: %v", err)
+	}
+
+}