github.com/wlan0/docker@v1.5.0/pkg/httputils/resumablerequestreader.go (about)

     1  package httputils
     2  
     3  import (
     4  	"fmt"
     5  	"io"
     6  	"net/http"
     7  	"time"
     8  
     9  	log "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  func ResumableRequestReaderWithInitialResponse(c *http.Client, r *http.Request, maxfail uint32, totalsize int64, initialResponse *http.Response) io.ReadCloser {
    30  	return &resumableRequestReader{client: c, request: r, maxFailures: maxfail, totalSize: totalsize, currentResponse: initialResponse}
    31  }
    32  
    33  func (r *resumableRequestReader) Read(p []byte) (n int, err error) {
    34  	if r.client == nil || r.request == nil {
    35  		return 0, fmt.Errorf("client and request can't be nil\n")
    36  	}
    37  	isFreshRequest := false
    38  	if r.lastRange != 0 && r.currentResponse == nil {
    39  		readRange := fmt.Sprintf("bytes=%d-%d", r.lastRange, r.totalSize)
    40  		r.request.Header.Set("Range", readRange)
    41  		time.Sleep(5 * time.Second)
    42  	}
    43  	if r.currentResponse == nil {
    44  		r.currentResponse, err = r.client.Do(r.request)
    45  		isFreshRequest = true
    46  	}
    47  	if err != nil && r.failures+1 != r.maxFailures {
    48  		r.cleanUpResponse()
    49  		r.failures++
    50  		time.Sleep(5 * time.Duration(r.failures) * time.Second)
    51  		return 0, nil
    52  	} else if err != nil {
    53  		r.cleanUpResponse()
    54  		return 0, err
    55  	}
    56  	if r.currentResponse.StatusCode == 416 && r.lastRange == r.totalSize && r.currentResponse.ContentLength == 0 {
    57  		r.cleanUpResponse()
    58  		return 0, io.EOF
    59  	} else if r.currentResponse.StatusCode != 206 && r.lastRange != 0 && isFreshRequest {
    60  		r.cleanUpResponse()
    61  		return 0, fmt.Errorf("the server doesn't support byte ranges")
    62  	}
    63  	if r.totalSize == 0 {
    64  		r.totalSize = r.currentResponse.ContentLength
    65  	} else if r.totalSize <= 0 {
    66  		r.cleanUpResponse()
    67  		return 0, fmt.Errorf("failed to auto detect content length")
    68  	}
    69  	n, err = r.currentResponse.Body.Read(p)
    70  	r.lastRange += int64(n)
    71  	if err != nil {
    72  		r.cleanUpResponse()
    73  	}
    74  	if err != nil && err != io.EOF {
    75  		log.Infof("encountered error during pull and clearing it before resume: %s", err)
    76  		err = nil
    77  	}
    78  	return n, err
    79  }
    80  
    81  func (r *resumableRequestReader) Close() error {
    82  	r.cleanUpResponse()
    83  	r.client = nil
    84  	r.request = nil
    85  	return nil
    86  }
    87  
    88  func (r *resumableRequestReader) cleanUpResponse() {
    89  	if r.currentResponse != nil {
    90  		r.currentResponse.Body.Close()
    91  		r.currentResponse = nil
    92  	}
    93  }