You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
 
 
 

102 lines
2.4 KiB

  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_test
  5. import (
  6. "context"
  7. "crypto/md5"
  8. "fmt"
  9. "io/ioutil"
  10. "log"
  11. "os"
  12. "path/filepath"
  13. "golang.org/x/sync/errgroup"
  14. )
  15. // Pipeline demonstrates the use of a Group to implement a multi-stage
  16. // pipeline: a version of the MD5All function with bounded parallelism from
  17. // https://blog.golang.org/pipelines.
  18. func ExampleGroup_pipeline() {
  19. m, err := MD5All(context.Background(), ".")
  20. if err != nil {
  21. log.Fatal(err)
  22. }
  23. for k, sum := range m {
  24. fmt.Printf("%s:\t%x\n", k, sum)
  25. }
  26. }
  27. type result struct {
  28. path string
  29. sum [md5.Size]byte
  30. }
  31. // MD5All reads all the files in the file tree rooted at root and returns a map
  32. // from file path to the MD5 sum of the file's contents. If the directory walk
  33. // fails or any read operation fails, MD5All returns an error.
  34. func MD5All(ctx context.Context, root string) (map[string][md5.Size]byte, error) {
  35. // ctx is canceled when g.Wait() returns. When this version of MD5All returns
  36. // - even in case of error! - we know that all of the goroutines have finished
  37. // and the memory they were using can be garbage-collected.
  38. g, ctx := errgroup.WithContext(ctx)
  39. paths := make(chan string)
  40. g.Go(func() error {
  41. defer close(paths)
  42. return filepath.Walk(root, func(path string, info os.FileInfo, err error) error {
  43. if err != nil {
  44. return err
  45. }
  46. if !info.Mode().IsRegular() {
  47. return nil
  48. }
  49. select {
  50. case paths <- path:
  51. case <-ctx.Done():
  52. return ctx.Err()
  53. }
  54. return nil
  55. })
  56. })
  57. // Start a fixed number of goroutines to read and digest files.
  58. c := make(chan result)
  59. const numDigesters = 20
  60. for i := 0; i < numDigesters; i++ {
  61. g.Go(func() error {
  62. for path := range paths {
  63. data, err := ioutil.ReadFile(path)
  64. if err != nil {
  65. return err
  66. }
  67. select {
  68. case c <- result{path, md5.Sum(data)}:
  69. case <-ctx.Done():
  70. return ctx.Err()
  71. }
  72. }
  73. return nil
  74. })
  75. }
  76. go func() {
  77. g.Wait()
  78. close(c)
  79. }()
  80. m := make(map[string][md5.Size]byte)
  81. for r := range c {
  82. m[r.path] = r.sum
  83. }
  84. // Check whether any of the goroutines failed. Since g is accumulating the
  85. // errors, we don't need to send them (or check for them) in the individual
  86. // results sent on the channel.
  87. if err := g.Wait(); err != nil {
  88. return nil, err
  89. }
  90. return m, nil
  91. }