go.chromium.org/luci@v0.0.0-20240309015107-7cdc2e660f33/swarming/server/rbe/dial.go (about)

     1  // Copyright 2023 The LUCI Authors.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //      http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package rbe
    16  
    17  import (
    18  	"context"
    19  	"fmt"
    20  
    21  	"go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc"
    22  	"google.golang.org/grpc"
    23  	"google.golang.org/grpc/credentials"
    24  
    25  	"go.chromium.org/luci/common/errors"
    26  	"go.chromium.org/luci/common/logging"
    27  	"go.chromium.org/luci/grpc/grpcmon"
    28  	"go.chromium.org/luci/server/auth"
    29  
    30  	"go.chromium.org/luci/swarming/internal/remoteworkers"
    31  )
    32  
    33  // Instructs the gRPC client how to retry.
    34  var retryPolicy = fmt.Sprintf(`{
    35    "methodConfig": [{
    36      "name": [
    37        {"service": "%s"},
    38        {"service": "%s"}
    39      ],
    40      "waitForReady": true,
    41      "retryPolicy": {
    42        "MaxAttempts": 5,
    43        "InitialBackoff": "0.01s",
    44        "MaxBackoff": "1s",
    45        "BackoffMultiplier": 2.0,
    46        "RetryableStatusCodes": [
    47          "INTERNAL",
    48          "UNAVAILABLE"
    49        ]
    50      }
    51    }]
    52  }`,
    53  	remoteworkers.Bots_ServiceDesc.ServiceName,
    54  	remoteworkers.Reservations_ServiceDesc.ServiceName,
    55  )
    56  
    57  // Dial dials RBE backend with proper authentication.
    58  //
    59  // Returns multiple identical clients each representing a separate HTTP2
    60  // connection.
    61  func Dial(ctx context.Context, count int) ([]grpc.ClientConnInterface, error) {
    62  	creds, err := auth.GetPerRPCCredentials(ctx,
    63  		auth.AsSelf,
    64  		auth.WithScopes(auth.CloudOAuthScopes...),
    65  	)
    66  	if err != nil {
    67  		return nil, errors.Annotate(err, "failed to get credentials").Err()
    68  	}
    69  	logging.Infof(ctx, "Dialing %d RBE backend connections...", count)
    70  	conns := make([]grpc.ClientConnInterface, count)
    71  	for i := 0; i < count; i++ {
    72  		conns[i], err = grpc.DialContext(ctx, "remotebuildexecution.googleapis.com:443",
    73  			grpc.WithTransportCredentials(credentials.NewTLS(nil)),
    74  			grpc.WithPerRPCCredentials(creds),
    75  			grpc.WithBlock(),
    76  			grpc.WithStatsHandler(&grpcmon.ClientRPCStatsMonitor{}),
    77  			grpc.WithUnaryInterceptor(otelgrpc.UnaryClientInterceptor()),
    78  			grpc.WithStreamInterceptor(otelgrpc.StreamClientInterceptor()),
    79  			grpc.WithDefaultServiceConfig(retryPolicy),
    80  		)
    81  		if err != nil {
    82  			return nil, errors.Annotate(err, "failed to dial RBE backend").Err()
    83  		}
    84  	}
    85  	return conns, nil
    86  }