🐛 Fix race conditions

errgroup_test.go

@@ -255,6 +255,45 @@ func TestTryGo(t *testing.T) {
 	g.Wait()
 }
 
+func TestTryGoError(t *testing.T) {
+	g := &errgroup.Group{}
+	ch := make(chan struct{})
+	var count int32 = 0
+	err1 := errors.New("group_test: err1")
+	err2 := errors.New("group_test: err2")
+
+	ran := g.TryGo(func() error {
+		atomic.AddInt32(&count, 1)
+		ch <- struct{}{}
+		return err1
+	})
+
+	if !ran {
+		t.Error("TryGo should succeed but failed first run")
+	}
+
+	<-ch
+
+	ran = g.TryGo(func() error {
+		atomic.AddInt32(&count, 1)
+		return err2
+	})
+
+	if ran {
+		t.Error("TryGo should failed but succeeded second run")
+	}
+
+	if count != 1 {
+		t.Errorf("TryGo should have run 1 time, but ran %d times", count)
+	}
+
+	err := g.Wait()
+
+	if err != err1 {
+		t.Errorf("g.Wait() = %v, wanted %v", err, err1)
+	}
+}
+
 func TestGoLimit(t *testing.T) {
 	const limit = 10
 
@@ -277,6 +316,111 @@ func TestGoLimit(t *testing.T) {
 	}
 }
 
+func TestSetLimit(t *testing.T) {
+	g := &errgroup.Group{}
+	ch := make(chan struct{})
+	var count int32 = 0
+
+	g.SetLimit(0)
+	ran := g.TryGo(func() error {
+		atomic.AddInt32(&count, 1)
+		return nil
+	})
+	if ran {
+		t.Fatal("TryGo should fail but succeeded first run")
+	}
+
+	g.SetLimit(-1)
+	ran = g.TryGo(func() error {
+		atomic.AddInt32(&count, 1)
+		ch <- struct{}{}
+		return nil
+	})
+	if !ran {
+		t.Fatal("TryGo should succeed but failed second run")
+	}
+
+	<-ch
+
+	g.SetLimit(1)
+	ran = g.TryGo(func() error {
+		atomic.AddInt32(&count, 1)
+		ch <- struct{}{}
+		return nil
+	})
+	if !ran {
+		t.Fatal("TryGo should succeed but failed third run")
+	}
+	ran = g.TryGo(func() error {
+		atomic.AddInt32(&count, 1)
+		return nil
+	})
+	if ran {
+		t.Fatal("TryGo should fail but succeeded fourth run")
+	}
+
+	<-ch
+
+	if count != 2 {
+		t.Errorf("TryGo should have run 2 times, but ran %d times", count)
+	}
+
+	err := g.Wait()
+
+	if err != nil {
+		t.Errorf("g.Wait() = %v, wanted %v", err, nil)
+	}
+
+}
+
+func TestLimitPanic(t *testing.T) {
+	g := &errgroup.Group{}
+	ch := make(chan struct{})
+
+	g.SetLimit(2)
+
+	ran := g.TryGo(func() error {
+		ch <- struct{}{}
+		return nil
+	})
+	if !ran {
+		t.Fatal("TryGo should succeed but failed first run")
+	}
+
+	var err error
+	func() {
+		defer func() {
+			rec := recover()
+			if rec == nil {
+				t.Fatal("SetLimit should have panicked, but didn't")
+			}
+
+			er, ok := rec.(error)
+			if !ok {
+				t.Fatalf("SetLimit should have panicked with an error, %T received", rec)
+			}
+
+			err = er
+		}()
+		g.SetLimit(5)
+	}()
+
+	<-ch
+
+	glimErr := &errgroup.ErrgroupLimitError{}
+	if !errors.As(err, &glimErr) {
+		t.Fatalf("panicked error should have been ErrgroupLimitError. %T received", err)
+	}
+
+	if glimErr.Size != 1 {
+		t.Fatalf("ErrgroupLimitError should have had a size of 1. Got %d", glimErr.Size)
+	}
+	expErrMsg := "errgroup: modify limit while 1 goroutines in the group are still active"
+	if errStr := glimErr.Error(); errStr != expErrMsg {
+		t.Fatalf("error message should have been %s. Got %s", expErrMsg, errStr)
+	}
+}
+
 func TestCancelCause(t *testing.T) {
 	errDoom := errors.New("group_test: doomed")