multiclient.go 1.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172
  1. package httpu
  2. import (
  3. "net/http"
  4. "golang.org/x/sync/errgroup"
  5. )
  6. // MultiClient dispatches requests out to all the delegated clients.
  7. type MultiClient struct {
  8. // The HTTPU clients to delegate to.
  9. delegates []ClientInterface
  10. }
  11. var _ ClientInterface = &MultiClient{}
  12. // NewMultiClient creates a new MultiClient that delegates to all the given
  13. // clients.
  14. func NewMultiClient(delegates []ClientInterface) *MultiClient {
  15. return &MultiClient{delegates: delegates}
  16. }
  17. func (mc *MultiClient) Close() error {
  18. for _, d := range mc.delegates {
  19. d.Close()
  20. }
  21. return nil
  22. }
  23. // Do implements ClientInterface.Do.
  24. func (mc *MultiClient) Do(
  25. req *http.Request,
  26. numSends int,
  27. ) ([]*http.Response, error) {
  28. tasks := &errgroup.Group{}
  29. results := make(chan []*http.Response)
  30. tasks.Go(func() error {
  31. defer close(results)
  32. return mc.sendRequests(results, req, numSends)
  33. })
  34. var responses []*http.Response
  35. tasks.Go(func() error {
  36. for rs := range results {
  37. responses = append(responses, rs...)
  38. }
  39. return nil
  40. })
  41. return responses, tasks.Wait()
  42. }
  43. func (mc *MultiClient) sendRequests(
  44. results chan<- []*http.Response,
  45. req *http.Request,
  46. numSends int,
  47. ) error {
  48. tasks := &errgroup.Group{}
  49. for _, d := range mc.delegates {
  50. d := d // copy for closure
  51. tasks.Go(func() error {
  52. responses, err := d.Do(req, numSends)
  53. if err != nil {
  54. return err
  55. }
  56. results <- responses
  57. return nil
  58. })
  59. }
  60. return tasks.Wait()
  61. }