github.com/olljanat/moby@v1.13.1/pkg/httputils/resumablerequestreader.go (about)

     1  package httputils
     2  
     3  import (
     4  	"fmt"
     5  	"io"
     6  	"net/http"
     7  	"time"
     8  
     9  	"github.com/Sirupsen/logrus"
    10  )
    11  
    12  type resumableRequestReader struct {
    13  	client          *http.Client
    14  	request         *http.Request
    15  	lastRange       int64
    16  	totalSize       int64
    17  	currentResponse *http.Response
    18  	failures        uint32
    19  	maxFailures     uint32
    20  }
    21  
    22  // ResumableRequestReader makes it possible to resume reading a request's body transparently
    23  // maxfail is the number of times we retry to make requests again (not resumes)
    24  // totalsize is the total length of the body; auto detect if not provided
    25  func ResumableRequestReader(c *http.Client, r *http.Request, maxfail uint32, totalsize int64) io.ReadCloser {
    26  	return &resumableRequestReader{client: c, request: r, maxFailures: maxfail, totalSize: totalsize}
    27  }
    28  
    29  // ResumableRequestReaderWithInitialResponse makes it possible to resume
    30  // reading the body of an already initiated request.
    31  func ResumableRequestReaderWithInitialResponse(c *http.Client, r *http.Request, maxfail uint32, totalsize int64, initialResponse *http.Response) io.ReadCloser {
    32  	return &resumableRequestReader{client: c, request: r, maxFailures: maxfail, totalSize: totalsize, currentResponse: initialResponse}
    33  }
    34  
    35  func (r *resumableRequestReader) Read(p []byte) (n int, err error) {
    36  	if r.client == nil || r.request == nil {
    37  		return 0, fmt.Errorf("client and request can't be nil\n")
    38  	}
    39  	isFreshRequest := false
    40  	if r.lastRange != 0 && r.currentResponse == nil {
    41  		readRange := fmt.Sprintf("bytes=%d-%d", r.lastRange, r.totalSize)
    42  		r.request.Header.Set("Range", readRange)
    43  		time.Sleep(5 * time.Second)
    44  	}
    45  	if r.currentResponse == nil {
    46  		r.currentResponse, err = r.client.Do(r.request)
    47  		isFreshRequest = true
    48  	}
    49  	if err != nil && r.failures+1 != r.maxFailures {
    50  		r.cleanUpResponse()
    51  		r.failures++
    52  		time.Sleep(5 * time.Duration(r.failures) * time.Second)
    53  		return 0, nil
    54  	} else if err != nil {
    55  		r.cleanUpResponse()
    56  		return 0, err
    57  	}
    58  	if r.currentResponse.StatusCode == 416 && r.lastRange == r.totalSize && r.currentResponse.ContentLength == 0 {
    59  		r.cleanUpResponse()
    60  		return 0, io.EOF
    61  	} else if r.currentResponse.StatusCode != 206 && r.lastRange != 0 && isFreshRequest {
    62  		r.cleanUpResponse()
    63  		return 0, fmt.Errorf("the server doesn't support byte ranges")
    64  	}
    65  	if r.totalSize == 0 {
    66  		r.totalSize = r.currentResponse.ContentLength
    67  	} else if r.totalSize <= 0 {
    68  		r.cleanUpResponse()
    69  		return 0, fmt.Errorf("failed to auto detect content length")
    70  	}
    71  	n, err = r.currentResponse.Body.Read(p)
    72  	r.lastRange += int64(n)
    73  	if err != nil {
    74  		r.cleanUpResponse()
    75  	}
    76  	if err != nil && err != io.EOF {
    77  		logrus.Infof("encountered error during pull and clearing it before resume: %s", err)
    78  		err = nil
    79  	}
    80  	return n, err
    81  }
    82  
    83  func (r *resumableRequestReader) Close() error {
    84  	r.cleanUpResponse()
    85  	r.client = nil
    86  	r.request = nil
    87  	return nil
    88  }
    89  
    90  func (r *resumableRequestReader) cleanUpResponse() {
    91  	if r.currentResponse != nil {
    92  		r.currentResponse.Body.Close()
    93  		r.currentResponse = nil
    94  	}
    95  }