From a05a8e1e9eeced76018788736b9d49ec4c7f4a9d Mon Sep 17 00:00:00 2001 From: Honnix Date: Thu, 2 Jun 2022 20:01:03 +0200 Subject: [PATCH] Grpc default service config (#301) * add load balancer config Signed-off-by: Babis Kiosidis * add policies doc link Signed-off-by: Babis Kiosidis * add available load balancing policies comment Signed-off-by: Babis Kiosidis * add descriptive comment about missing balancer value Signed-off-by: Babis Kiosidis * describe load balancing policy behaviour Signed-off-by: Babis Kiosidis * import balancers Signed-off-by: Babis Kiosidis * skip linting for empty imports Signed-off-by: Babis Kiosidis * lint nolint Signed-off-by: Babis Kiosidis * rely on grpc client functionality for the configs and remove todo Signed-off-by: Babis Kiosidis * dont repeat the comment Signed-off-by: Babis Kiosidis * change load balancing config to serviceConfig Signed-off-by: Babis Kiosidis * Change config name and remove package preloading Signed-off-by: Hongxin Liang * Test it Signed-off-by: Hongxin Liang Co-authored-by: Babis Kiosidis Signed-off-by: Eduardo Apolinario --- flyteidl/clients/go/admin/client.go | 4 +++ flyteidl/clients/go/admin/client_test.go | 37 ++++++++++++++++++++++++ flyteidl/clients/go/admin/config.go | 11 +++++-- 3 files changed, 49 insertions(+), 3 deletions(-) diff --git a/flyteidl/clients/go/admin/client.go b/flyteidl/clients/go/admin/client.go index 05cfd907fa5..95bbfdaf922 100644 --- a/flyteidl/clients/go/admin/client.go +++ b/flyteidl/clients/go/admin/client.go @@ -192,6 +192,10 @@ func initializeClients(ctx context.Context, cfg *Config, tokenCache pkce.TokenCa opts = append(opts, authOpt) } + if cfg.DefaultServiceConfig != "" { + opts = append(opts, grpc.WithDefaultServiceConfig(cfg.DefaultServiceConfig)) + } + adminConnection, err := NewAdminConnection(ctx, cfg, opts...) if err != nil { logger.Panicf(ctx, "failed to initialize Admin connection. Err: %s", err.Error()) diff --git a/flyteidl/clients/go/admin/client_test.go b/flyteidl/clients/go/admin/client_test.go index b063e20e42f..7e66079acda 100644 --- a/flyteidl/clients/go/admin/client_test.go +++ b/flyteidl/clients/go/admin/client_test.go @@ -20,6 +20,7 @@ import ( "github.com/stretchr/testify/assert" "github.com/stretchr/testify/mock" "golang.org/x/oauth2" + _ "google.golang.org/grpc/balancer/roundrobin" //nolint ) func TestInitializeAndGetAdminClient(t *testing.T) { @@ -262,6 +263,42 @@ func Test_getPkceAuthTokenSource(t *testing.T) { }) } +func TestGetDefaultServiceConfig(t *testing.T) { + u, _ := url.Parse("localhost:8089") + adminServiceConfig := &Config{ + Endpoint: config.URL{URL: *u}, + DefaultServiceConfig: `{"loadBalancingConfig": [{"round_robin":{}}]}`, + } + + assert.NoError(t, SetConfig(adminServiceConfig)) + + ctx := context.Background() + t.Run("legal", func(t *testing.T) { + u, err := url.Parse("http://localhost:8089") + assert.NoError(t, err) + clientSet, err := ClientSetBuilder().WithConfig(&Config{Endpoint: config.URL{URL: *u}, DefaultServiceConfig: `{"loadBalancingConfig": [{"round_robin":{}}]}`}).Build(ctx) + assert.NoError(t, err) + assert.NotNil(t, clientSet) + assert.NotNil(t, clientSet.AdminClient()) + assert.NotNil(t, clientSet.AuthMetadataClient()) + assert.NotNil(t, clientSet.IdentityClient()) + assert.NotNil(t, clientSet.HealthServiceClient()) + }) + t.Run("illegal default service config", func(t *testing.T) { + defer func() { + if r := recover(); r == nil { + t.Errorf("The code did not panic") + } + }() + + u, err := url.Parse("http://localhost:8089") + assert.NoError(t, err) + clientSet, err := ClientSetBuilder().WithConfig(&Config{Endpoint: config.URL{URL: *u}, DefaultServiceConfig: `{"loadBalancingConfig": [{"foo":{}}]}`}).Build(ctx) + assert.Error(t, err) + assert.Nil(t, clientSet) + }) +} + func ExampleClientSetBuilder() { ctx := context.Background() // Create a client set that initializes the connection with flyte admin and sets up Auth (if needed). diff --git a/flyteidl/clients/go/admin/config.go b/flyteidl/clients/go/admin/config.go index dd1e1bce30e..511f94bf69e 100644 --- a/flyteidl/clients/go/admin/config.go +++ b/flyteidl/clients/go/admin/config.go @@ -1,7 +1,6 @@ // Initializes an Admin Client that exposes all implemented services by FlyteAdmin server. The library supports different -// authentication flows (see AuthType). It initializes the grpc connection once and reuses it. The gRPC connection is -// sticky (it hogs one server and keeps the connection alive). For better load balancing against the server, place a -// proxy service in between instead. +// authentication flows (see AuthType). It initializes the grpc connection once and reuses it. A grpc load balancing policy +// can be configured as well. package admin import ( @@ -68,6 +67,12 @@ type Config struct { PkceConfig pkce.Config `json:"pkceConfig" pflag:",Config for Pkce authentication flow."` Command []string `json:"command" pflag:",Command for external authentication token generation"` + + // Set the gRPC service config formatted as a json string https://github.com/grpc/grpc/blob/master/doc/service_config.md + // eg. {"loadBalancingConfig": [{"round_robin":{}}], "methodConfig": [{"name":[{"service": "foo", "method": "bar"}, {"service": "baz"}], "timeout": "1.000000001s"}]} + // find the full schema here https://github.com/grpc/grpc-proto/blob/master/grpc/service_config/service_config.proto#L625 + // Note that required packages may need to be preloaded to support certain service config. For example "google.golang.org/grpc/balancer/roundrobin" should be preloaded to have round-robin policy supported. + DefaultServiceConfig string `json:"defaultServiceConfig" pdflag:",Set the default service config for the admin gRPC client"` } var (