Merge pull request #4313 from tonistiigi/session-grpc-message-size-fix

session: fix grpc message size limits for tar streams

session/filesync/diffcopy.go

@@ -47,6 +47,22 @@ type streamWriterCloser struct {
 }
 
 func (wc *streamWriterCloser) Write(dt []byte) (int, error) {
+	// grpc-go has a 4MB limit on messages by default. Split large messages
+	// so we don't get close to that limit.
+	const maxChunkSize = 3 * 1024 * 1024
+	if len(dt) > maxChunkSize {
+		n1, err := wc.Write(dt[:maxChunkSize])
+		if err != nil {
+			return n1, err
+		}
+		dt = dt[maxChunkSize:]
+		var n2 int
+		if n2, err = wc.Write(dt); err != nil {
+			return n1 + n2, err
+		}
+		return n1 + n2, nil
+	}
+
 	if err := wc.ClientStream.SendMsg(&BytesMessage{Data: dt}); err != nil {
 		// SendMsg return EOF on remote errors
 		if errors.Is(err, io.EOF) {

session/grpc.go

@@ -7,6 +7,7 @@ import (
 	"sync/atomic"
 	"time"
 
+	"github.com/containerd/containerd/defaults"
 	grpc_middleware "github.com/grpc-ecosystem/go-grpc-middleware"
 	"github.com/moby/buildkit/util/bklog"
 	"github.com/moby/buildkit/util/grpcerrors"
@@ -44,6 +45,8 @@ func grpcClientConn(ctx context.Context, conn net.Conn) (context.Context, *grpc.
 	dialOpts := []grpc.DialOption{
 		dialer,
 		grpc.WithTransportCredentials(insecure.NewCredentials()),
+		grpc.WithDefaultCallOptions(grpc.MaxCallRecvMsgSize(defaults.DefaultMaxRecvMsgSize)),
+		grpc.WithDefaultCallOptions(grpc.MaxCallSendMsgSize(defaults.DefaultMaxSendMsgSize)),
 	}
 
 	if span := trace.SpanFromContext(ctx); span.SpanContext().IsValid() {

session/upload/uploadprovider/provider.go

@@ -59,6 +59,20 @@ type writer struct {
 }
 
 func (w *writer) Write(dt []byte) (int, error) {
+	// avoid sending too big messages on grpc stream
+	const maxChunkSize = 3 * 1024 * 1024
+	if len(dt) > maxChunkSize {
+		n1, err := w.Write(dt[:maxChunkSize])
+		if err != nil {
+			return n1, err
+		}
+		dt = dt[maxChunkSize:]
+		var n2 int
+		if n2, err := w.Write(dt); err != nil {
+			return n1 + n2, err
+		}
+		return n1 + n2, nil
+	}
 	if err := w.SendMsg(&upload.BytesMessage{Data: dt}); err != nil {
 		return 0, err
 	}