gateway: compress frontend input payloads

Signed-off-by: CrazyMax <1951866+crazy-max@users.noreply.github.com>

Author: crazy-max

frontend/dockerfile/dockerfile_test.go

@@ -19,12 +19,15 @@ import (
 	"runtime"
 	"slices"
 	"strings"
+	"sync/atomic"
 	"testing"
 	"time"
 
 	cacheimporttypes "github.com/moby/buildkit/cache/remotecache/v1/types"
 	"github.com/tonistiigi/fsutil"
 	"golang.org/x/sync/errgroup"
+	"google.golang.org/grpc"
+	grpcgzip "google.golang.org/grpc/encoding/gzip"
 	"google.golang.org/grpc/status"
 	"google.golang.org/protobuf/proto"
 	"google.golang.org/protobuf/types/known/anypb"
@@ -44,6 +47,7 @@ import (
 	"github.com/moby/buildkit/frontend/dockerfile/builder"
 	"github.com/moby/buildkit/frontend/dockerui"
 	gateway "github.com/moby/buildkit/frontend/gateway/client"
+	gatewayapi "github.com/moby/buildkit/frontend/gateway/pb"
 	"github.com/moby/buildkit/frontend/subrequests"
 	"github.com/moby/buildkit/identity"
 	"github.com/moby/buildkit/session"
@@ -141,6 +145,7 @@ var allTests = integration.TestFuncs(
 	testFrontendUseForwardedSolveResults,
 	testFrontendEvaluate,
 	testFrontendInputs,
+	testFrontendInputsCompressed,
 	testErrorsSourceMap,
 	testMultiArgs,
 	testFrontendSubrequests,
@@ -7340,6 +7345,72 @@ COPY foo foo2
 	require.Equal(t, expected, actual)
 }
 
+type inputsCompressionInterceptor struct {
+	calls      atomic.Int64
+	compressed atomic.Bool
+}
+
+func (i *inputsCompressionInterceptor) unary(ctx context.Context, method string, req any, reply any, cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error {
+	if method == gatewayapi.LLBBridge_Inputs_FullMethodName {
+		i.calls.Add(1)
+		for _, opt := range opts {
+			if compressor, ok := opt.(grpc.CompressorCallOption); ok && compressor.CompressorType == grpcgzip.Name {
+				i.compressed.Store(true)
+			}
+		}
+	}
+	return invoker(ctx, method, req, reply, cc, opts...)
+}
+
+func testFrontendInputsCompressed(t *testing.T, sb integration.Sandbox) {
+	integration.SkipOnPlatform(t, "windows")
+	f := getFrontend(t, sb)
+	if _, ok := f.(*clientFrontend); !ok {
+		t.Skip("only test with client frontend")
+	}
+
+	inputsInterceptor := &inputsCompressionInterceptor{}
+	c, err := client.New(sb.Context(), sb.Address(), client.WithGRPCDialOption(grpc.WithChainUnaryInterceptor(inputsInterceptor.unary)))
+	require.NoError(t, err)
+	defer c.Close()
+
+	payload := bytes.Repeat([]byte("a"), 64<<10)
+	input := llb.Scratch().File(llb.Mkfile("/payload", 0600, payload))
+
+	destDir := t.TempDir()
+	dockerfile := []byte(`
+FROM scratch
+COPY payload payload
+`)
+	dir := integration.Tmpdir(
+		t,
+		fstest.CreateFile("Dockerfile", dockerfile, 0600),
+	)
+
+	_, err = f.Solve(sb.Context(), c, client.SolveOpt{
+		Exports: []client.ExportEntry{
+			{
+				Type:      client.ExporterLocal,
+				OutputDir: destDir,
+			},
+		},
+		LocalMounts: map[string]fsutil.FS{
+			dockerui.DefaultLocalNameDockerfile: dir,
+		},
+		FrontendInputs: map[string]llb.State{
+			dockerui.DefaultLocalNameContext: input,
+		},
+	}, nil)
+	require.NoError(t, err)
+
+	actual, err := os.ReadFile(filepath.Join(destDir, "payload"))
+	require.NoError(t, err)
+	require.Equal(t, payload, actual)
+
+	require.NotZero(t, inputsInterceptor.calls.Load(), "expected an Inputs request")
+	require.True(t, inputsInterceptor.compressed.Load(), "expected Inputs request to use gzip compression")
+}
+
 func testFrontendSubrequests(t *testing.T, sb integration.Sandbox) {
 	f := getFrontend(t, sb)
 	if _, ok := f.(*clientFrontend); !ok {

frontend/gateway/gateway.go

@@ -57,6 +57,7 @@ import (
 	spb "google.golang.org/genproto/googleapis/rpc/status"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
+	_ "google.golang.org/grpc/encoding/gzip" // register gzip compressor for gateway RPC responses
 	"google.golang.org/grpc/health"
 	"google.golang.org/grpc/health/grpc_health_v1"
 	"google.golang.org/grpc/status"

frontend/gateway/grpcclient/client.go

@@ -35,6 +35,7 @@ import (
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/credentials/insecure"
+	"google.golang.org/grpc/encoding/gzip"
 	"google.golang.org/grpc/status"
 )
 
@@ -412,7 +413,12 @@ func (c *grpcClient) Solve(ctx context.Context, creq client.SolveRequest) (res *
 		}
 	}
 
-	resp, err := c.client.Solve(ctx, req)
+	var opts []grpc.CallOption
+	if len(req.FrontendInputs) > 0 {
+		opts = append(opts, grpc.UseCompressor(gzip.Name))
+	}
+
+	resp, err := c.client.Solve(ctx, req, opts...)
 	if err != nil {
 		return nil, err
 	}
@@ -800,7 +806,7 @@ func (c *grpcClient) Inputs(ctx context.Context) (map[string]llb.State, error) {
 		return nil, err
 	}
 
-	resp, err := c.client.Inputs(ctx, &pb.InputsRequest{})
+	resp, err := c.client.Inputs(ctx, &pb.InputsRequest{}, grpc.UseCompressor(gzip.Name))
 	if err != nil {
 		return nil, err
 	}