xds: generic xds client resource watching e2e

xds/internal/clients/config.go

@@ -59,6 +59,16 @@ type ServerIdentifier struct {
 	//
 	// For example, a custom TransportBuilder might use this field to
 	// configure a specific security credentials.
+	//
+	// Extensions may be any type that is comparable, as they are used as map
+	// keys internally. If Extensions are not able to be used as a map key,
+	// the client may panic.
+	//
+	// See: https://go.dev/ref/spec#Comparison_operators
+	//
+	// Any equivalent extensions in all ServerIdentifiers present in a single
+	// client's configuration should have the same value. Not following this
+	// restriction may result in excess resource usage.
 	Extensions any
 }
 

xds/internal/clients/grpctransport/examples_test.go

@@ -0,0 +1,39 @@
+/*
+ *
+ * Copyright 2025 gRPC authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package grpctransport_test
+
+import (
+	"fmt"
+
+	"google.golang.org/grpc/xds/internal/clients"
+	"google.golang.org/grpc/xds/internal/clients/grpctransport"
+)
+
+// ExampleServerIdentifierExtension demonstrates how to create
+// clients.ServerIdentifier with grpctransport.ServerIdentifierExtension as
+// its extensions.
+//
+// This example is creating clients.ServerIdentifier to connect to server at
+// localhost:5678 using the credentials named "local". Note that "local" must
+// exist as an entry in the provided credentials to grpctransport.Builder.
+func ExampleServerIdentifierExtension() {
+	// Note the Extensions field is set by value and not by pointer.
+	fmt.Printf("%+v", clients.ServerIdentifier{ServerURI: "localhost:5678", Extensions: grpctransport.ServerIdentifierExtension{Credentials: "local"}})
+	// Output: {ServerURI:localhost:5678 Extensions:{Credentials:local}}
+}

xds/internal/clients/grpctransport/grpc_transport.go

@@ -23,25 +23,55 @@ package grpctransport
 import (
 	"context"
 	"fmt"
+	"sync"
 	"time"
 
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/credentials"
+	"google.golang.org/grpc/grpclog"
 	"google.golang.org/grpc/keepalive"
 	"google.golang.org/grpc/xds/internal/clients"
 )
 
+var (
+	logger = grpclog.Component("grpctransport")
+)
+
 // ServerIdentifierExtension holds settings for connecting to a gRPC server,
 // such as an xDS management or an LRS server.
+//
+// It must be set by value (not pointer) in the
+// clients.ServerIdentifier.Extensions field (See Example).
 type ServerIdentifierExtension struct {
-	// Credentials will be used for all gRPC transports. If it is unset,
-	// transport creation will fail.
-	Credentials credentials.Bundle
+	// Credentials is name of the credentials to use for this transport to the
+	// server. It must be present in the map passed to NewBuilder.
+	Credentials string
 }
 
 // Builder creates gRPC-based Transports. It must be paired with ServerIdentifiers
 // that contain an Extension field of type ServerIdentifierExtension.
-type Builder struct{}
+type Builder struct {
+	// credentials is a map of credentials names to credentials.Bundle which
+	// can be used to connect to the server.
+	credentials map[string]credentials.Bundle
+
+	mu sync.Mutex
+	// connections is a map of clients.ServerIdentifiers in use by the Builder
+	// to connect to different servers.
+	connections map[clients.ServerIdentifier]*grpc.ClientConn
+	refs        map[clients.ServerIdentifier]int
+}
+
+// NewBuilder provides a builder for creating gRPC-based Transports using
+// the credentials from provided map of credentials names to
+// credentials.Bundle.
+func NewBuilder(credentials map[string]credentials.Bundle) *Builder {
+	return &Builder{
+		credentials: credentials,
+		connections: make(map[clients.ServerIdentifier]*grpc.ClientConn),
+		refs:        make(map[clients.ServerIdentifier]int),
+	}
+}
 
 // Build returns a gRPC-based clients.Transport.
 //
@@ -57,14 +87,24 @@ func (b *Builder) Build(si clients.ServerIdentifier) (clients.Transport, error)
 	if !ok {
 		return nil, fmt.Errorf("grpctransport: Extensions field is %T, but must be %T in ServerIdentifier", si.Extensions, ServerIdentifierExtension{})
 	}
-	if sce.Credentials == nil {
-		return nil, fmt.Errorf("grptransport: Credentials field is not set in ServerIdentifierExtension")
+
+	creds, ok := b.credentials[sce.Credentials]
+	if !ok {
+		return nil, fmt.Errorf("grpctransport: unknown credentials type %q specified in extensions", sce.Credentials)
 	}
 
-	// TODO: Incorporate reference count map for existing transports and
-	// deduplicate transports based on the provided ServerIdentifier so that
-	// transport channel to same server can be shared between xDS and LRS
-	// client.
+	b.mu.Lock()
+	defer b.mu.Unlock()
+
+	if cc, ok := b.connections[si]; ok {
+		if logger.V(2) {
+			logger.Info("Reusing existing connection to the server for ServerIdentifier: %v", si)
+		}
+		b.refs[si]++
+		tr := &grpcTransport{cc: cc}
+		tr.cleanup = b.cleanupFunc(si, tr)
+		return tr, nil
+	}
 
 	// Create a new gRPC client/channel for the server with the provided
 	// credentials, server URI, and a byte codec to send and receive messages.
@@ -74,16 +114,50 @@ func (b *Builder) Build(si clients.ServerIdentifier) (clients.Transport, error)
 		Time:    5 * time.Minute,
 		Timeout: 20 * time.Second,
 	})
-	cc, err := grpc.NewClient(si.ServerURI, kpCfg, grpc.WithCredentialsBundle(sce.Credentials), grpc.WithDefaultCallOptions(grpc.ForceCodec(&byteCodec{})))
+	cc, err := grpc.NewClient(si.ServerURI, kpCfg, grpc.WithCredentialsBundle(creds), grpc.WithDefaultCallOptions(grpc.ForceCodec(&byteCodec{})))
 	if err != nil {
-		return nil, fmt.Errorf("grpctransport: failed to create transport to server %q: %v", si.ServerURI, err)
+		return nil, fmt.Errorf("grpctransport: failed to create connection to server %q: %v", si.ServerURI, err)
 	}
+	tr := &grpcTransport{cc: cc}
+	// Register a cleanup function that decrements the refs to the gRPC
+	// transport each time Close() is called to close it and remove from
+	// transports and connections map if last reference is being released.
+	tr.cleanup = b.cleanupFunc(si, tr)
+
+	// Add the newly created connection to the maps to re-use the transport
+	// channel and track references.
+	b.connections[si] = cc
+	b.refs[si] = 1
+
+	if logger.V(2) {
+		logger.Info("Created a new transport to the server for ServerIdentifier: %v", si)
+	}
+	return tr, nil
+}
 
-	return &grpcTransport{cc: cc}, nil
+func (b *Builder) cleanupFunc(si clients.ServerIdentifier, tr *grpcTransport) func() {
+	return sync.OnceFunc(func() {
+		b.mu.Lock()
+		defer b.mu.Unlock()
+
+		b.refs[si]--
+		if b.refs[si] != 0 {
+			return
+		}
+
+		tr.cc.Close()
+		tr.cc = nil
+		delete(b.connections, si)
+		delete(b.refs, si)
+	})
 }
 
 type grpcTransport struct {
 	cc *grpc.ClientConn
+
+	// cleanup is the function to be invoked for releasing the references to
+	// the gRPC transport each time Close() is called.
+	cleanup func()
 }
 
 // NewStream creates a new gRPC stream to the server for the specified method.
@@ -96,8 +170,8 @@ func (g *grpcTransport) NewStream(ctx context.Context, method string) (clients.S
 }
 
 // Close closes the gRPC channel to the server.
-func (g *grpcTransport) Close() error {
-	return g.cc.Close()
+func (g *grpcTransport) Close() {
+	g.cleanup()
 }
 
 type stream struct {