From 920665ad444cf125ce0dccf5a91cc18746e85d19 Mon Sep 17 00:00:00 2001 From: Robert Lin Date: Tue, 16 Jul 2024 13:38:53 -0700 Subject: [PATCH] chore/lib/telemetrygateway: fixup Dial helper (#63862) Makes it a little easier to use ## Test plan n/a --- lib/telemetrygateway/v1/dial.go | 20 +++++++++++++++----- 1 file changed, 15 insertions(+), 5 deletions(-) diff --git a/lib/telemetrygateway/v1/dial.go b/lib/telemetrygateway/v1/dial.go index 70e771f4ac8..274251f8746 100644 --- a/lib/telemetrygateway/v1/dial.go +++ b/lib/telemetrygateway/v1/dial.go @@ -10,8 +10,9 @@ import ( "google.golang.org/grpc" "google.golang.org/grpc/keepalive" - "github.com/sourcegraph/sourcegraph/lib/errors" "google.golang.org/grpc/credentials/oauth" + + "github.com/sourcegraph/sourcegraph/lib/errors" ) // Experimental: some additional networking options to account for some odd @@ -34,6 +35,12 @@ var cloudRunDialOptions = []grpc.DialOption{ grpc.WithIdleTimeout(1 * time.Minute), } +type SimpleClient struct { + TelemeteryGatewayServiceClient + // Conn can be closed to ensure clean shutdown. + Conn *grpc.ClientConn +} + // Dial establishes a connection to the Telemetry Gateway gRPC service with // the given configuration. The oauth2.TokenSource should provide SAMS credentials, // for example: @@ -63,20 +70,23 @@ var cloudRunDialOptions = []grpc.DialOption{ // Dial is intended for simple, standard production use cases. If you need // to customize the way you connect to Telemetry Gateway, you should create your // own dial setup. -func Dial(ctx context.Context, logger log.Logger, addr *url.URL, ts oauth2.TokenSource, dialOpts ...grpc.DialOption) (*grpc.ClientConn, error) { +func Dial(ctx context.Context, logger log.Logger, addr *url.URL, ts oauth2.TokenSource, dialOpts ...grpc.DialOption) (*SimpleClient, error) { insecureTarget := addr.Scheme != "https" if insecureTarget { return nil, errors.Newf("insecure target Telemetry Gateway %q", addr.String()) } dialOpts = append(dialOpts, grpc.WithPerRPCCredentials(oauth.TokenSource{TokenSource: ts})) dialOpts = append(dialOpts, cloudRunDialOptions...) - logger.Info("dialing Enterprise Portal gRPC service", + logger.Info("dialing Telemetry Gateway gRPC service", log.String("host", addr.Host), log.Bool("insecureTarget", insecureTarget)) //lint:ignore SA1019 DialContext will be supported throughout 1.x conn, err := grpc.NewClient(addr.Host, dialOpts...) if err != nil { - return nil, errors.Wrapf(err, "failed to connect to Enterprise Portal gRPC service at %s", addr.String()) + return nil, errors.Wrapf(err, "failed to connect to Telemetry Gateway gRPC service at %s", addr.String()) } - return conn, nil + return &SimpleClient{ + TelemeteryGatewayServiceClient: NewTelemeteryGatewayServiceClient(conn), + Conn: conn, + }, nil }