optimize(transport): optimize resumable transport initial request

Author: derekhjray

pkg/v1/remote/transport/resumable.go

@@ -29,43 +29,20 @@ type resumableTransport struct {
 }
 
 func (rt *resumableTransport) RoundTrip(in *http.Request) (*http.Response, error) {
-	if in.Method != http.MethodGet {
-		return rt.inner.RoundTrip(in)
-	}
-
-	req := in.Clone(in.Context())
-	req.Header.Set("Range", "bytes=0-")
-	resp, err := rt.inner.RoundTrip(req)
+	resp, err := rt.inner.RoundTrip(in)
 	if err != nil {
 		return resp, err
 	}
 
-	switch resp.StatusCode {
-	case http.StatusPartialContent:
-	case http.StatusRequestedRangeNotSatisfiable:
-		// fallback to previous behavior
-		resp.Body.Close()
-		return rt.inner.RoundTrip(in)
-	default:
+	if in.Method != http.MethodGet || resp.StatusCode != http.StatusOK || resp.ContentLength <= 0 {
 		return resp, nil
 	}
 
-	var contentLength int64
-	if _, _, contentLength, err = parseContentRange(resp.Header.Get("Content-Range")); err != nil || contentLength <= 0 {
-		// fallback to previous behavior
-		resp.Body.Close()
-		return rt.inner.RoundTrip(in)
-	}
-
-	// modify response status to 200, ensure caller error checking works
-	resp.StatusCode = http.StatusOK
-	resp.Status = "200 OK"
-	resp.ContentLength = contentLength
 	resp.Body = &resumableBody{
 		rc:          resp.Body,
 		inner:       rt.inner,
-		req:         req,
-		total:       contentLength,
+		req:         in,
+		total:       resp.ContentLength,
 		transferred: 0,
 	}
 

pkg/v1/remote/transport/resumable_test.go

@@ -31,42 +31,46 @@ func handleResumableLayer(data []byte, w http.ResponseWriter, r *http.Request, t
 		return
 	}
 
+	var (
+		contentLength, start, end int64
+		statusCode                = http.StatusOK
+		err                       error
+	)
+
+	contentLength = int64(len(data))
+	end = contentLength - 1
 	contentRange := r.Header.Get("Range")
-	if contentRange == "" {
-		w.WriteHeader(http.StatusBadRequest)
-		return
-	}
-
-	matches := rangeRe.FindStringSubmatch(contentRange)
-	if len(matches) != 3 {
-		w.WriteHeader(http.StatusBadRequest)
-		return
-	}
-
-	contentLength := int64(len(data))
-	start, err := strconv.ParseInt(matches[1], 10, 64)
-	if err != nil || start < 0 {
-		w.WriteHeader(http.StatusBadRequest)
-		return
-	}
-
-	if start >= int64(contentLength) {
-		w.WriteHeader(http.StatusRequestedRangeNotSatisfiable)
-		return
-	}
+	if contentRange != "" {
+		matches := rangeRe.FindStringSubmatch(contentRange)
+		if len(matches) != 3 {
+			w.WriteHeader(http.StatusBadRequest)
+			return
+		}
 
-	var end = int64(contentLength) - 1
-	if matches[2] != "" {
-		end, err = strconv.ParseInt(matches[2], 10, 64)
-		if err != nil || end < 0 {
+		if start, err = strconv.ParseInt(matches[1], 10, 64); err != nil || start < 0 {
 			w.WriteHeader(http.StatusBadRequest)
 			return
 		}
 
-		if end >= int64(contentLength) {
+		if start >= int64(contentLength) {
 			w.WriteHeader(http.StatusRequestedRangeNotSatisfiable)
 			return
 		}
+
+		if matches[2] != "" {
+			end, err = strconv.ParseInt(matches[2], 10, 64)
+			if err != nil || end < 0 {
+				w.WriteHeader(http.StatusBadRequest)
+				return
+			}
+
+			if end >= int64(contentLength) {
+				w.WriteHeader(http.StatusRequestedRangeNotSatisfiable)
+				return
+			}
+		}
+
+		statusCode = http.StatusPartialContent
 	}
 
 	var currentContentLength = end - start + 1
@@ -91,9 +95,14 @@ func handleResumableLayer(data []byte, w http.ResponseWriter, r *http.Request, t
 
 	end = start + currentContentLength - 1
 
-	w.Header().Set("Content-Range", fmt.Sprintf("bytes %d-%d/%d", start, end, contentLength))
-	w.Header().Set("Content-Length", strconv.FormatInt(currentContentLength, 10))
-	w.WriteHeader(http.StatusPartialContent)
+	if statusCode == http.StatusPartialContent {
+		w.Header().Set("Content-Length", strconv.FormatInt(currentContentLength, 10))
+		w.Header().Set("Content-Range", fmt.Sprintf("bytes %d-%d/%d", start, end, contentLength))
+	} else {
+		w.Header().Set("Content-Length", strconv.FormatInt(contentLength, 10))
+	}
+
+	w.WriteHeader(statusCode)
 	w.Write(data[start : end+1])
 	time.Sleep(time.Second)
 }