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,16 @@ import (
 	"runtime"
 	"slices"
 	"strings"
+	"sync"
 	"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/stats"
 	"google.golang.org/grpc/status"
 	"google.golang.org/protobuf/proto"
 	"google.golang.org/protobuf/types/known/anypb"
@@ -141,6 +145,7 @@ var allTests = integration.TestFuncs(
 	testFrontendUseForwardedSolveResults,
 	testFrontendEvaluate,
 	testFrontendInputs,
+	testFrontendInputsCompressed,
 	testErrorsSourceMap,
 	testMultiArgs,
 	testFrontendSubrequests,
@@ -7340,6 +7345,115 @@ COPY foo foo2
 	require.Equal(t, expected, actual)
 }
 
+type inputsPayloadStats struct {
+	mu              sync.Mutex
+	payloads        int
+	inCompressions  []string
+	outCompressions []string
+}
+
+type inputsPayloadStatsMethodKey struct{}
+
+func (s *inputsPayloadStats) TagRPC(ctx context.Context, info *stats.RPCTagInfo) context.Context {
+	return context.WithValue(ctx, inputsPayloadStatsMethodKey{}, info.FullMethodName)
+}
+
+func (s *inputsPayloadStats) HandleRPC(ctx context.Context, stat stats.RPCStats) {
+	method, _ := ctx.Value(inputsPayloadStatsMethodKey{}).(string)
+	if !strings.HasSuffix(method, "/LLBBridge/Inputs") {
+		return
+	}
+	switch stat := stat.(type) {
+	case *stats.InHeader:
+		if stat.Client {
+			s.mu.Lock()
+			defer s.mu.Unlock()
+			s.inCompressions = append(s.inCompressions, stat.Compression)
+		}
+	case *stats.OutHeader:
+		if stat.Client {
+			s.mu.Lock()
+			defer s.mu.Unlock()
+			s.outCompressions = append(s.outCompressions, stat.Compression)
+		}
+	case *stats.InPayload:
+		if stat.Client {
+			s.mu.Lock()
+			defer s.mu.Unlock()
+			s.payloads++
+		}
+	}
+}
+
+func (s *inputsPayloadStats) TagConn(ctx context.Context, _ *stats.ConnTagInfo) context.Context {
+	return ctx
+}
+
+func (s *inputsPayloadStats) HandleConn(context.Context, stats.ConnStats) {}
+
+func (s *inputsPayloadStats) inputsPayloadCount() int {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	return s.payloads
+}
+
+func (s *inputsPayloadStats) inputsCompressions() (in, out []string) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	return slices.Clone(s.inCompressions), slices.Clone(s.outCompressions)
+}
+
+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")
+	}
+
+	payloadStats := &inputsPayloadStats{}
+	c, err := client.New(sb.Context(), sb.Address(), client.WithGRPCDialOption(grpc.WithStatsHandler(payloadStats)))
+	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)
+
+	inCompressions, outCompressions := payloadStats.inputsCompressions()
+	require.Contains(t, outCompressions, grpcgzip.Name, "expected Inputs request to use gzip compression")
+	require.Contains(t, inCompressions, grpcgzip.Name, "expected Inputs response to use gzip compression")
+	require.NotZero(t, payloadStats.inputsPayloadCount(), "expected an Inputs response payload")
+}
+
 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
 	}