github.com/toplink-cn/moby@v0.0.0-20240305205811-460b4aebdf81/registry/resumable/resumablerequestreader.go (about)

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