azure-storage-blob-go/azblob/zc_retry_reader.go

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
}