forked from huin/goupnp
/
multiclient.go
132 lines (114 loc) · 2.86 KB
/
multiclient.go
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
127
128
129
130
131
132
package httpu
import (
"net/http"
"time"
"golang.org/x/sync/errgroup"
)
// MultiClient dispatches requests out to all the delegated clients.
type MultiClient struct {
// The HTTPU clients to delegate to.
delegates []ClientInterface
}
var _ ClientInterface = &MultiClient{}
// NewMultiClient creates a new MultiClient that delegates to all the given
// clients.
func NewMultiClient(delegates []ClientInterface) *MultiClient {
return &MultiClient{
delegates: delegates,
}
}
// Do implements ClientInterface.Do.
func (mc *MultiClient) Do(
req *http.Request,
timeout time.Duration,
numSends int,
) ([]*http.Response, error) {
tasks := &errgroup.Group{}
results := make(chan []*http.Response)
tasks.Go(func() error {
defer close(results)
return mc.sendRequests(results, req, timeout, numSends)
})
var responses []*http.Response
tasks.Go(func() error {
for rs := range results {
responses = append(responses, rs...)
}
return nil
})
return responses, tasks.Wait()
}
func (mc *MultiClient) sendRequests(
results chan<- []*http.Response,
req *http.Request,
timeout time.Duration,
numSends int,
) error {
tasks := &errgroup.Group{}
for _, d := range mc.delegates {
d := d // copy for closure
tasks.Go(func() error {
responses, err := d.Do(req, timeout, numSends)
if err != nil {
return err
}
results <- responses
return nil
})
}
return tasks.Wait()
}
// MultiClientCtx dispatches requests out to all the delegated clients.
type MultiClientCtx struct {
// The HTTPU clients to delegate to.
delegates []ClientInterfaceCtx
}
var _ ClientInterfaceCtx = &MultiClientCtx{}
// NewMultiClient creates a new MultiClient that delegates to all the given
// clients.
func NewMultiClientCtx(delegates []ClientInterfaceCtx) *MultiClientCtx {
return &MultiClientCtx{
delegates: delegates,
}
}
// DoWithContext implements ClientInterfaceCtx.DoWithContext.
func (mc *MultiClientCtx) DoWithContext(
req *http.Request,
numSends int,
) ([]*http.Response, error) {
tasks, ctx := errgroup.WithContext(req.Context())
req = req.WithContext(ctx) // so we cancel if the errgroup errors
results := make(chan []*http.Response)
// For each client, send the request to it and collect results.
tasks.Go(func() error {
defer close(results)
return mc.sendRequestsCtx(results, req, numSends)
})
var responses []*http.Response
tasks.Go(func() error {
for rs := range results {
responses = append(responses, rs...)
}
return nil
})
return responses, tasks.Wait()
}
func (mc *MultiClientCtx) sendRequestsCtx(
results chan<- []*http.Response,
req *http.Request,
numSends int,
) error {
tasks := &errgroup.Group{}
for _, d := range mc.delegates {
d := d // copy for closure
tasks.Go(func() error {
responses, err := d.DoWithContext(req, numSends)
if err != nil {
return err
}
results <- responses
return nil
})
}
return tasks.Wait()
}