errgroup.go 3.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151
  1. // Copyright 2016 The Go Authors. All rights reserved.
  2. // Use of this source code is governed by a BSD-style
  3. // license that can be found in the LICENSE file.
  4. // Package errgroup provides synchronization, error propagation, and Context
  5. // cancelation for groups of goroutines working on subtasks of a common task.
  6. //
  7. // [errgroup.Group] is related to [sync.WaitGroup] but adds handling of tasks
  8. // returning errors.
  9. package errgroup
  10. import (
  11. "context"
  12. "fmt"
  13. "sync"
  14. )
  15. type token struct{}
  16. // A Group is a collection of goroutines working on subtasks that are part of
  17. // the same overall task. A Group should not be reused for different tasks.
  18. //
  19. // A zero Group is valid, has no limit on the number of active goroutines,
  20. // and does not cancel on error.
  21. type Group struct {
  22. cancel func(error)
  23. wg sync.WaitGroup
  24. sem chan token
  25. errOnce sync.Once
  26. err error
  27. }
  28. func (g *Group) done() {
  29. if g.sem != nil {
  30. <-g.sem
  31. }
  32. g.wg.Done()
  33. }
  34. // WithContext returns a new Group and an associated Context derived from ctx.
  35. //
  36. // The derived Context is canceled the first time a function passed to Go
  37. // returns a non-nil error or the first time Wait returns, whichever occurs
  38. // first.
  39. func WithContext(ctx context.Context) (*Group, context.Context) {
  40. ctx, cancel := context.WithCancelCause(ctx)
  41. return &Group{cancel: cancel}, ctx
  42. }
  43. // Wait blocks until all function calls from the Go method have returned, then
  44. // returns the first non-nil error (if any) from them.
  45. func (g *Group) Wait() error {
  46. g.wg.Wait()
  47. if g.cancel != nil {
  48. g.cancel(g.err)
  49. }
  50. return g.err
  51. }
  52. // Go calls the given function in a new goroutine.
  53. //
  54. // The first call to Go must happen before a Wait.
  55. // It blocks until the new goroutine can be added without the number of
  56. // goroutines in the group exceeding the configured limit.
  57. //
  58. // The first goroutine in the group that returns a non-nil error will
  59. // cancel the associated Context, if any. The error will be returned
  60. // by Wait.
  61. func (g *Group) Go(f func() error) {
  62. if g.sem != nil {
  63. g.sem <- token{}
  64. }
  65. g.wg.Add(1)
  66. go func() {
  67. defer g.done()
  68. // It is tempting to propagate panics from f()
  69. // up to the goroutine that calls Wait, but
  70. // it creates more problems than it solves:
  71. // - it delays panics arbitrarily,
  72. // making bugs harder to detect;
  73. // - it turns f's panic stack into a mere value,
  74. // hiding it from crash-monitoring tools;
  75. // - it risks deadlocks that hide the panic entirely,
  76. // if f's panic leaves the program in a state
  77. // that prevents the Wait call from being reached.
  78. // See #53757, #74275, #74304, #74306.
  79. if err := f(); err != nil {
  80. g.errOnce.Do(func() {
  81. g.err = err
  82. if g.cancel != nil {
  83. g.cancel(g.err)
  84. }
  85. })
  86. }
  87. }()
  88. }
  89. // TryGo calls the given function in a new goroutine only if the number of
  90. // active goroutines in the group is currently below the configured limit.
  91. //
  92. // The return value reports whether the goroutine was started.
  93. func (g *Group) TryGo(f func() error) bool {
  94. if g.sem != nil {
  95. select {
  96. case g.sem <- token{}:
  97. // Note: this allows barging iff channels in general allow barging.
  98. default:
  99. return false
  100. }
  101. }
  102. g.wg.Add(1)
  103. go func() {
  104. defer g.done()
  105. if err := f(); err != nil {
  106. g.errOnce.Do(func() {
  107. g.err = err
  108. if g.cancel != nil {
  109. g.cancel(g.err)
  110. }
  111. })
  112. }
  113. }()
  114. return true
  115. }
  116. // SetLimit limits the number of active goroutines in this group to at most n.
  117. // A negative value indicates no limit.
  118. // A limit of zero will prevent any new goroutines from being added.
  119. //
  120. // Any subsequent call to the Go method will block until it can add an active
  121. // goroutine without exceeding the configured limit.
  122. //
  123. // The limit must not be modified while any goroutines in the group are active.
  124. func (g *Group) SetLimit(n int) {
  125. if n < 0 {
  126. g.sem = nil
  127. return
  128. }
  129. if len(g.sem) != 0 {
  130. panic(fmt.Errorf("errgroup: modify limit while %v goroutines in the group are still active", len(g.sem)))
  131. }
  132. g.sem = make(chan token, n)
  133. }