move the wait group to the caller

fixes #74

README.md

@@ -110,26 +110,29 @@ See `example/push` for the complete listing.
 HTTP/2 can send multiple requests over a single connection, but `service.Push` waits for a response before returning. Instead, you can wrap a `Service` in a queue to handle responses independently, allowing you to send multiple notifications at once.
 
 ```go
+var wg sync.WaitGroup
 queue := push.NewQueue(service, numWorkers)
-done := make(chan bool)
 
 // process responses (responses may be received in any order)
 go func() {
 	for resp := range queue.Responses {
 		log.Println(resp)
+		// done receiving and processing one response
+		wg.Done()
 	}
-	done <- true
 }()
 
 // send the notifications
 for i := 0; i < 100; i++ {
+	// increment count of notifications sent and queue it
+	wg.Add(1)
 	queue.Push(deviceToken, nil, b)
 }
 
-// done sending notifications, wait for all responses and shutdown
-queue.Wait()
-// all responses received, wait for all responses to be processed
-<-done
+// wait for all responses to be processed
+wg.Wait()
+// shutdown the channels and workers for the queue
+queue.Close()
 ```
 
 It's important to set up a goroutine to handle responses before sending any notifications, otherwise Push will block waiting for room to return a Response.

example/concurrent/main.go

@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"log"
 	"os"
+	"sync"
 	"time"
 
 	"github.com/RobotsAndPencils/buford/certificate"
@@ -60,7 +61,7 @@ func main() {
 	exitOnError(err)
 	service := push.NewService(client, host)
 	queue := push.NewQueue(service, workers)
-	done := make(chan bool)
+	var wg sync.WaitGroup
 
 	// process responses
 	// NOTE: Responses may be received in any order.
@@ -73,8 +74,8 @@ func main() {
 				log.Printf("(%d) device: %s, apns-id: %s", count, resp.DeviceToken, resp.ID)
 			}
 			count++
+			wg.Done()
 		}
-		done <- true
 	}()
 
 	// prepare notification(s) to send
@@ -87,11 +88,12 @@ func main() {
 	// send notifications:
 	start := time.Now()
 	for i := 0; i < number; i++ {
+		wg.Add(1)
 		queue.Push(deviceToken, nil, b)
 	}
 	// done sending notifications, wait for all responses and shutdown:
-	queue.Wait()
-	<-done
+	wg.Wait()
+	queue.Close()
 	elapsed := time.Since(start)
 
 	log.Printf("Time for %d responses: %s (%s ea.)", number, elapsed, elapsed/time.Duration(number))

push/queue.go

@@ -1,13 +1,10 @@
 package push
 
-import "sync"
-
 // Queue up notifications without waiting for the response.
 type Queue struct {
 	service       *Service
 	notifications chan notification
 	Responses     chan Response
-	wg            sync.WaitGroup
 }
 
 // notification to send.
@@ -46,17 +43,15 @@ func (q *Queue) Push(deviceToken string, headers *Headers, payload []byte) {
 		Headers:     headers,
 		Payload:     payload,
 	}
-	q.wg.Add(1)
 	q.notifications <- n
 }
 
-// Wait for all responses to be handled and then close channels.
-func (q *Queue) Wait() {
+// Close the channels for notifications and Responses and shutdown workers.
+// You should only call this after all responses have been received.
+func (q *Queue) Close() {
 	// Stop accepting new notifications and shutdown workers after existing notifications
 	// are processed:
 	close(q.notifications)
-	// Wait for all responses to be handled:
-	q.wg.Wait()
 	// Close responses channel to clean up:
 	close(q.Responses)
 }
@@ -65,6 +60,5 @@ func worker(q *Queue) {
 	for n := range q.notifications {
 		id, err := q.service.Push(n.DeviceToken, n.Headers, n.Payload)
 		q.Responses <- Response{DeviceToken: n.DeviceToken, ID: id, Err: err}
-		q.wg.Done()
 	}
 }

push/queue_test.go

@@ -5,6 +5,7 @@ import (
 	"net/http"
 	"net/http/httptest"
 	"strings"
+	"sync"
 	"testing"
 
 	"github.com/RobotsAndPencils/buford/push"
@@ -28,6 +29,7 @@ func TestQueuePush(t *testing.T) {
 
 	service := push.NewService(http.DefaultClient, server.URL)
 	queue := push.NewQueue(service, workers)
+	var wg sync.WaitGroup
 
 	go func() {
 		for resp := range queue.Responses {
@@ -37,11 +39,14 @@ func TestQueuePush(t *testing.T) {
 			if resp.ID != resp.DeviceToken {
 				t.Errorf("Expected %q == %q.", resp.ID, resp.DeviceToken)
 			}
+			wg.Done()
 		}
 	}()
 
 	for i := 0; i < number; i++ {
+		wg.Add(1)
 		queue.Push(fmt.Sprintf("%04d", i), nil, payload)
 	}
-	queue.Wait()
+	wg.Wait()
+	queue.Close()
 }