siso: split reapi casDataConn from conn, casConn To avoid gRPC stream saturation in two phase caching, have separate casDataConn for long gRPC stream request for bytestream. Bug: b/539291862 Change-Id: I38823011bd712e2ab45488fc0abaa67fc65eb95b Reviewed-on: https://chromium-review.googlesource.com/c/build/+/8165662 Auto-Submit: Fumitoshi Ukai <ukai@google.com> Reviewed-by: Yuke Liao <liaoyuke@google.com> Commit-Queue: Fumitoshi Ukai <ukai@google.com>
diff --git a/siso/reapi/cache.go b/siso/reapi/cache.go index 45fbdf3..4d4e15d 100644 --- a/siso/reapi/cache.go +++ b/siso/reapi/cache.go
@@ -133,7 +133,7 @@ // or if an error occurs during setup, otherwise we leak the stream until the // parent context (step context) is canceled. ctx, cancel := context.WithCancel(ctx) - r, err := bytestreamio.Open(ctx, bpb.NewByteStreamClient(s.c.casConn), s.c.resourceName(s.d)) + r, err := bytestreamio.Open(ctx, bpb.NewByteStreamClient(s.c.casDataConn), s.c.resourceName(s.d)) if err != nil { cancel() s.c.m.ReadDone(0, err)
diff --git a/siso/reapi/cas.go b/siso/reapi/cas.go index 7a4cc43..fd503d3 100644 --- a/siso/reapi/cas.go +++ b/siso/reapi/cas.go
@@ -428,7 +428,7 @@ err := retry.Do(ctx, func() error { ctx, cancel := fetch.ContextWithTimeout(ctx, d) defer cancel() - r, err := bytestreamio.Open(ctx, bpb.NewByteStreamClient(c.casConn), resourceName) + r, err := bytestreamio.Open(ctx, bpb.NewByteStreamClient(c.casDataConn), resourceName) if err != nil { c.m.ReadDone(0, err) return err @@ -1044,7 +1044,7 @@ func (c *Client) uploadWithByteStream(ctx context.Context, digests []digest.Digest, uploads map[digest.Digest]*uploadOp, ds *blob.Store) []missingBlob { clog.Infof(ctx, "upload by streaming %d", len(digests)) - bsClient := bpb.NewByteStreamClient(c.casConn) + bsClient := bpb.NewByteStreamClient(c.casDataConn) var ( mu sync.Mutex missing []missingBlob
diff --git a/siso/reapi/conn_close_test.go b/siso/reapi/conn_close_test.go index d28044c..54d45a4 100644 --- a/siso/reapi/conn_close_test.go +++ b/siso/reapi/conn_close_test.go
@@ -41,7 +41,8 @@ ctx := t.Context() conn := &fakeConn{} casConn := &fakeConn{} - c, err := NewFromConn(ctx, Option{}, cred.Cred{}, conn, casConn) + casDataConn := &fakeConn{} + c, err := NewFromConn(ctx, Option{}, cred.Cred{}, conn, casConn, casDataConn) if err != nil { t.Fatalf("NewFromConn: %v", err) } @@ -54,12 +55,15 @@ if got, want := casConn.closed, 1; got != want { t.Errorf("casConn.closed = %d; want %d", got, want) } + if got, want := casDataConn.closed, 1; got != want { + t.Errorf("casDataConn.closed = %d; want %d", got, want) + } } func TestClose_SharedConnClosedOnce(t *testing.T) { ctx := t.Context() conn := &fakeConn{} - c, err := NewFromConn(ctx, Option{}, cred.Cred{}, conn, conn) + c, err := NewFromConn(ctx, Option{}, cred.Cred{}, conn, conn, conn) if err != nil { t.Fatalf("NewFromConn: %v", err) } @@ -75,7 +79,8 @@ ctx := t.Context() conn := &fakeConn{} casConn := &fakeConn{} - _, err := NewFromConn(ctx, Option{DigestFunction: "bogus"}, cred.Cred{}, conn, casConn) + casDataConn := &fakeConn{} + _, err := NewFromConn(ctx, Option{DigestFunction: "bogus"}, cred.Cred{}, conn, casConn, casDataConn) if err == nil { t.Fatal("NewFromConn succeeded; want digest function error") } @@ -85,13 +90,17 @@ if got, want := casConn.closed, 1; got != want { t.Errorf("casConn.closed = %d; want %d", got, want) } + if got, want := casDataConn.closed, 1; got != want { + t.Errorf("casDataConn.closed = %d; want %d", got, want) + } } func TestInit_CapabilitiesErrorClosesBothConns(t *testing.T) { ctx := t.Context() conn := &fakeConn{invokeErr: status.Error(codes.InvalidArgument, "no capabilities")} casConn := &fakeConn{} - c, err := NewFromConn(ctx, Option{}, cred.Cred{}, conn, casConn) + casDataConn := &fakeConn{} + c, err := NewFromConn(ctx, Option{}, cred.Cred{}, conn, casConn, casDataConn) if err != nil { t.Fatalf("NewFromConn: %v", err) } @@ -104,6 +113,9 @@ if got, want := casConn.closed, 1; got != want { t.Errorf("casConn.closed = %d; want %d", got, want) } + if got, want := casDataConn.closed, 1; got != want { + t.Errorf("casDataConn.closed = %d; want %d", got, want) + } } func TestInit_DigestFunctionValidationErrorClosesBothConns(t *testing.T) { @@ -112,7 +124,8 @@ // function other than sha256 fails validation. conn := &fakeConn{} casConn := &fakeConn{} - c, err := NewFromConn(ctx, Option{DigestFunction: "sha1"}, cred.Cred{}, conn, casConn) + casDataConn := &fakeConn{} + c, err := NewFromConn(ctx, Option{DigestFunction: "sha1"}, cred.Cred{}, conn, casConn, casDataConn) if err != nil { t.Fatalf("NewFromConn: %v", err) } @@ -125,4 +138,7 @@ if got, want := casConn.closed, 1; got != want { t.Errorf("casConn.closed = %d; want %d", got, want) } + if got, want := casDataConn.closed, 1; got != want { + t.Errorf("casDataConn.closed = %d; want %d", got, want) + } }
diff --git a/siso/reapi/proxy.go b/siso/reapi/proxy.go index 082416b..2b0b6e0 100644 --- a/siso/reapi/proxy.go +++ b/siso/reapi/proxy.go
@@ -132,7 +132,7 @@ ep := &executionProxy{client: rpb.NewExecutionClient(p.client.conn)} rpb.RegisterExecutionServer(server, ep) - bsp := &byteStreamProxy{client: bspb.NewByteStreamClient(p.client.casConn)} + bsp := &byteStreamProxy{client: bspb.NewByteStreamClient(p.client.casDataConn)} bspb.RegisterByteStreamServer(server, bsp) errCh := make(chan error, 1)
diff --git a/siso/reapi/reapi.go b/siso/reapi/reapi.go index ef562fd..aff144f 100644 --- a/siso/reapi/reapi.go +++ b/siso/reapi/reapi.go
@@ -108,6 +108,7 @@ WalkDirStream bool ConnPool int + DataConnPool int KeepAliveParams keepalive.ClientParameters // StatsHandler, if non-nil, is wired into every gRPC ClientConn via @@ -209,6 +210,7 @@ fs.BoolVar(&o.WalkDirStream, o.Prefix+"_walkdir_stream", false, "use GetTree for WalkDir") fs.IntVar(&o.ConnPool, o.Prefix+"_grpc_conn_pool", 25, "grpc connection pool") + fs.IntVar(&o.DataConnPool, o.Prefix+"_grpc_data_conn_pool", 25, "grpc connection pool for data (bytestream Read, Write)") // https://grpc.io/docs/guides/keepalive/#keepalive-configuration-specification // b/286237547 - RBE suggests 30s @@ -323,11 +325,12 @@ // Client is a remote exec API client. type Client struct { - opt Option - digestFn digest.Function - cred cred.Cred - conn grpcClientConn - casConn grpcClientConn + opt Option + digestFn digest.Function + cred cred.Cred + conn grpcClientConn + casConn grpcClientConn + casDataConn grpcClientConn mu sync.Mutex capabilities *rpb.ServerCapabilities @@ -416,24 +419,38 @@ opt.BatchCompressedBlob = 0 clog.Warningf(ctx, "disabling blob compression because grpc compression is enabled") } - clog.Infof(ctx, "address: %q instance: %q", opt.Address, opt.Instance) - conn, err := newConn(ctx, opt.Address, cred, opt) + address := opt.Address + clog.Infof(ctx, "address: %q (%d) instance: %q", address, opt.ConnPool, opt.Instance) + conn, err := newConn(ctx, address, cred, opt.ConnPool, opt) if err != nil { return nil, DialError{Err: err} } casConn := conn if opt.CASAddress != "" { - clog.Infof(ctx, "cas address: %q", opt.CASAddress) - casConn, err = newConn(ctx, opt.CASAddress, cred, opt) + address = opt.CASAddress + clog.Infof(ctx, "cas address: %q (%d)", address, opt.ConnPool) + casConn, err = newConn(ctx, address, cred, opt.ConnPool, opt) if err != nil { conn.Close() return nil, DialError{Err: err} } } - return NewFromConn(ctx, opt, cred, conn, casConn) + casDataConn := casConn + if opt.DataConnPool > 0 { + clog.Infof(ctx, "data address: %q (%d)", address, opt.DataConnPool) + casDataConn, err = newConn(ctx, address, cred, opt.DataConnPool, opt) + if err != nil { + conn.Close() + if conn != casConn { + casConn.Close() + } + return nil, DialError{Err: err} + } + } + return NewFromConn(ctx, opt, cred, conn, casConn, casDataConn) } -func newConn(ctx context.Context, addr string, cred cred.Cred, opt Option) (grpcClientConn, error) { +func newConn(ctx context.Context, addr string, cred cred.Cred, pool int, opt Option) (grpcClientConn, error) { // Force the gRPC DNS resolver by prefixing "dns:///" when the caller // did not supply a scheme. gtransport.DialPool rides the deprecated // grpc.DialContext path whose default resolver is "passthrough", @@ -449,7 +466,7 @@ } copts := []option.ClientOption{ option.WithEndpoint(endpoint), - option.WithGRPCConnectionPool(opt.ConnPool), + option.WithGRPCConnectionPool(pool), } if !isGoogleRBE(addr) { // disable Google Application Default for non RBE backend. @@ -550,8 +567,8 @@ return conn, nil } -// NewFromConn creates new remote exec API client from conn and casConn. -func NewFromConn(ctx context.Context, opt Option, cred cred.Cred, conn, casConn grpcClientConn) (*Client, error) { +// NewFromConn creates new remote exec API client from conn, casConn, and casDataConn. +func NewFromConn(ctx context.Context, opt Option, cred cred.Cred, conn, casConn, casDataConn grpcClientConn) (*Client, error) { zstdDecoderPool := &sync.Pool{} zstdDecoderPool.New = func() any { d, err := zstd.NewReader(nil, zstdDecoderOpts...) @@ -570,6 +587,9 @@ if casConn != nil && casConn != conn { casConn.Close() } + if casDataConn != nil && casDataConn != casConn { + casDataConn.Close() + } return nil, err } c := &Client{ @@ -578,6 +598,7 @@ cred: cred, conn: conn, casConn: casConn, + casDataConn: casDataConn, zstdDecoderPool: zstdDecoderPool, m: iometrics.New("reapi"), } @@ -724,6 +745,12 @@ err = cerr } } + if c.casDataConn != nil && c.casDataConn != c.casConn { + cerr := c.casDataConn.Close() + if err == nil { + err = cerr + } + } return err } @@ -785,6 +812,7 @@ // Decide before cancel() overwrites the deadline cause. keep := keepFirstAttempt(callCtx, cause, err) cancel() + clog.Infof(ctx, "GetActionResult keep=%t err=%v", keep, err) if keep { c.m.OpsDone(err) return result, err
diff --git a/siso/reapi/reapi_test.go b/siso/reapi/reapi_test.go index ea2d763..256b652 100644 --- a/siso/reapi/reapi_test.go +++ b/siso/reapi/reapi_test.go
@@ -352,7 +352,7 @@ Instance: "test", KeepExecStream: true, DigestFunction: "blake3", - }, cred.Cred{}, conn, conn) + }, cred.Cred{}, conn, conn, conn) if err != nil { t.Fatal(err) } @@ -398,7 +398,7 @@ cl, err := reapi.NewFromConn(ctx, reapi.Option{ Instance: "test", KeepExecStream: true, - }, cred.Cred{}, conn, conn) + }, cred.Cred{}, conn, conn, conn) if err != nil { t.Fatal(err) }
diff --git a/siso/reapi/reapitest/localexec.go b/siso/reapi/reapitest/localexec.go index d436a79..16744b9 100644 --- a/siso/reapi/reapitest/localexec.go +++ b/siso/reapi/reapitest/localexec.go
@@ -101,7 +101,7 @@ Address: "bufconn", Instance: "projects/siso-test/instances/default_instance", } - client, err := reapi.NewFromConn(ctx, opt, cred.Cred{}, conn, conn) + client, err := reapi.NewFromConn(ctx, opt, cred.Cred{}, conn, conn, conn) if err != nil { t.Fatal(err) }
diff --git a/siso/reapi/reapitest/reapitest.go b/siso/reapi/reapitest/reapitest.go index d5068d5..bf00515 100644 --- a/siso/reapi/reapitest/reapitest.go +++ b/siso/reapi/reapitest/reapitest.go
@@ -160,7 +160,7 @@ if err != nil { t.Fatal(err) } - client, err := reapi.NewFromConn(ctx, opt, cred.Cred{}, conn, conn) + client, err := reapi.NewFromConn(ctx, opt, cred.Cred{}, conn, conn, conn) if err != nil { t.Fatal(err) }