Improve image pull performance from http 1.1 container registries

amrmahdi · amrmahdi · commit 289130b8a776 · 2020-10-26T05:05:09.000Z
Private registries that does not support http 2.0 such as Azure Container Registry streams back content in a max of 16KB chunks (max TLS record size). The small chunks introduce an overhead when copying the layers to the content store sine each chunk incurs the overhead of  grpc message that has to be sent to the content store.

This change reduces this overhead by buffering the chunks into 1MB chunks and only then writes a message to the content store.

Below is a per comparsion between the 2 approaches using a couple of large images that are being pulled from the docker hub (http 2.0) and a private Azure CR (http 1.1) in seconds.

image                                                   | Buffered copy | master
-------                                                 |---------------|----------
docker.io/pytorch/pytorch:latest                        |  55.63        | 58.33
docker.io/nvidia/cuda:latest                            |  72.05        | 75.98
containerdpulltest.azurecr.io/pytorch/pytorch:latest    | 61.45         | 77.1
containerdpulltest.azurecr.io/nvidia/cuda:latest        | 77.13         | 85.47

Signed-off-by: Amr Mahdi &lt;amramahdi@gmail.com&gt;
diff --git a/content/helpers.go b/content/helpers.go
@@ -230,8 +230,31 @@ func seekReader(r io.Reader, offset, size int64) (io.Reader, error) {
 }
 
 func copyWithBuffer(dst io.Writer, src io.Reader) (written int64, err error) {
-	buf := bufPool.Get().(*[]byte)
-	written, err = io.CopyBuffer(dst, src, *buf)
-	bufPool.Put(buf)
+	bufRef := bufPool.Get().(*[]byte)
+	defer bufPool.Put(bufRef)
+	buf := *bufRef
+	for {
+		nr, er := io.ReadAtLeast(src, buf, len(buf))
+		if nr > 0 {
+			nw, ew := dst.Write(buf[0:nr])
+			if nw > 0 {
+				written += int64(nw)
+			}
+			if ew != nil {
+				err = ew
+				break
+			}
+			if nr != nw {
+				err = io.ErrShortWrite
+				break
+			}
+		}
+		if er != nil {
+			if er != io.EOF && er != io.ErrUnexpectedEOF {
+				err = er
+			}
+			break
+		}
+	}
 	return
 }
