~evanj/errgroupcount

ref: 45e72d8ba8fcca5ae47d06e98648bb2047cfadc6 errgroupcount/errgroupcount.go -rw-r--r-- 3.5 KiB
45e72d8bEvan M Jones Feat(*): Project init. Breaking into its own module, outside 1 year, 8 months ago
                                                                                
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
// Why? Isn't errgroup.Group{} enough?
//
// No.
//
// It has happened for too many times now where I have to start up a list of
// goroutines, have to keep track of their errors, but don't care about ALL
// their errors. So I've made this. A canonical example is something like the
// following (this is taken from my own personal blog). The requirements are:
// (1) Given a "slug" value we check to see if there's a piece of content in the
// database of with a "slug" of that value.
// (2) The piece of content can either have a type of "page" or "post"
// (3) We only want to return an http.StatusNotFound page if we cannot find a
// "page" or "post" with the slug value. E.G. We only care if we hit two errors
// (not found) instead of just one (as only the one case is supported by the
// normal errgroup.
//
// Aside: for this personal case about "slug" clashing is not cared about.
//
// Example:
//
// func (s SearchEndpoint) ServeHTTP(w http.ResponseWriter, r *http.Request) {
// 	slug := s.Param(r, "slug")
// 	eg, _ := errgroupcount.WithContext(r.Context())
// 	once := sync.Once{}
//
// 	eg.Go(func() error {
// 		item, err := s.content.Search("page", slug)
// 		if err != nil {
// 			return err
// 		}
// 		once.Do(func() {
// 			s.HTML(w, r, http.StatusOK, "item.html", map[string]interface{}{"Item": item})
// 		})
// 		return nil
// 	})
//
// 	eg.Go(func() error {
// 		item, err := s.content.Search("post", slug)
// 		if err != nil {
// 			return err
// 		}
// 		once.Do(func() {
// 			s.HTML(w, r, http.StatusOK, "item.html", map[string]interface{}{"Item": item})
// 		})
// 		return nil
// 	})
//
// 	if err := eg.WaitCount(2); err != nil {
// 		s.ErrorString(w, r, http.StatusNotFound, "failed to find content")
// 	}
// }
//
// Enjoy.
//
package errgroupcount

import (
	"context"
	"sync"
)

// Package errgroupcount provides synchronization, error propagation, and Context
// cancelation for groups of goroutines working on subtasks of a common task.
//
// TYPES
//
// A Group is a collection of goroutines working on subtasks that are part of
// the same overall task.
//
// A zero Group is valid and does not cancel on error.
type Group struct {
	cancel func()

	wg sync.WaitGroup

	errMutex sync.Mutex
	err      []error
}

// WithContext returns a new Group and an associated Context derived from ctx.
//
// The derived Context is canceled the first time a function passed to Go
// returns a non-nil error or the first time Wait returns, whichever occurs
// first.
func WithContext(ctx context.Context) (*Group, context.Context) {
	ctx, cancel := context.WithCancel(ctx)
	return &Group{cancel: cancel}, ctx
}

// Go calls the given function in a new goroutine.
//
// The first call to return a non-nil error cancels the group; its error will be
// returned by Wait.
func (g *Group) Go(f func() error) {
	g.wg.Add(1)

	go func() {
		defer g.wg.Done()

		if err := f(); err != nil {
			g.errMutex.Lock()
			g.err = append(g.err, err)
			if g.cancel != nil {
				g.cancel()
			}
			g.errMutex.Unlock()
		}
	}()
}

// WaitCount blocks until all function calls from the Go method have returned,
// then will pop shift and return the first error received if and only if the
// total error count is greater than the int provided to WaitCount. If the total
// error count is less than the value supplied to WaitCount nil is returned.
// `top int` must be greated than zero, otherise nil is always returned.
func (g *Group) WaitCount(top int) (err error) {
	g.wg.Wait()
	if g.cancel != nil {
		g.cancel()
	}
	if top > 0 && len(g.err) >= top {
		err, g.err = g.err[0], g.err[1:]
		return err
	}
	return nil
}