errgroup/errgroup_test.go

// Copyright 2016 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package errgroup_test

import (
	"context"
	"errors"
	"fmt"
	"net/http"
	"os"
	"sync/atomic"
	"testing"
	"time"

	"codeberg.org/danjones000/errgroup"
)

var (
	Web   = fakeSearch("web")
	Image = fakeSearch("image")
	Video = fakeSearch("video")
)

type Result string
type Search func(ctx context.Context, query string) (Result, error)

func fakeSearch(kind string) Search {
	return func(_ context.Context, query string) (Result, error) {
		return Result(fmt.Sprintf("%s result for %q", kind, query)), nil
	}
}

// JustErrors illustrates the use of a Group in place of a sync.WaitGroup to
// simplify goroutine counting and error handling. This example is derived from
// the sync.WaitGroup example at https://golang.org/pkg/sync/#example_WaitGroup.
func ExampleGroup_justErrors() {
	g := new(errgroup.Group)
	var urls = []string{
		"http://www.golang.org/",
		"http://www.google.com/",
		"http://www.somestupidname.com/",
	}
	for _, url := range urls {
		g.Go(func() error {
			// Fetch the URL.
			resp, err := http.Get(url)
			if err == nil {
				resp.Body.Close()
			}
			return err
		})
	}
	// Wait for all HTTP fetches to complete.
	if err := g.Wait(); err == nil {
		fmt.Println("Successfully fetched all URLs.")
	}
}

// Parallel illustrates the use of a Group for synchronizing a simple parallel
// task: the "Google Search 2.0" function from
// https://talks.golang.org/2012/concurrency.slide#46, augmented with a Context
// and error-handling.
func ExampleGroup_parallel() {
	Google := func(ctx context.Context, query string) ([]Result, error) {
		g, ctx := errgroup.WithContext(ctx)

		searches := []Search{Web, Image, Video}
		results := make([]Result, len(searches))
		for i, search := range searches {
			g.Go(func() error {
				result, err := search(ctx, query)
				if err == nil {
					results[i] = result
				}
				return err
			})
		}
		if err := g.Wait(); err != nil {
			return nil, err
		}
		return results, nil
	}

	results, err := Google(context.Background(), "golang")
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		return
	}
	for _, result := range results {
		fmt.Println(result)
	}

	// Output:
	// web result for "golang"
	// image result for "golang"
	// video result for "golang"
}

// FirstError demonstrates that g.Go becomes a no-op if a previous g.Go
// has returned an error.
func ExampleGroup_firstError() {
	err1 := errors.New("errgroup_test: 1")
	err2 := errors.New("errgroup_test: 2")

	g := new(errgroup.Group)

	ch := make(chan struct{})

	g.Go(func() error {
		fmt.Printf("Returning %s\n", err1)
		ch <- struct{}{}
		return err1
	})

	<-ch

	g.Go(func() error {
		// This should never run
		fmt.Printf("Returning %s\n", err2)
		return err2
	})

	err := g.Wait()
	fmt.Printf("Got %s\n", err)

	// Output:
	// Returning errgroup_test: 1
	// Got errgroup_test: 1
}

func TestZeroGroup(t *testing.T) {
	err1 := errors.New("errgroup_test: 1")
	err2 := errors.New("errgroup_test: 2")

	cases := []struct {
		errs []error
	}{
		{errs: []error{}},
		{errs: []error{nil}},
		{errs: []error{err1}},
		{errs: []error{err1, nil}},
		{errs: []error{err1, nil, err2}},
	}

	for _, tc := range cases {
		g := new(errgroup.Group)

		var firstErr error
		for i, err := range tc.errs {
			g.Go(func() error { return err })

			if firstErr == nil && err != nil {
				firstErr = err
			}

			if gErr := g.Wait(); gErr != firstErr {
				t.Errorf("after %T.Go(func() error { return err }) for err in %v\n"+
					"g.Wait() = %v; want %v",
					g, tc.errs[:i+1], err, firstErr)
			}
		}
	}
}

func TestWithContext(t *testing.T) {
	errDoom := errors.New("group_test: doomed")

	cases := []struct {
		errs []error
		want error
	}{
		{want: nil},
		{errs: []error{nil}, want: nil},
		{errs: []error{errDoom}, want: errDoom},
		{errs: []error{errDoom, nil}, want: errDoom},
	}

	for _, tc := range cases {
		g, ctx := errgroup.WithContext(context.Background())

		for _, err := range tc.errs {
			g.Go(func() error { return err })
		}

		if err := g.Wait(); err != tc.want {
			t.Errorf("after %T.Go(func() error { return err }) for err in %v\n"+
				"g.Wait() = %v; want %v",
				g, tc.errs, err, tc.want)
		}

		canceled := false
		select {
		case <-ctx.Done():
			canceled = true
		default:
		}
		if !canceled {
			t.Errorf("after %T.Go(func() error { return err }) for err in %v\n"+
				"ctx.Done() was not closed",
				g, tc.errs)
		}
	}
}

func TestTryGo(t *testing.T) {
	g := &errgroup.Group{}
	n := 42
	g.SetLimit(42)
	ch := make(chan struct{})
	fn := func() error {
		ch <- struct{}{}
		return nil
	}
	for i := 0; i < n; i++ {
		if !g.TryGo(fn) {
			t.Fatalf("TryGo should succeed but got fail at %d-th call.", i)
		}
	}
	if g.TryGo(fn) {
		t.Fatalf("TryGo is expected to fail but succeeded.")
	}
	go func() {
		for i := 0; i < n; i++ {
			<-ch
		}
	}()
	g.Wait()

	if !g.TryGo(fn) {
		t.Fatalf("TryGo should success but got fail after all goroutines.")
	}
	go func() { <-ch }()
	g.Wait()

	// Switch limit.
	g.SetLimit(1)
	if !g.TryGo(fn) {
		t.Fatalf("TryGo should success but got failed.")
	}
	if g.TryGo(fn) {
		t.Fatalf("TryGo should fail but succeeded.")
	}
	go func() { <-ch }()
	g.Wait()

	// Block all calls.
	g.SetLimit(0)
	for i := 0; i < 1<<10; i++ {
		if g.TryGo(fn) {
			t.Fatalf("TryGo should fail but got succeded.")
		}
	}
	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

	g := &errgroup.Group{}
	g.SetLimit(limit)
	var active int32
	for i := 0; i <= 1<<10; i++ {
		g.Go(func() error {
			n := atomic.AddInt32(&active, 1)
			if n > limit {
				return fmt.Errorf("saw %d active goroutines; want ≤ %d", n, limit)
			}
			time.Sleep(1 * time.Microsecond) // Give other goroutines a chance to increment active.
			atomic.AddInt32(&active, -1)
			return nil
		})
	}
	if err := g.Wait(); err != nil {
		t.Fatal(err)
	}
}

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")

	cases := []struct {
		errs []error
		want error
	}{
		{want: nil},
		{errs: []error{nil}, want: nil},
		{errs: []error{errDoom}, want: errDoom},
		{errs: []error{errDoom, nil}, want: errDoom},
		{errs: []error{nil, errDoom}, want: errDoom},
	}

	for _, tc := range cases {
		g, ctx := errgroup.WithContext(context.Background())

		for _, err := range tc.errs {
			g.TryGo(func() error { return err })
		}

		if err := g.Wait(); err != tc.want {
			t.Errorf("after %T.TryGo(func() error { return err }) for err in %v\n"+
				"g.Wait() = %v; want %v",
				g, tc.errs, err, tc.want)
		}

		if tc.want == nil {
			tc.want = context.Canceled
		}

		if err := context.Cause(ctx); err != tc.want {
			t.Errorf("after %T.TryGo(func() error { return err }) for err in %v\n"+
				"context.Cause(ctx) = %v; tc.want %v",
				g, tc.errs, err, tc.want)
		}
	}
}

func BenchmarkGo(b *testing.B) {
	fn := func() {}
	g := &errgroup.Group{}
	b.ResetTimer()
	b.ReportAllocs()
	for i := 0; i < b.N; i++ {
		g.Go(func() error { fn(); return nil })
	}
	g.Wait()
}