Skip to content

Commit

Permalink
Ensure per-try timeout is properly cancelled (#7541)
Browse files Browse the repository at this point in the history
Explicitly call cancel the per-try timeout when the response body has
been read/closed by the body download policy.
When the response body is returned to the caller for reading/closing,
wrap it in a responseBodyReader that will cancel the timeout when the
body is closed.
Logger.Should() will return false if no listener is set.
  • Loading branch information
jhendrixMSFT authored Feb 25, 2020
1 parent ee7c545 commit cbb75bf
Show file tree
Hide file tree
Showing 3 changed files with 40 additions and 3 deletions.
9 changes: 8 additions & 1 deletion sdk/azcore/log.go
Original file line number Diff line number Diff line change
Expand Up @@ -51,8 +51,15 @@ func (l *Logger) SetListener(lst Listener) {
}

// Should returns true if the specified log classification should be written to the log.
// TODO: explain why you would want to call this
// By default all log classifications will be logged. Call SetClassification() to limit
// the log classifications for logging.
// If no listener has been set this will return false.
// Calling this method is useful when the message to log is computationally expensive
// and you want to avoid the overhead if its log classification is not enabled.
func (l *Logger) Should(cls LogClassification) bool {
if l.lst == nil {
return false
}
if l.cls == nil || len(l.cls) == 0 {
return true
}
Expand Down
27 changes: 25 additions & 2 deletions sdk/azcore/policy_retry.go
Original file line number Diff line number Diff line change
Expand Up @@ -143,10 +143,18 @@ func (p *retryPolicy) Do(ctx context.Context, req *Request) (resp *Response, err
return
}

// Set the time for this particular retry operation and then Do the operation.
// Set the per-try time for this particular retry operation and then Do the operation.
tryCtx, tryCancel := context.WithTimeout(ctx, options.TryTimeout)
resp, err = req.Next(tryCtx) // Make the request
tryCancel()
if req.bodyDownloadEnabled() {
// if auto-downloading of the response body is enabled then
// it's been read and closed by this point so cancel the timeout
tryCancel()
} else {
// wrap the response body in a responseBodyReader.
// closing the responseBodyReader will cancel the timeout.
resp.Body = &responseBodyReader{rb: resp.Body, cancelFunc: tryCancel}
}
if shouldLog {
Log().Write(LogRetryPolicy, fmt.Sprintf("Err=%v, response=%v\n", err, resp))
}
Expand Down Expand Up @@ -216,3 +224,18 @@ func (b *retryableRequestBody) realClose() error {
}
return nil
}

// used when returning the response body to the caller for reading/closing
type responseBodyReader struct {
rb io.ReadCloser
cancelFunc context.CancelFunc
}

func (r *responseBodyReader) Read(p []byte) (int, error) {
return r.rb.Read(p)
}

func (r *responseBodyReader) Close() error {
r.cancelFunc()
return r.rb.Close()
}
7 changes: 7 additions & 0 deletions sdk/azcore/request.go
Original file line number Diff line number Diff line change
Expand Up @@ -144,6 +144,13 @@ func (req *Request) SkipBodyDownload() {
req.SetOperationValue(bodyDownloadPolicyOpValues{skip: true})
}

// returns true if auto-body download policy is enabled
func (req *Request) bodyDownloadEnabled() bool {
var opValues bodyDownloadPolicyOpValues
req.OperationValue(&opValues)
return !opValues.skip
}

// RewindBody seeks the request's Body stream back to the beginning so it can be resent when retrying an operation.
func (req *Request) RewindBody() error {
if req.Body != nil {
Expand Down

0 comments on commit cbb75bf

Please sign in to comment.