Merge pull request #53 from craigpastro/add-wait-safe-method

Add WaitAndRecover method to WaitGroup

waitgroup.go

@@ -41,3 +41,12 @@ func (h *WaitGroup) Wait() {
 	// Propagate a panic if we caught one from a child goroutine.
 	h.pc.Repanic()
 }
+
+// WaitAndRecover will block until all goroutines spawned with Go exit and
+// will return a *panics.RecoveredPanic if one of the child goroutines panics.
+func (h *WaitGroup) WaitAndRecover() *panics.RecoveredPanic {
+	h.wg.Wait()
+
+	// Return a recovered panic if we caught one from a child goroutine.
+	return h.pc.Recovered()
+}

waitgroup_test.go

@@ -98,5 +98,46 @@ func TestWaitGroup(t *testing.T) {
 			require.Panics(t, wg.Wait)
 			require.Equal(t, int64(2), i.Load())
 		})
+
+		t.Run("is caught by waitandrecover", func(t *testing.T) {
+			t.Parallel()
+			var wg WaitGroup
+			wg.Go(func() {
+				panic("super bad thing")
+			})
+			p := wg.WaitAndRecover()
+			require.Equal(t, p.Value, "super bad thing")
+		})
+
+		t.Run("one is caught by waitandrecover", func(t *testing.T) {
+			t.Parallel()
+			var wg WaitGroup
+			wg.Go(func() {
+				panic("super bad thing")
+			})
+			wg.Go(func() {
+				panic("super badder thing")
+			})
+			p := wg.WaitAndRecover()
+			require.NotNil(t, p)
+		})
+
+		t.Run("nonpanics run successfully with waitandrecover", func(t *testing.T) {
+			t.Parallel()
+			var wg WaitGroup
+			var i atomic.Int64
+			wg.Go(func() {
+				i.Add(1)
+			})
+			wg.Go(func() {
+				panic("super bad thing")
+			})
+			wg.Go(func() {
+				i.Add(1)
+			})
+			p := wg.WaitAndRecover()
+			require.Equal(t, p.Value, "super bad thing")
+			require.Equal(t, int64(2), i.Load())
+		})
 	})
 }