Open fewer connections to rex servers

docs/config.html

@@ -963,6 +963,19 @@ <h3 class="mt1 f6 lh-title" id="remote.numexecutors">
         <p>Defines the number of remote executors to use simultaneously.</p>
       </div>
     </li>
+    <li>
+      <div>
+        <h3 class="mt1 f6 lh-title" id="remote.casconnections">
+          CASConnections <span class="normal">(int)</span>
+        </h3>
+
+        <p>
+          Number of connections to open to the CAS server. More connections can
+          improve throughput for large transfers, at the cost of more load on the
+          server. Defaults to 4.
+        </p>
+      </div>
+    </li>
     <li>
       <div>
         <h3 class="mt1 f6 lh-title" id="remote.name">Name</h3>

src/core/config.go

@@ -407,6 +407,7 @@ func DefaultConfiguration() *Configuration {
 	config.Display.MaxWorkers = 40
 	config.Display.ColourScheme = "dark"
 	config.Remote.NumExecutors = 20 // kind of arbitrary
+	config.Remote.CASConnections = 4
 	config.Remote.Secure = true
 	config.Remote.VerifyOutputs = true
 	config.Remote.UploadDirs = true
@@ -566,6 +567,7 @@ type Configuration struct {
 		CASURL                  string       `help:"URL for the CAS service, if it is different to the main one."`
 		AssetURL                string       `help:"URL for the remote asset server, if it is different to the main one."`
 		NumExecutors            int          `help:"Maximum number of remote executors to use simultaneously."`
+		CASConnections          int          `help:"Number of connections to open to the CAS server. More connections can improve throughput for large transfers, at the cost of more load on the server."`
 		Instance                string       `help:"Remote instance name to request; depending on the server this may be required."`
 		Name                    string       `help:"A name for this worker instance. This is attached to artifacts uploaded to remote storage." example:"agent-001"`
 		DisplayURL              string       `help:"A URL to browse the remote server with (e.g. using buildbarn-browser). Only used when printing hashes."`

src/remote/BUILD

@@ -10,6 +10,7 @@ go_library(
         "///third_party/go/cloud.google.com_go_longrunning//autogen/longrunningpb",
         "///third_party/go/github.com_alessio_shellescape//:shellescape",
         "///third_party/go/github.com_bazelbuild_remote-apis-sdks//go/api/command",
+        "///third_party/go/github.com_bazelbuild_remote-apis-sdks//go/pkg/balancer",
         "///third_party/go/github.com_bazelbuild_remote-apis-sdks//go/pkg/client",
         "///third_party/go/github.com_bazelbuild_remote-apis-sdks//go/pkg/command",
         "///third_party/go/github.com_bazelbuild_remote-apis-sdks//go/pkg/digest",

src/remote/remote.go

@@ -15,6 +15,7 @@ import (
 	"time"
 
 	"cloud.google.com/go/longrunning/autogen/longrunningpb"
+	"github.com/bazelbuild/remote-apis-sdks/go/pkg/balancer"
 	"github.com/bazelbuild/remote-apis-sdks/go/pkg/client"
 	"github.com/bazelbuild/remote-apis-sdks/go/pkg/digest"
 	"github.com/bazelbuild/remote-apis-sdks/go/pkg/filemetadata"
@@ -207,13 +208,11 @@ func (c *Client) initExec() error {
 	if err != nil {
 		return err
 	}
-	client, err := client.NewClient(context.Background(), c.instance, client.DialParams{
-		Service:            c.state.Config.Remote.URL,
-		CASService:         c.state.Config.Remote.CASURL,
-		NoSecurity:         !c.state.Config.Remote.Secure,
-		TransportCredsOnly: c.state.Config.Remote.Secure,
-		DialOpts:           dialOpts,
-	}, client.UseBatchOps(true), &client.TreeSymlinkOpts{Preserved: true}, client.RetryTransient(), client.RPCTimeouts(map[string]time.Duration{
+	execConn, casConn, err := c.dialConnections(dialOpts)
+	if err != nil {
+		return err
+	}
+	client, err := client.NewClientFromConnection(context.Background(), c.instance, execConn, casConn, client.UseBatchOps(true), &client.TreeSymlinkOpts{Preserved: true}, client.RetryTransient(), client.RPCTimeouts(map[string]time.Duration{
 		"default":          time.Duration(c.state.Config.Remote.Timeout),
 		"GetCapabilities":  5 * time.Second,
 		"BatchUpdateBlobs": time.Minute,
@@ -284,6 +283,52 @@ func (c *Client) initExec() error {
 	return nil
 }
 
+// dialConnections opens the connections to the execution and CAS servers.
+// The two have quite different traffic so we size them separately; execution is a relatively
+// small number of long-lived but mostly idle streams, whereas CAS has many short-lived
+// RPCs which can benefit from multiple connections.
+func (c *Client) dialConnections(dialOpts []grpc.DialOption) (execConn, casConn client.GrpcClientConn, err error) {
+	ctx := context.Background()
+	opts, _, err := client.OptsFromParams(ctx, client.DialParams{
+		NoSecurity:         !c.state.Config.Remote.Secure,
+		TransportCredsOnly: c.state.Config.Remote.Secure,
+		DialOpts:           dialOpts,
+	})
+	if err != nil {
+		return nil, nil, fmt.Errorf("failed to prepare gRPC dial options: %w", err)
+	}
+	dial := func(url string, n int) (client.GrpcClientConn, error) {
+		pool, err := balancer.NewRRConnPool(ctx, n, func(ctx context.Context) (*grpc.ClientConn, error) {
+			return grpc.NewClient(url, opts...)
+		})
+		if err != nil {
+			return nil, fmt.Errorf("failed to connect to %s: %w", url, err)
+		}
+		return pool, nil
+	}
+	// gRPC won't open more than one HTTP/2 connection per ClientConn, so if the server (or some
+	// proxy in front of it) limits concurrent streams per connection, further Execute calls block
+	// until one finishes. Common limits are 100-128; 64 per connection gives some headroom since
+	// round-robin doesn't balance long-lived streams perfectly.
+	numExecConns := max((c.state.Config.NumRemoteExecutors()+63)/64, 1)
+	numCASConns := max(c.state.Config.Remote.CASConnections, 1)
+	url := c.state.Config.Remote.URL
+	casURL := c.state.Config.Remote.CASURL
+	if casURL == "" || casURL == url {
+		// Both kinds of traffic share these connections, so we need the capacity of both.
+		conn, err := dial(url, numExecConns+numCASConns)
+		return conn, conn, err
+	}
+	if execConn, err = dial(url, numExecConns); err != nil {
+		return nil, nil, err
+	}
+	if casConn, err = dial(casURL, numCASConns); err != nil {
+		execConn.Close()
+		return nil, nil, err
+	}
+	return execConn, casConn, nil
+}
+
 // initFetch initialises the remote fetch server.
 func (c *Client) initFetch() error {
 	dialOpts, err := c.dialOpts()