proxy: break up writes from the remote writer to avoid grpc limits

jsternberg · jsternberg · commit eaa7ca80dcc1 · 2025-02-28T11:29:51.000-06:00
The remote content writer proxy already has the capability to break up large files into multiple writes, but the current API doesn't recognize when it's about to exceed the limits and attempts to send the data over grpc in one message instead of breaking it into multiple messages. This changes the behavior of `Write` to automatically break up the size of the content based on the max send message size. Signed-off-by: Jonathan A. Sternberg <jonathan.sternberg@docker.com> (cherry picked from commit f25f36c)
diff --git a/core/content/proxy/content_writer.go b/core/content/proxy/content_writer.go
@@ -26,6 +26,7 @@ import (
 	digest "github.com/opencontainers/go-digest"
 
 	"github.com/containerd/containerd/v2/core/content"
+	"github.com/containerd/containerd/v2/defaults"
 	"github.com/containerd/containerd/v2/pkg/protobuf"
 )
 
@@ -76,27 +77,37 @@ func (rw *remoteWriter) Digest() digest.Digest {
 }
 
 func (rw *remoteWriter) Write(p []byte) (n int, err error) {
-	offset := rw.offset
+	const maxBufferSize = defaults.DefaultMaxSendMsgSize >> 1
+	for i := 0; i < len(p); i += maxBufferSize {
+		offset := rw.offset
 
-	resp, err := rw.send(&contentapi.WriteContentRequest{
-		Action: contentapi.WriteAction_WRITE,
-		Offset: offset,
-		Data:   p,
-	})
-	if err != nil {
-		return 0, fmt.Errorf("failed to send write: %w", errgrpc.ToNative(err))
-	}
+		end := i + maxBufferSize
+		if end > len(p) {
+			end = len(p)
+		}
+		data := p[i:end]
+
+		resp, err := rw.send(&contentapi.WriteContentRequest{
+			Action: contentapi.WriteAction_WRITE,
+			Offset: offset,
+			Data:   data,
+		})
+		if err != nil {
+			return 0, fmt.Errorf("failed to send write: %w", errgrpc.ToNative(err))
+		}
 
-	n = int(resp.Offset - offset)
-	if n < len(p) {
-		err = io.ErrShortWrite
-	}
+		written := int(resp.Offset - offset)
+		rw.offset += int64(written)
+		if resp.Digest != "" {
+			rw.digest = digest.Digest(resp.Digest)
+		}
+		n += written
 
-	rw.offset += int64(n)
-	if resp.Digest != "" {
-		rw.digest = digest.Digest(resp.Digest)
+		if written < len(data) {
+			return n, io.ErrShortWrite
+		}
 	}
-	return
+	return n, nil
 }
 
 func (rw *remoteWriter) Commit(ctx context.Context, size int64, expected digest.Digest, opts ...content.Opt) (err error) {
diff --git a/core/content/testsuite/testsuite.go b/core/content/testsuite/testsuite.go
@@ -175,6 +175,13 @@ func checkContentStoreWriter(ctx context.Context, t *testing.T, cs content.Store
 	}
 	defer w4.Close()
 
+	c5, d5 := createContent(16 << 21)
+	w5, err := content.OpenWriter(ctx, cs, content.WithRef("c5"), content.WithDescriptor(ocispec.Descriptor{Size: int64(len(c5)), Digest: d5}))
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer w5.Close()
+
 	smallbuf := make([]byte, 32)
 	for _, s := range []struct {
 		content []byte
@@ -201,6 +208,11 @@ func checkContentStoreWriter(ctx context.Context, t *testing.T, cs content.Store
 			digest:  d4,
 			writer:  w4,
 		},
+		{
+			content: c5,
+			digest:  d5,
+			writer:  w5,
+		},
 	} {
 		n, err := io.CopyBuffer(s.writer, bytes.NewReader(s.content), smallbuf)
 		if err != nil {
@@ -647,7 +659,6 @@ func checkLabels(ctx context.Context, t *testing.T, cs content.Store) {
 	if err := checkInfo(ctx, cs, d1, info, preCommit, postCommit, preUpdate, postUpdate); err != nil {
 		t.Fatalf("Check info failed: %+v", err)
 	}
-
 }
 
 func checkResume(rf func(context.Context, content.Writer, []byte, int64, int64, digest.Digest) error) func(ctx context.Context, t *testing.T, cs content.Store) {
@@ -734,7 +745,6 @@ func resumeDiscard(ctx context.Context, w content.Writer, b []byte, written, siz
 	}
 	if err := w.Commit(ctx, size, dgst); err != nil {
 		return fmt.Errorf("commit failed: %w", err)
-
 	}
 	return nil
 }
@@ -931,7 +941,6 @@ func checkCrossNSAppend(ctx context.Context, t *testing.T, cs content.Store) {
 	if err := checkContent(ctx2, cs, d2, info2, t1, t3, t1, t3); err != nil {
 		t.Fatal(err)
 	}
-
 }
 
 func checkCrossNSIsolate(ctx context.Context, t *testing.T, cs content.Store) {
@@ -1129,6 +1138,7 @@ func checkInfo(ctx context.Context, cs content.Store, d digest.Digest, expected
 
 	return nil
 }
+
 func checkContent(ctx context.Context, cs content.Store, d digest.Digest, expected content.Info, c1, c2, u1, u2 time.Time) error {
 	if err := checkInfo(ctx, cs, d, expected, c1, c2, u1, u2); err != nil {
 		return err

Original file line number	Diff line number	Diff line change
`@@ -175,6 +175,13 @@ func checkContentStoreWriter(ctx context.Context, t *testing.T, cs content.Store`
`175`	`175`	`}`
`176`	`176`	`defer w4.Close()`
`177`	`177`
	`178`	`+ c5, d5 := createContent(16 << 21)`
	`179`	`+ w5, err := content.OpenWriter(ctx, cs, content.WithRef("c5"), content.WithDescriptor(ocispec.Descriptor{Size: int64(len(c5)), Digest: d5}))`
	`180`	`+ if err != nil {`
	`181`	`+ t.Fatal(err)`
	`182`	`+ }`
	`183`	`+ defer w5.Close()`
	`184`	`+`
`178`	`185`	`smallbuf := make([]byte, 32)`
`179`	`186`	`for _, s := range []struct {`
`180`	`187`	`content []byte`
`@@ -201,6 +208,11 @@ func checkContentStoreWriter(ctx context.Context, t *testing.T, cs content.Store`
`201`	`208`	`digest: d4,`
`202`	`209`	`writer: w4,`
`203`	`210`	`},`
	`211`	`+ {`
	`212`	`+ content: c5,`
	`213`	`+ digest: d5,`
	`214`	`+ writer: w5,`
	`215`	`+ },`
`204`	`216`	`} {`
`205`	`217`	`n, err := io.CopyBuffer(s.writer, bytes.NewReader(s.content), smallbuf)`
`206`	`218`	`if err != nil {`
`@@ -647,7 +659,6 @@ func checkLabels(ctx context.Context, t *testing.T, cs content.Store) {`
`647`	`659`	`if err := checkInfo(ctx, cs, d1, info, preCommit, postCommit, preUpdate, postUpdate); err != nil {`
`648`	`660`	`t.Fatalf("Check info failed: %+v", err)`
`649`	`661`	`}`
`650`		`-`
`651`	`662`	`}`
`652`	`663`
`653`	`664`	`func checkResume(rf func(context.Context, content.Writer, []byte, int64, int64, digest.Digest) error) func(ctx context.Context, t *testing.T, cs content.Store) {`
`@@ -734,7 +745,6 @@ func resumeDiscard(ctx context.Context, w content.Writer, b []byte, written, siz`
`734`	`745`	`}`
`735`	`746`	`if err := w.Commit(ctx, size, dgst); err != nil {`
`736`	`747`	`return fmt.Errorf("commit failed: %w", err)`
`737`		`-`
`738`	`748`	`}`
`739`	`749`	`return nil`
`740`	`750`	`}`
`@@ -931,7 +941,6 @@ func checkCrossNSAppend(ctx context.Context, t *testing.T, cs content.Store) {`
`931`	`941`	`if err := checkContent(ctx2, cs, d2, info2, t1, t3, t1, t3); err != nil {`
`932`	`942`	`t.Fatal(err)`
`933`	`943`	`}`
`934`		`-`
`935`	`944`	`}`
`936`	`945`
`937`	`946`	`func checkCrossNSIsolate(ctx context.Context, t *testing.T, cs content.Store) {`
`@@ -1129,6 +1138,7 @@ func checkInfo(ctx context.Context, cs content.Store, d digest.Digest, expected`
`1129`	`1138`
`1130`	`1139`	`return nil`
`1131`	`1140`	`}`
	`1141`	`+`
`1132`	`1142`	`func checkContent(ctx context.Context, cs content.Store, d digest.Digest, expected content.Info, c1, c2, u1, u2 time.Time) error {`
`1133`	`1143`	`if err := checkInfo(ctx, cs, d, expected, c1, c2, u1, u2); err != nil {`
`1134`	`1144`	`return err`