fix(connector): add gRPC keepalive and default UNAVAILABLE retry policy

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

@@ -4,13 +4,15 @@ import (
 	"context"
 	"crypto/x509"
 	"strings"
+	"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"
@@ -20,6 +22,42 @@ import (
 
 const defaultTaskTypeVersion = 0
 
+// defaultGRPCServiceConfig is the gRPC service config applied to a connector
+// connection when its Deployment does not set DefaultServiceConfig. It enables
+// round-robin load balancing across connector replicas and retries transient
+// UNAVAILABLE failures so a single dropped connection does not fail the task.
+//
+// Connector traffic frequently flows through an L7 gateway (e.g. the
+// Knative/Kourier Envoy gateway), which reaps idle HTTP/2 connections and sends
+// GOAWAY/drain on routing changes. The cached *grpc.ClientConn then surfaces
+// "error reading server preface: ... use of closed network connection" on the
+// next RPC; retrying UNAVAILABLE absorbs that blip in-band. retryThrottling
+// stops retries from hammering a connector that is genuinely down.
+const defaultGRPCServiceConfig = `{
+  "loadBalancingConfig": [{"round_robin":{}}],
+  "methodConfig": [{
+    "name": [{"service": "flyteidl2.connector.AsyncConnectorService"}],
+    "retryPolicy": {
+      "maxAttempts": 4,
+      "initialBackoff": "0.2s",
+      "maxBackoff": "3s",
+      "backoffMultiplier": 2.0,
+      "retryableStatusCodes": ["UNAVAILABLE"]
+    }
+  }],
+  "retryThrottling": {"maxTokens": 10, "tokenRatio": 0.1}
+}`
+
+// gRPC client keepalive parameters. Pinging below the gateway's idle timeout
+// keeps long-lived connector connections from being reaped while idle (the
+// connector metadata poll and task RPCs are bursty), and proactively detects a
+// half-dead connection instead of discovering it on the next RPC.
+var connectorKeepalive = keepalive.ClientParameters{
+	Time:                30 * time.Second,
+	Timeout:             10 * time.Second,
+	PermitWithoutStream: true,
+}
+
 type Connector struct {
 	// ConnectorDeployment is the connector deployment where this connector is running.
 	ConnectorDeployment *Deployment
@@ -50,9 +88,13 @@ func getGrpcConnection(ctx context.Context, connector *Deployment) (*grpc.Client
 		opts = append(opts, grpc.WithTransportCredentials(creds))
 	}
 
-	if len(connector.DefaultServiceConfig) != 0 {
-		opts = append(opts, grpc.WithDefaultServiceConfig(connector.DefaultServiceConfig))
+	serviceConfig := connector.DefaultServiceConfig
+	if len(serviceConfig) == 0 {
+		serviceConfig = defaultGRPCServiceConfig
 	}
+	opts = append(opts, grpc.WithDefaultServiceConfig(serviceConfig))
+
+	opts = append(opts, grpc.WithKeepaliveParams(connectorKeepalive))
 
 	var err error
 	conn, err := grpc.Dial(connector.Endpoint, opts...) //nolint: staticcheck

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

@@ -2,9 +2,12 @@ package connector
 
 import (
 	"context"
+	"encoding/json"
+	"strings"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
 func TestInitializeClients(t *testing.T) {
@@ -26,3 +29,59 @@ func TestInitializeClients(t *testing.T) {
 	_, ok = cs.asyncConnectorClients["x"]
 	assert.True(t, ok)
 }
+
+func TestDefaultGRPCServiceConfig(t *testing.T) {
+	// Must be valid JSON — grpc.WithDefaultServiceConfig silently ignores a
+	// malformed config, so a typo here would disable LB/retry without any error.
+	assert.True(t, json.Valid([]byte(defaultGRPCServiceConfig)), "defaultGRPCServiceConfig must be valid JSON")
+
+	var parsed struct {
+		LoadBalancingConfig []map[string]any `json:"loadBalancingConfig"`
+		MethodConfig        []struct {
+			Name []struct {
+				Service string `json:"service"`
+			} `json:"name"`
+			RetryPolicy struct {
+				MaxAttempts          int      `json:"maxAttempts"`
+				RetryableStatusCodes []string `json:"retryableStatusCodes"`
+			} `json:"retryPolicy"`
+		} `json:"methodConfig"`
+		RetryThrottling map[string]any `json:"retryThrottling"`
+	}
+	require.NoError(t, json.Unmarshal([]byte(defaultGRPCServiceConfig), &parsed))
+
+	require.Len(t, parsed.LoadBalancingConfig, 1)
+	_, hasRoundRobin := parsed.LoadBalancingConfig[0]["round_robin"]
+	assert.True(t, hasRoundRobin, "expected round_robin load balancing")
+
+	require.Len(t, parsed.MethodConfig, 1)
+	mc := parsed.MethodConfig[0]
+	require.Len(t, mc.Name, 1)
+	assert.Equal(t, "flyteidl2.connector.AsyncConnectorService", mc.Name[0].Service)
+	assert.Greater(t, mc.RetryPolicy.MaxAttempts, 1)
+	assert.Equal(t, []string{"UNAVAILABLE"}, mc.RetryPolicy.RetryableStatusCodes)
+	assert.NotEmpty(t, parsed.RetryThrottling, "expected retryThrottling to bound retry storms")
+}
+
+func TestGetGrpcConnection(t *testing.T) {
+	ctx := context.Background()
+
+	// Empty DefaultServiceConfig must fall back to defaultGRPCServiceConfig
+	// (round-robin + retry) rather than no service config at all.
+	conn, err := getGrpcConnection(ctx, &Deployment{Endpoint: "x", Insecure: true})
+	require.NoError(t, err)
+	require.NotNil(t, conn)
+	assert.NoError(t, conn.Close())
+
+	// A deployment-specific DefaultServiceConfig still takes precedence.
+	custom := `{"loadBalancingConfig": [{"round_robin":{}}]}`
+	conn, err = getGrpcConnection(ctx, &Deployment{Endpoint: "y", Insecure: true, DefaultServiceConfig: custom})
+	require.NoError(t, err)
+	require.NotNil(t, conn)
+	assert.NoError(t, conn.Close())
+}
+
+// ensure the constant referenced from config.go and the LB choice stay in sync
+func TestDefaultGRPCServiceConfigMentionsRoundRobin(t *testing.T) {
+	assert.True(t, strings.Contains(defaultGRPCServiceConfig, "round_robin"))
+}

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

@@ -41,10 +41,13 @@ var (
 			},
 		},
 		DefaultConnector: Deployment{
-			Endpoint:             "",
-			Insecure:             true,
-			DefaultTimeout:       config.Duration{Duration: 10 * time.Second},
-			DefaultServiceConfig: `{"loadBalancingConfig": [{"round_robin":{}}]}`,
+			Endpoint:       "",
+			Insecure:       true,
+			DefaultTimeout: config.Duration{Duration: 10 * time.Second},
+			// DefaultServiceConfig is left empty so getGrpcConnection falls back to
+			// defaultGRPCServiceConfig (round-robin LB + UNAVAILABLE retry). Set this
+			// per-deployment only to override that default.
+			DefaultServiceConfig: "",
 		},
 		ConnectorDeployments:  map[string]*Deployment{},
 		ConnectorForTaskTypes: map[string]string{},

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

@@ -53,7 +53,9 @@ func TestDefaultAgentConfig(t *testing.T) {
 	assert.Equal(t, "", cfg.DefaultConnector.Endpoint)
 	assert.True(t, cfg.DefaultConnector.Insecure)
 	assert.Equal(t, 10*time.Second, cfg.DefaultConnector.DefaultTimeout.Duration)
-	assert.Equal(t, `{"loadBalancingConfig": [{"round_robin":{}}]}`, cfg.DefaultConnector.DefaultServiceConfig)
+	// DefaultServiceConfig defaults to empty; getGrpcConnection falls back to
+	// defaultGRPCServiceConfig when a deployment does not override it.
+	assert.Empty(t, cfg.DefaultConnector.DefaultServiceConfig)
 
 	assert.Empty(t, cfg.DefaultConnector.Timeouts)