Ensure per-try timeout is properly cancelled

sdk/azcore/log.go

@@ -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
 	}

sdk/azcore/policy_retry.go

@@ -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))
 		}
@@ -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()
+}

sdk/azcore/request.go

@@ -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 {