fix(connector): reuse persistent connections for ListConnectors polling (#7705)

Signed-off-by: Kevin Su <pingsutw@apache.org>
Signed-off-by: Kevin Su <pingsutw@gmail.com>
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

Author: pingsutw

flyteplugins/go/tasks/plugins/webapi/connector/client.go

@@ -4,13 +4,16 @@ import (
 	"context"
 	"crypto/x509"
 	"strings"
+	"sync"
+	"time"
 
 	"golang.org/x/exp/maps"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/credentials"
 	"google.golang.org/grpc/credentials/insecure"
 	"google.golang.org/grpc/grpclog"
+	"google.golang.org/grpc/keepalive"
 	"google.golang.org/grpc/status"
 
 	"github.com/flyteorg/flyte/v2/flytestdlib/config"
@@ -32,10 +35,66 @@ type Connector struct {
 
 // ClientSet contains the clients exposed to communicate with various connector services.
 type ClientSet struct {
+	mu                       sync.RWMutex
 	asyncConnectorClients    map[string]connector.AsyncConnectorServiceClient    // map[endpoint] => AsyncConnectorServiceClient
 	connectorMetadataClients map[string]connector.ConnectorMetadataServiceClient // map[endpoint] => ConnectorMetadataServiceClient
 }
 
+// getOrDialAsyncClient returns the cached AsyncConnectorService client for the
+// endpoint, dialing (and caching) a persistent connection on first use.
+func (cs *ClientSet) getOrDialAsyncClient(ctx context.Context, deployment *Deployment) (connector.AsyncConnectorServiceClient, error) {
+	cs.mu.Lock()
+	defer cs.mu.Unlock()
+
+	if cs.asyncConnectorClients == nil {
+		cs.asyncConnectorClients = make(map[string]connector.AsyncConnectorServiceClient)
+	}
+	if cs.connectorMetadataClients == nil {
+		cs.connectorMetadataClients = make(map[string]connector.ConnectorMetadataServiceClient)
+	}
+
+	if client, ok := cs.asyncConnectorClients[deployment.Endpoint]; ok {
+		return client, nil
+	}
+
+	conn, err := getGrpcConnection(ctx, deployment)
+	if err != nil {
+		return nil, err
+	}
+
+	asyncClient := connector.NewAsyncConnectorServiceClient(conn)
+	cs.asyncConnectorClients[deployment.Endpoint] = asyncClient
+	if _, ok := cs.connectorMetadataClients[deployment.Endpoint]; !ok {
+		cs.connectorMetadataClients[deployment.Endpoint] = connector.NewConnectorMetadataServiceClient(conn)
+	}
+	return asyncClient, nil
+}
+
+// getOrDialMetadataClient returns the cached ConnectorMetadataService client for
+// the endpoint, dialing (and caching) a persistent connection on first use.
+func (cs *ClientSet) getOrDialMetadataClient(ctx context.Context, deployment *Deployment) (connector.ConnectorMetadataServiceClient, error) {
+	cs.mu.Lock()
+	defer cs.mu.Unlock()
+	if client, ok := cs.connectorMetadataClients[deployment.Endpoint]; ok {
+		return client, nil
+	}
+	conn, err := getGrpcConnection(ctx, deployment)
+	if err != nil {
+		return nil, err
+	}
+	client := connector.NewConnectorMetadataServiceClient(conn)
+	cs.connectorMetadataClients[deployment.Endpoint] = client
+	return client, nil
+}
+
+// metadataClient returns the cached ConnectorMetadataService client for the endpoint.
+func (cs *ClientSet) metadataClient(endpoint string) (connector.ConnectorMetadataServiceClient, bool) {
+	cs.mu.RLock()
+	defer cs.mu.RUnlock()
+	client, ok := cs.connectorMetadataClients[endpoint]
+	return client, ok
+}
+
 func getGrpcConnection(ctx context.Context, connector *Deployment) (*grpc.ClientConn, error) {
 	var opts []grpc.DialOption
 
@@ -59,6 +118,13 @@ func getGrpcConnection(ctx context.Context, connector *Deployment) (*grpc.Client
 	opts = append(opts,
 		grpc.WithChainUnaryInterceptor(clientMetrics.UnaryClientInterceptor()),
 		grpc.WithChainStreamInterceptor(clientMetrics.StreamClientInterceptor()),
+		// Keepalive lets gRPC notice a dead/half-open transport within seconds
+		// (via HTTP/2 PINGs during active RPCs) instead of relying on the OS TCP
+		// timeout.
+		grpc.WithKeepaliveParams(keepalive.ClientParameters{
+			Time:    30 * time.Second,
+			Timeout: 10 * time.Second,
+		}),
 	)
 
 	var err error
@@ -102,7 +168,7 @@ func updateRegistry(
 	isConnectorApp bool,
 ) {
 	for connectorID, connectorDeployment := range connectorDeployments {
-		client, ok := cs.connectorMetadataClients[connectorDeployment.Endpoint]
+		client, ok := cs.metadataClient(connectorDeployment.Endpoint)
 		if !ok {
 			logger.Warningf(ctx, "Connector client not found in the clientSet for the endpoint: %v", connectorDeployment.Endpoint)
 			continue
@@ -199,28 +265,29 @@ func getConnectorRegistry(ctx context.Context, cs *ClientSet) Registry {
 	return newConnectorRegistry
 }
 
-func getConnectorClientSets(ctx context.Context) *ClientSet {
-	clientSet := &ClientSet{
-		asyncConnectorClients:    make(map[string]connector.AsyncConnectorServiceClient),
-		connectorMetadataClients: make(map[string]connector.ConnectorMetadataServiceClient),
-	}
-
+// allConnectorDeployments returns every configured connector endpoint: the
+// default connector, explicit deployments, and connector apps.
+func allConnectorDeployments(cfg *Config) []*Deployment {
 	var connectorDeployments []*Deployment
-	cfg := GetConfig()
-
 	if len(cfg.DefaultConnector.Endpoint) != 0 {
 		connectorDeployments = append(connectorDeployments, &cfg.DefaultConnector)
 	}
-
 	for _, deployment := range cfg.ConnectorDeployments {
 		connectorDeployments = append(connectorDeployments, deployment)
 	}
-
 	for _, deployment := range cfg.ConnectorApps {
 		connectorDeployments = append(connectorDeployments, deployment)
 	}
+	return connectorDeployments
+}
+
+func getConnectorClientSets(ctx context.Context) *ClientSet {
+	clientSet := &ClientSet{
+		asyncConnectorClients:    make(map[string]connector.AsyncConnectorServiceClient),
+		connectorMetadataClients: make(map[string]connector.ConnectorMetadataServiceClient),
+	}
 
-	for _, connectorDeployment := range connectorDeployments {
+	for _, connectorDeployment := range allConnectorDeployments(GetConfig()) {
 		if _, ok := clientSet.connectorMetadataClients[connectorDeployment.Endpoint]; ok {
 			logger.Infof(ctx, "Connector client already initialized for [%v]", connectorDeployment.Endpoint)
 			continue

flyteplugins/go/tasks/plugins/webapi/connector/client_test.go

@@ -5,6 +5,10 @@ import (
 	"testing"
 
 	"github.com/stretchr/testify/assert"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/credentials/insecure"
+
+	connectorpb "github.com/flyteorg/flyte/v2/gen/go/flyteidl2/connector"
 )
 
 func TestInitializeClients(t *testing.T) {
@@ -26,3 +30,45 @@ func TestInitializeClients(t *testing.T) {
 	_, ok = cs.asyncConnectorClients["x"]
 	assert.True(t, ok)
 }
+
+func TestAllConnectorDeployments(t *testing.T) {
+	cfg := defaultConfig
+	cfg.DefaultConnector = Deployment{Endpoint: "default"}
+	cfg.ConnectorDeployments = map[string]*Deployment{"x": {Endpoint: "x"}}
+	cfg.ConnectorApps = map[string]*Deployment{"app": {Endpoint: "app"}}
+
+	endpoints := make([]string, 0)
+	for _, d := range allConnectorDeployments(&cfg) {
+		endpoints = append(endpoints, d.Endpoint)
+	}
+	assert.ElementsMatch(t, []string{"default", "x", "app"}, endpoints)
+}
+
+func TestGetOrDialMetadataClientReusesConnection(t *testing.T) {
+	// Cancelable context so getGrpcConnection's close goroutine is cleaned up
+	// when the test finishes.
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	// A pre-existing cached client for endpoint "ep".
+	conn, err := grpc.NewClient("passthrough:///ep", grpc.WithTransportCredentials(insecure.NewCredentials()))
+	assert.NoError(t, err)
+	t.Cleanup(func() { _ = conn.Close() })
+	existing := connectorpb.NewConnectorMetadataServiceClient(conn)
+
+	cs := &ClientSet{
+		asyncConnectorClients:    map[string]connectorpb.AsyncConnectorServiceClient{},
+		connectorMetadataClients: map[string]connectorpb.ConnectorMetadataServiceClient{"ep": existing},
+	}
+
+	// Already cached -> must reuse the same client, not re-dial/replace it.
+	got, err := cs.getOrDialMetadataClient(ctx, &Deployment{Endpoint: "ep"})
+	assert.NoError(t, err)
+	assert.True(t, existing == got, "cached connection should be reused, not re-dialed")
+
+	// New endpoint -> dialed and cached. The passthrough scheme avoids DNS.
+	_, err = cs.getOrDialMetadataClient(ctx, &Deployment{Endpoint: "passthrough:///new", Insecure: true})
+	assert.NoError(t, err)
+	_, ok := cs.metadataClient("passthrough:///new")
+	assert.True(t, ok, "new endpoint should be dialed and cached")
+}

flyteplugins/go/tasks/plugins/webapi/connector/plugin.go

@@ -374,24 +374,17 @@ func (p *Plugin) Status(ctx context.Context, taskCtx webapi.StatusContext) (phas
 }
 
 func (p *Plugin) getAsyncConnectorClient(ctx context.Context, connector *Deployment) (connectorPb.AsyncConnectorServiceClient, error) {
-	client, ok := p.cs.asyncConnectorClients[connector.Endpoint]
-	if !ok {
-		conn, err := getGrpcConnection(ctx, connector)
-		if err != nil {
-			return nil, err
-		}
-		client = connectorPb.NewAsyncConnectorServiceClient(conn)
-		p.cs.asyncConnectorClients[connector.Endpoint] = client
-	}
-	return client, nil
+	return p.cs.getOrDialAsyncClient(ctx, connector)
 }
 
 func (p *Plugin) watchConnectors(ctx context.Context, connectorService *ConnectorService) {
 	go wait.Until(func() {
-		childCtx, cancel := context.WithCancel(ctx)
-		defer cancel()
-		clientSet := getConnectorClientSets(childCtx)
-		connectorRegistry := getConnectorRegistry(childCtx, clientSet)
+		for _, deployment := range allConnectorDeployments(GetConfig()) {
+			if _, err := p.cs.getOrDialMetadataClient(ctx, deployment); err != nil {
+				logger.Errorf(ctx, "failed to connect to connector [%v]: %v", deployment.Endpoint, err)
+			}
+		}
+		connectorRegistry := getConnectorRegistry(ctx, p.cs)
 		p.setRegistry(connectorRegistry)
 		connectorService.SetSupportedTaskType(connectorRegistry.getSupportedTaskTypes())
 	}, p.cfg.PollInterval.Duration, ctx.Done())