/
zc_retry_reader.go
178 lines (154 loc) · 7.15 KB
/
zc_retry_reader.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
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
package azblob
import (
"context"
"io"
"net"
"net/http"
"strings"
"sync"
)
const CountToEnd = 0
// HTTPGetter is a function type that refers to a method that performs an HTTP GET operation.
type HTTPGetter func(ctx context.Context, i HTTPGetterInfo) (*http.Response, error)
// HTTPGetterInfo is passed to an HTTPGetter function passing it parameters
// that should be used to make an HTTP GET request.
type HTTPGetterInfo struct {
// Offset specifies the start offset that should be used when
// creating the HTTP GET request's Range header
Offset int64
// Count specifies the count of bytes that should be used to calculate
// the end offset when creating the HTTP GET request's Range header
Count int64
// ETag specifies the resource's etag that should be used when creating
// the HTTP GET request's If-Match header
ETag ETag
}
// FailedReadNotifier is a function type that represents the notification function called when a read fails
type FailedReadNotifier func(failureCount int, lastError error, offset int64, count int64, willRetry bool)
// RetryReaderOptions contains properties which can help to decide when to do retry.
type RetryReaderOptions struct {
// MaxRetryRequests specifies the maximum number of HTTP GET requests that will be made
// while reading from a RetryReader. A value of zero means that no additional HTTP
// GET requests will be made.
MaxRetryRequests int
doInjectError bool
doInjectErrorRound int
// NotifyFailedRead is called, if non-nil, after any failure to read. Expected usage is diagnostic logging.
NotifyFailedRead FailedReadNotifier
// TreatEarlyCloseAsError can be set to true to prevent retries after "read on closed response body". By default,
// retryReader has the following special behaviour: closing the response body before it is all read is treated as a
// retryable error. This is to allow callers to force a retry by closing the body from another goroutine (e.g. if the =
// read is too slow, caller may want to force a retry in the hope that the retry will be quicker). If
// TreatEarlyCloseAsError is true, then retryReader's special behaviour is suppressed, and "read on closed body" is instead
// treated as a fatal (non-retryable) error.
// Note that setting TreatEarlyCloseAsError only guarantees that Closing will produce a fatal error if the Close happens
// from the same "thread" (goroutine) as Read. Concurrent Close calls from other goroutines may instead produce network errors
// which will be retried.
TreatEarlyCloseAsError bool
}
// retryReader implements io.ReaderCloser methods.
// retryReader tries to read from response, and if there is retriable network error
// returned during reading, it will retry according to retry reader option through executing
// user defined action with provided data to get a new response, and continue the overall reading process
// through reading from the new response.
type retryReader struct {
ctx context.Context
info HTTPGetterInfo
countWasBounded bool
o RetryReaderOptions
getter HTTPGetter
// we support Close-ing during Reads (from other goroutines), so we protect the shared state, which is response
responseMu *sync.Mutex
response *http.Response
}
// NewRetryReader creates a retry reader.
func NewRetryReader(ctx context.Context, initialResponse *http.Response,
info HTTPGetterInfo, o RetryReaderOptions, getter HTTPGetter) io.ReadCloser {
return &retryReader{
ctx: ctx,
getter: getter,
info: info,
countWasBounded: info.Count != CountToEnd,
response: initialResponse,
responseMu: &sync.Mutex{},
o: o}
}
func (s *retryReader) setResponse(r *http.Response) {
s.responseMu.Lock()
defer s.responseMu.Unlock()
s.response = r
}
func (s *retryReader) Read(p []byte) (n int, err error) {
for try := 0; ; try++ {
//fmt.Println(try) // Comment out for debugging.
if s.countWasBounded && s.info.Count == CountToEnd {
// User specified an original count and the remaining bytes are 0, return 0, EOF
return 0, io.EOF
}
s.responseMu.Lock()
resp := s.response
s.responseMu.Unlock()
if resp == nil { // We don't have a response stream to read from, try to get one.
newResponse, err := s.getter(s.ctx, s.info)
if err != nil {
return 0, err
}
// Successful GET; this is the network stream we'll read from.
s.setResponse(newResponse)
resp = newResponse
}
n, err := resp.Body.Read(p) // Read from the stream (this will return non-nil err if forceRetry is called, from another goroutine, while it is running)
// Injection mechanism for testing.
if s.o.doInjectError && try == s.o.doInjectErrorRound {
err = &net.DNSError{IsTemporary: true}
}
// We successfully read data or end EOF.
if err == nil || err == io.EOF {
s.info.Offset += int64(n) // Increments the start offset in case we need to make a new HTTP request in the future
if s.info.Count != CountToEnd {
s.info.Count -= int64(n) // Decrement the count in case we need to make a new HTTP request in the future
}
return n, err // Return the return to the caller
}
s.Close() // Error, close stream
s.setResponse(nil) // Our stream is no longer good
// Check the retry count and error code, and decide whether to retry.
retriesExhausted := try >= s.o.MaxRetryRequests
_, isNetError := err.(net.Error)
willRetry := (isNetError || s.wasRetryableEarlyClose(err)) && !retriesExhausted
// Notify, for logging purposes, of any failures
if s.o.NotifyFailedRead != nil {
failureCount := try + 1 // because try is zero-based
s.o.NotifyFailedRead(failureCount, err, s.info.Offset, s.info.Count, willRetry)
}
if willRetry {
continue
// Loop around and try to get and read from new stream.
}
return n, err // Not retryable, or retries exhausted, so just return
}
}
// By default, we allow early Closing, from another concurrent goroutine, to be used to force a retry
// Is this safe, to close early from another goroutine? Early close ultimately ends up calling
// net.Conn.Close, and that is documented as "Any blocked Read or Write operations will be unblocked and return errors"
// which is exactly the behaviour we want.
// NOTE: that if caller has forced an early Close from a separate goroutine (separate from the Read)
// then there are two different types of error that may happen - either the one one we check for here,
// or a net.Error (due to closure of connection). Which one happens depends on timing. We only need this routine
// to check for one, since the other is a net.Error, which our main Read retry loop is already handing.
func (s *retryReader) wasRetryableEarlyClose(err error) bool {
if s.o.TreatEarlyCloseAsError {
return false // user wants all early closes to be errors, and so not retryable
}
// unfortunately, http.errReadOnClosedResBody is private, so the best we can do here is to check for its text
return strings.HasSuffix(err.Error(), ReadOnClosedBodyMessage)
}
const ReadOnClosedBodyMessage = "read on closed response body"
func (s *retryReader) Close() error {
s.responseMu.Lock()
defer s.responseMu.Unlock()
if s.response != nil && s.response.Body != nil {
return s.response.Body.Close()
}
return nil
}