refactor(broadcast): pass get wait ch to Wait callback

Signed-off-by: Christian Stewart <christian@aperture.us>

broadcast/broadcast.go

@@ -48,9 +48,10 @@ func (c *Broadcast) HoldLockMaybeAsync(cb func(broadcast func(), getWaitCh func(
 // Wait waits for the cb to return true or an error before returning.
 // When the broadcast channel is broadcasted, re-calls cb again to re-check the value.
 // cb is called while the mutex is locked.
-// Returns false, context.Canceled if ctx is canceled.
-// Return nil if and only if cb returned true, nil.
-func (c *Broadcast) Wait(ctx context.Context, cb func(broadcast func()) (bool, error)) error {
+// Returns the wait channel and any error.
+// Returns context.Canceled if ctx is canceled.
+// Return waitCh, nil if and only if cb returned true, nil.
+func (c *Broadcast) Wait(ctx context.Context, cb func(broadcast func(), getWaitCh func() <-chan struct{}) (bool, error)) error {
 	if cb == nil || ctx == nil {
 		return errors.New("cb and ctx must be set")
 	}
@@ -65,7 +66,7 @@ func (c *Broadcast) Wait(ctx context.Context, cb func(broadcast func()) (bool, e
 		var done bool
 		var err error
 		c.HoldLock(func(broadcast func(), getWaitCh func() <-chan struct{}) {
-			done, err = cb(broadcast)
+			done, err = cb(broadcast, getWaitCh)
 			if !done && err == nil {
 				waitCh = getWaitCh()
 			}

broadcast/broadcast_test.go

@@ -62,10 +62,14 @@ func ExampleBroadcast_Wait() {
 
 	ctx := context.Background()
 	var gotValue int
-	b.Wait(ctx, func(broadcast func()) (bool, error) {
+	err := b.Wait(ctx, func(broadcast func(), getWaitCh func() <-chan struct{}) (bool, error) {
 		gotValue = currValue
 		return gotValue == 9, nil
 	})
+	if err != nil {
+		fmt.Printf("failed to wait for value: %v", err.Error())
+		return
+	}
 
 	fmt.Printf("waited for value to increment: %v\n", gotValue)
 	// Output: waited for value to increment: 9