fix(clients): cache GetPublicClientConfig response to avoid redundant API calls

flyteidl/clients/go/admin/cached_metadata_client.go

@@ -0,0 +1,46 @@
+package admin
+
+import (
+	"context"
+	"sync"
+
+	"google.golang.org/grpc"
+
+	"github.com/flyteorg/flyte/flyteidl/gen/pb-go/flyteidl/service"
+)
+
+// cachedMetadataClient wraps an AuthMetadataServiceClient and caches the
+// GetPublicClientConfig response. The server returns static configuration
+// that does not change at runtime, so a single successful fetch is enough.
+// Errors are not cached: a failed call (for example a transient RPC error or
+// a cancelled context) is retried on the next call.
+type cachedMetadataClient struct {
+	service.AuthMetadataServiceClient
+	mu     sync.Mutex
+	cached *service.PublicClientAuthConfigResponse
+}
+
+// NewCachedMetadataClient returns an AuthMetadataServiceClient that calls
+// GetPublicClientConfig until it succeeds once and returns the cached
+// response thereafter.
+func NewCachedMetadataClient(inner service.AuthMetadataServiceClient) service.AuthMetadataServiceClient {
+	return &cachedMetadataClient{AuthMetadataServiceClient: inner}
+}
+
+func (c *cachedMetadataClient) GetPublicClientConfig(
+	ctx context.Context,
+	in *service.PublicClientAuthConfigRequest,
+	opts ...grpc.CallOption,
+) (*service.PublicClientAuthConfigResponse, error) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	if c.cached != nil {
+		return c.cached, nil
+	}
+	resp, err := c.AuthMetadataServiceClient.GetPublicClientConfig(ctx, in, opts...)
+	if err != nil {
+		return nil, err
+	}
+	c.cached = resp
+	return resp, nil
+}

flyteidl/clients/go/admin/cached_metadata_client_test.go

@@ -0,0 +1,76 @@
+package admin
+
+import (
+	"context"
+	"errors"
+	"sync"
+	"sync/atomic"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"google.golang.org/grpc"
+
+	"github.com/flyteorg/flyte/flyteidl/gen/pb-go/flyteidl/service"
+)
+
+type countingMetadataClient struct {
+	service.AuthMetadataServiceClient
+	calls atomic.Int32
+	errs  []error
+	resp  *service.PublicClientAuthConfigResponse
+}
+
+func (c *countingMetadataClient) GetPublicClientConfig(
+	_ context.Context,
+	_ *service.PublicClientAuthConfigRequest,
+	_ ...grpc.CallOption,
+) (*service.PublicClientAuthConfigResponse, error) {
+	n := int(c.calls.Add(1))
+	if n <= len(c.errs) {
+		return nil, c.errs[n-1]
+	}
+	return c.resp, nil
+}
+
+func TestCachedMetadataClient_CallsInnerOnceForConcurrentCallers(t *testing.T) {
+	inner := &countingMetadataClient{resp: &service.PublicClientAuthConfigResponse{ClientId: "flytectl"}}
+	client := NewCachedMetadataClient(inner)
+
+	var wg sync.WaitGroup
+	for i := 0; i < 50; i++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			resp, err := client.GetPublicClientConfig(context.Background(), &service.PublicClientAuthConfigRequest{})
+			assert.NoError(t, err)
+			assert.Equal(t, "flytectl", resp.GetClientId())
+		}()
+	}
+	wg.Wait()
+
+	assert.Equal(t, int32(1), inner.calls.Load())
+}
+
+func TestCachedMetadataClient_RetriesAfterError(t *testing.T) {
+	inner := &countingMetadataClient{
+		errs: []error{errors.New("unavailable"), context.DeadlineExceeded},
+		resp: &service.PublicClientAuthConfigResponse{ClientId: "flytectl"},
+	}
+	client := NewCachedMetadataClient(inner)
+	ctx := context.Background()
+	req := &service.PublicClientAuthConfigRequest{}
+
+	_, err := client.GetPublicClientConfig(ctx, req)
+	assert.Error(t, err)
+	_, err = client.GetPublicClientConfig(ctx, req)
+	assert.ErrorIs(t, err, context.DeadlineExceeded)
+
+	resp, err := client.GetPublicClientConfig(ctx, req)
+	assert.NoError(t, err)
+	assert.Equal(t, "flytectl", resp.GetClientId())
+
+	resp, err = client.GetPublicClientConfig(ctx, req)
+	assert.NoError(t, err)
+	assert.Equal(t, "flytectl", resp.GetClientId())
+	assert.Equal(t, int32(3), inner.calls.Load())
+}

flyteidl/clients/go/admin/client.go

@@ -121,7 +121,7 @@ func InitializeAuthMetadataClient(ctx context.Context, cfg *Config, proxyCredent
 		return nil, fmt.Errorf("failed to initialized admin connection. Error: %w", err)
 	}
 
-	return service.NewAuthMetadataServiceClient(authMetadataConnection), nil
+	return NewCachedMetadataClient(service.NewAuthMetadataServiceClient(authMetadataConnection)), nil
 }
 
 func NewAdminConnection(ctx context.Context, cfg *Config, proxyCredentialsFuture *PerRPCCredentialsFuture, opts ...grpc.DialOption) (*grpc.ClientConn, error) {