add options for using custom heal/reetry clients for registry

Signed-off-by: denis-tingaikin <[email protected]>
networkservicemesh · Dec 18, 2024 · a6546f5 · a6546f5
1 parent 1e611de
commit a6546f5
Show file tree

Hide file tree

Showing 3 changed files with 39 additions and 6 deletions.
diff --git a/pkg/registry/chains/client/ns_client.go b/pkg/registry/chains/client/ns_client.go
@@ -43,6 +43,8 @@ func NewNetworkServiceRegistryClient(ctx context.Context, opts ...Option) regist
 		nsClientURLResolver:       null.NewNetworkServiceRegistryClient(),
 		authorizeNSRegistryClient: authorize.NewNetworkServiceRegistryClient(authorize.Any()),
 		dialTimeout:               time.Millisecond * 300,
+		nsHealClient:              heal.NewNetworkServiceRegistryClient(ctx),
+		nsRetryClient:             retry.NewNetworkServiceRegistryClient(ctx),
 	}
 	for _, opt := range opts {
 		opt(clientOpts)
@@ -53,9 +55,9 @@ func NewNetworkServiceRegistryClient(ctx context.Context, opts ...Option) regist
 			[]registry.NetworkServiceRegistryClient{
 				begin.NewNetworkServiceRegistryClient(),
 				metadata.NewNetworkServiceClient(),
-				retry.NewNetworkServiceRegistryClient(ctx),
+				clientOpts.nsRetryClient,
 				clientOpts.authorizeNSRegistryClient,
-				heal.NewNetworkServiceRegistryClient(ctx),
+				clientOpts.nsHealClient,
 				clientOpts.nsClientURLResolver,
 				clientconn.NewNetworkServiceRegistryClient(),
 				grpcmetadata.NewNetworkServiceRegistryClient(),

diff --git a/pkg/registry/chains/client/nse_client.go b/pkg/registry/chains/client/nse_client.go
@@ -45,6 +45,8 @@ func NewNetworkServiceEndpointRegistryClient(ctx context.Context, opts ...Option
 		nseClientURLResolver:       null.NewNetworkServiceEndpointRegistryClient(),
 		authorizeNSERegistryClient: authorize.NewNetworkServiceEndpointRegistryClient(authorize.Any()),
 		dialTimeout:                time.Millisecond * 300,
+		nseHealClient:              heal.NewNetworkServiceEndpointRegistryClient(ctx),
+		nseRetryClient:             retry.NewNetworkServiceEndpointRegistryClient(ctx),
 	}
 	for _, opt := range opts {
 		opt(clientOpts)
@@ -55,8 +57,8 @@ func NewNetworkServiceEndpointRegistryClient(ctx context.Context, opts ...Option
 			[]registry.NetworkServiceEndpointRegistryClient{
 				begin.NewNetworkServiceEndpointRegistryClient(),
 				metadata.NewNetworkServiceEndpointClient(),
-				retry.NewNetworkServiceEndpointRegistryClient(ctx),
-				heal.NewNetworkServiceEndpointRegistryClient(ctx),
+				clientOpts.nseRetryClient,
+				clientOpts.nseHealClient,
 				refresh.NewNetworkServiceEndpointRegistryClient(ctx),
 				clientOpts.authorizeNSERegistryClient,
 				clientOpts.nseClientURLResolver,

diff --git a/pkg/registry/chains/client/option.go b/pkg/registry/chains/client/option.go
@@ -47,6 +47,34 @@ func WithNSClientURLResolver(c registry.NetworkServiceRegistryClient) Option {
 	}
 }
 
+// WithNSHealClient overrides default heal network service registry client
+func WithNSHealClient(c registry.NetworkServiceRegistryClient) Option {
+	return func(clientOpts *clientOptions) {
+		clientOpts.nsHealClient = c
+	}
+}
+
+// WithNSRetryClient overrides default retry network service registry client
+func WithNSRetryClient(c registry.NetworkServiceRegistryClient) Option {
+	return func(clientOpts *clientOptions) {
+		clientOpts.nsRetryClient = c
+	}
+}
+
+// WithNSERetryClient overrides default retry network service endpoint registry client
+func WithNSERetryClient(c registry.NetworkServiceEndpointRegistryClient) Option {
+	return func(clientOpts *clientOptions) {
+		clientOpts.nseRetryClient = c
+	}
+}
+
+// WithNSEHealClient overrides default heal network service endpoint registry client
+func WithNSEHealClient(c registry.NetworkServiceEndpointRegistryClient) Option {
+	return func(clientOpts *clientOptions) {
+		clientOpts.nseHealClient = c
+	}
+}
+
 // WithNSEClientURLResolver sets nse client URL resolver
 func WithNSEClientURLResolver(c registry.NetworkServiceEndpointRegistryClient) Option {
 	return func(clientOpts *clientOptions) {
@@ -103,8 +131,9 @@ func WithDialTimeout(dialTimeout time.Duration) Option {
 }
 
 type clientOptions struct {
-	nsClientURLResolver        registry.NetworkServiceRegistryClient
-	nseClientURLResolver       registry.NetworkServiceEndpointRegistryClient
+	nsHealClient, nsRetryClient, nsClientURLResolver    registry.NetworkServiceRegistryClient
+	nseClientURLResolver, nseHealClient, nseRetryClient registry.NetworkServiceEndpointRegistryClient
+
 	authorizeNSRegistryClient  registry.NetworkServiceRegistryClient
 	authorizeNSERegistryClient registry.NetworkServiceEndpointRegistryClient
 	nsAdditionalFunctionality  []registry.NetworkServiceRegistryClient