From e4ab1f2db3d28a68277eeb5802ef65ec0df3b318 Mon Sep 17 00:00:00 2001 From: denis-tingaikin Date: Tue, 10 Sep 2024 15:33:41 +0300 Subject: [PATCH 1/4] add limit chain elements Signed-off-by: denis-tingaikin --- pkg/networkservice/chains/client/client.go | 2 + pkg/networkservice/common/limit/client.go | 111 +++++++++++++++ .../common/limit/client_test.go | 105 ++++++++++++++ pkg/registry/chains/client/ns_client.go | 2 + pkg/registry/chains/client/nse_client.go | 2 + .../chains/proxydns/server_ns_test.go | 4 +- pkg/registry/common/limit/common.go | 32 +++++ pkg/registry/common/limit/ns_client.go | 130 +++++++++++++++++ pkg/registry/common/limit/ns_client_test.go | 124 +++++++++++++++++ pkg/registry/common/limit/nse_client.go | 131 ++++++++++++++++++ pkg/registry/common/limit/nse_client_test.go | 112 +++++++++++++++ 11 files changed, 753 insertions(+), 2 deletions(-) create mode 100644 pkg/networkservice/common/limit/client.go create mode 100644 pkg/networkservice/common/limit/client_test.go create mode 100644 pkg/registry/common/limit/common.go create mode 100644 pkg/registry/common/limit/ns_client.go create mode 100644 pkg/registry/common/limit/ns_client_test.go create mode 100644 pkg/registry/common/limit/nse_client.go create mode 100644 pkg/registry/common/limit/nse_client_test.go diff --git a/pkg/networkservice/chains/client/client.go b/pkg/networkservice/chains/client/client.go index a8467383d..3c0173e90 100644 --- a/pkg/networkservice/chains/client/client.go +++ b/pkg/networkservice/chains/client/client.go @@ -28,6 +28,7 @@ import ( "github.com/networkservicemesh/sdk/pkg/networkservice/common/clienturl" "github.com/networkservicemesh/sdk/pkg/networkservice/common/connect" "github.com/networkservicemesh/sdk/pkg/networkservice/common/dial" + "github.com/networkservicemesh/sdk/pkg/networkservice/common/limit" "github.com/networkservicemesh/sdk/pkg/networkservice/common/null" "github.com/networkservicemesh/sdk/pkg/networkservice/common/refresh" "github.com/networkservicemesh/sdk/pkg/networkservice/common/trimpath" @@ -63,6 +64,7 @@ func NewClient(ctx context.Context, clientOpts ...Option) networkservice.Network dial.WithDialOptions(opts.dialOptions...), dial.WithDialTimeout(opts.dialTimeout), ), + limit.NewClient(), }, append( opts.additionalFunctionality, diff --git a/pkg/networkservice/common/limit/client.go b/pkg/networkservice/common/limit/client.go new file mode 100644 index 000000000..c2b998a7f --- /dev/null +++ b/pkg/networkservice/common/limit/client.go @@ -0,0 +1,111 @@ +// Copyright (c) 2024 Cisco and/or its affiliates. +// +// SPDX-License-Identifier: Apache-2.0 +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at: +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Package limit provides a chain element that can set limits for the RPC calls. +package limit + +import ( + "context" + "time" + + "github.com/golang/protobuf/ptypes/empty" + "github.com/networkservicemesh/api/pkg/api/networkservice" + "github.com/networkservicemesh/sdk/pkg/networkservice/common/clientconn" + "github.com/networkservicemesh/sdk/pkg/networkservice/core/next" + "github.com/networkservicemesh/sdk/pkg/tools/log" + "google.golang.org/grpc" +) + +type Option func(c *limitClient) + +// WithDialLimit sets dial limit +func WithDialLimit(d time.Duration) Option { + return func(c *limitClient) { + c.dialLimit = d + } +} + +type limitClient struct { + dialLimit time.Duration +} + +func NewClient(opts ...Option) networkservice.NetworkServiceClient { + ret := &limitClient{ + dialLimit: time.Minute, + } + + for _, opt := range opts[:] { + opt(ret) + } + + return ret +} + +func (n *limitClient) Request(ctx context.Context, request *networkservice.NetworkServiceRequest, opts ...grpc.CallOption) (*networkservice.Connection, error) { + cc, ok := clientconn.Load(ctx) + if !ok { + return next.Server(ctx).Request(ctx, request) + } + + closer, ok := cc.(interface{ Close() error }) + if !ok { + return next.Server(ctx).Request(ctx, request) + } + + doneCh := make(chan struct{}) + defer close(doneCh) + + logger := log.FromContext(ctx).WithField("throttleClient", "Request") + + go func() { + select { + case <-time.After(n.dialLimit): + logger.Warn("Reached dial limit, closing conneciton...") + _ = closer.Close() + case <-doneCh: + return + } + }() + return next.Client(ctx).Request(ctx, request, opts...) +} + +func (n *limitClient) Close(ctx context.Context, conn *networkservice.Connection, opts ...grpc.CallOption) (*empty.Empty, error) { + cc, ok := clientconn.Load(ctx) + if !ok { + return next.Server(ctx).Close(ctx, conn) + } + + closer, ok := cc.(interface{ Close() error }) + if !ok { + return next.Server(ctx).Close(ctx, conn) + } + + doneCh := make(chan struct{}) + defer close(doneCh) + + logger := log.FromContext(ctx).WithField("throttleClient", "Close") + + go func() { + select { + case <-time.After(n.dialLimit): + logger.Warn("Reached dial limit, closing conneciton...") + _ = closer.Close() + case <-doneCh: + return + } + }() + return next.Client(ctx).Close(ctx, conn, opts...) +} diff --git a/pkg/networkservice/common/limit/client_test.go b/pkg/networkservice/common/limit/client_test.go new file mode 100644 index 000000000..2483c7646 --- /dev/null +++ b/pkg/networkservice/common/limit/client_test.go @@ -0,0 +1,105 @@ +package limit_test + +import ( + "context" + "sync/atomic" + "testing" + "time" + + "github.com/networkservicemesh/api/pkg/api/networkservice" + "github.com/networkservicemesh/sdk/pkg/networkservice/common/clientconn" + "github.com/networkservicemesh/sdk/pkg/networkservice/common/limit" + "github.com/networkservicemesh/sdk/pkg/networkservice/core/chain" + "github.com/networkservicemesh/sdk/pkg/networkservice/utils/checks/checkclose" + "github.com/networkservicemesh/sdk/pkg/networkservice/utils/checks/checkrequest" + "github.com/networkservicemesh/sdk/pkg/networkservice/utils/metadata" + "github.com/stretchr/testify/require" + "go.uber.org/goleak" + "google.golang.org/grpc" +) + +type myConnection struct { + closed atomic.Bool + grpc.ClientConnInterface +} + +func (cc *myConnection) Close() error { + cc.closed.Store(true) + return nil +} + +func Test_DialLimitShouldCalled_OnLimitReached_Request(t *testing.T) { + t.Cleanup(func() { goleak.VerifyNone(t) }) + + var cc = new(myConnection) + var myChain = chain.NewNetworkServiceClient( + metadata.NewClient(), + clientconn.NewClient(cc), + limit.NewClient(limit.WithDialLimit(time.Second/5)), + checkrequest.NewClient(t, func(t *testing.T, nsr *networkservice.NetworkServiceRequest) { + time.Sleep(time.Second / 4) + }), + ) + + _, _ = myChain.Request(context.Background(), &networkservice.NetworkServiceRequest{}) + + require.Eventually(t, func() bool { + return cc.closed.Load() + }, time.Second/2, time.Millisecond*75) +} + +func Test_DialLimitShouldCalled_OnLimitReached_Close(t *testing.T) { + t.Cleanup(func() { goleak.VerifyNone(t) }) + + var cc = new(myConnection) + var myChain = chain.NewNetworkServiceClient( + metadata.NewClient(), + clientconn.NewClient(cc), + limit.NewClient(limit.WithDialLimit(time.Second/5)), + checkclose.NewClient(t, func(t *testing.T, nsr *networkservice.Connection) { + time.Sleep(time.Second / 4) + }), + ) + + _, _ = myChain.Request(context.Background(), &networkservice.NetworkServiceRequest{}) + _, _ = myChain.Close(context.Background(), &networkservice.Connection{}) + + require.Eventually(t, func() bool { + return cc.closed.Load() + }, time.Second/2, time.Millisecond*75) +} + +func Test_DialLimitShouldNotBeCalled_OnSuccesRequest(t *testing.T) { + t.Cleanup(func() { goleak.VerifyNone(t) }) + + var cc = new(myConnection) + var myChain = chain.NewNetworkServiceClient( + metadata.NewClient(), + clientconn.NewClient(cc), + limit.NewClient(limit.WithDialLimit(time.Second/5)), + ) + + _, _ = myChain.Request(context.Background(), &networkservice.NetworkServiceRequest{}) + + require.Never(t, func() bool { + return cc.closed.Load() + }, time.Second/2, time.Millisecond*75) +} + +func Test_DialLimitShouldNotBeCalled_OnSuccessClose(t *testing.T) { + t.Cleanup(func() { goleak.VerifyNone(t) }) + + var cc = new(myConnection) + var myChain = chain.NewNetworkServiceClient( + metadata.NewClient(), + clientconn.NewClient(cc), + limit.NewClient(limit.WithDialLimit(time.Second/5)), + ) + + _, _ = myChain.Request(context.Background(), &networkservice.NetworkServiceRequest{}) + _, _ = myChain.Close(context.Background(), &networkservice.Connection{}) + + require.Never(t, func() bool { + return cc.closed.Load() + }, time.Second/2, time.Millisecond*75) +} diff --git a/pkg/registry/chains/client/ns_client.go b/pkg/registry/chains/client/ns_client.go index 749151548..bcdbdce72 100644 --- a/pkg/registry/chains/client/ns_client.go +++ b/pkg/registry/chains/client/ns_client.go @@ -31,6 +31,7 @@ import ( "github.com/networkservicemesh/sdk/pkg/registry/common/dial" "github.com/networkservicemesh/sdk/pkg/registry/common/grpcmetadata" "github.com/networkservicemesh/sdk/pkg/registry/common/heal" + "github.com/networkservicemesh/sdk/pkg/registry/common/limit" "github.com/networkservicemesh/sdk/pkg/registry/common/null" "github.com/networkservicemesh/sdk/pkg/registry/common/retry" "github.com/networkservicemesh/sdk/pkg/registry/core/chain" @@ -63,6 +64,7 @@ func NewNetworkServiceRegistryClient(ctx context.Context, opts ...Option) regist dial.WithDialTimeout(clientOpts.dialTimeout), dial.WithDialOptions(clientOpts.dialOptions...), ), + limit.NewNetworkServiceRegistryClient(), }, append( clientOpts.nsAdditionalFunctionality, diff --git a/pkg/registry/chains/client/nse_client.go b/pkg/registry/chains/client/nse_client.go index 7ef1d7ae5..a5cc0af39 100644 --- a/pkg/registry/chains/client/nse_client.go +++ b/pkg/registry/chains/client/nse_client.go @@ -32,6 +32,7 @@ import ( "github.com/networkservicemesh/sdk/pkg/registry/common/dial" "github.com/networkservicemesh/sdk/pkg/registry/common/grpcmetadata" "github.com/networkservicemesh/sdk/pkg/registry/common/heal" + "github.com/networkservicemesh/sdk/pkg/registry/common/limit" "github.com/networkservicemesh/sdk/pkg/registry/common/null" "github.com/networkservicemesh/sdk/pkg/registry/common/refresh" "github.com/networkservicemesh/sdk/pkg/registry/common/retry" @@ -66,6 +67,7 @@ func NewNetworkServiceEndpointRegistryClient(ctx context.Context, opts ...Option dial.WithDialTimeout(clientOpts.dialTimeout), dial.WithDialOptions(clientOpts.dialOptions...), ), + limit.NewNetworkServiceEndpointRegistryClient(), }, append( clientOpts.nseAdditionalFunctionality, diff --git a/pkg/registry/chains/proxydns/server_ns_test.go b/pkg/registry/chains/proxydns/server_ns_test.go index f7fba1294..316273c86 100644 --- a/pkg/registry/chains/proxydns/server_ns_test.go +++ b/pkg/registry/chains/proxydns/server_ns_test.go @@ -1,6 +1,6 @@ // Copyright (c) 2020-2022 Doc.ai and/or its affiliates. // -// Copyright (c) 2022 Cisco Systems, Inc. +// Copyright (c) 2022-2024 Cisco Systems, Inc. // // SPDX-License-Identifier: Apache-2.0 // @@ -143,7 +143,7 @@ func TestLocalDomain_NetworkServiceRegistry(t *testing.T) { registryclient.WithDialOptions(grpc.WithTransportCredentials(insecure.NewCredentials())), registryclient.WithClientURL(domain1.Registry.URL)) - stream, err := client2.Find(context.Background(), ®istryapi.NetworkServiceQuery{ + stream, err := client2.Find(ctx, ®istryapi.NetworkServiceQuery{ NetworkService: ®istryapi.NetworkService{ Name: expected.Name, }, diff --git a/pkg/registry/common/limit/common.go b/pkg/registry/common/limit/common.go new file mode 100644 index 000000000..c745daba7 --- /dev/null +++ b/pkg/registry/common/limit/common.go @@ -0,0 +1,32 @@ +// Copyright (c) 2024 Cisco and/or its affiliates. +// +// SPDX-License-Identifier: Apache-2.0 +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at: +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package limit + +import "time" + +type limitConfig struct { + dialLimit time.Duration +} + +type Option func(cfg *limitConfig) + +// WithDialLimit sets dial time limit +func WithDialLimit(t time.Duration) Option { + return Option(func(cfg *limitConfig) { + cfg.dialLimit = t + }) +} diff --git a/pkg/registry/common/limit/ns_client.go b/pkg/registry/common/limit/ns_client.go new file mode 100644 index 000000000..53d063f9e --- /dev/null +++ b/pkg/registry/common/limit/ns_client.go @@ -0,0 +1,130 @@ +// Copyright (c) 2024 Cisco and/or its affiliates. +// +// SPDX-License-Identifier: Apache-2.0 +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at: +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Package limit provides a chain element that can set limits for the RPC calls. +package limit + +import ( + "context" + "time" + + "github.com/golang/protobuf/ptypes/empty" + "github.com/networkservicemesh/api/pkg/api/registry" + "github.com/networkservicemesh/sdk/pkg/registry/common/clientconn" + "github.com/networkservicemesh/sdk/pkg/registry/core/next" + "github.com/networkservicemesh/sdk/pkg/tools/log" + "google.golang.org/grpc" +) + +type limitNSClient struct { + cfg *limitConfig +} + +func (n *limitNSClient) Register(ctx context.Context, in *registry.NetworkService, opts ...grpc.CallOption) (*registry.NetworkService, error) { + cc, ok := clientconn.Load(ctx) + if !ok { + return next.NetworkServiceRegistryClient(ctx).Register(ctx, in, opts...) + } + + closer, ok := cc.(interface{ Close() error }) + if !ok { + return next.NetworkServiceRegistryClient(ctx).Register(ctx, in, opts...) + } + + doneCh := make(chan struct{}) + defer close(doneCh) + + logger := log.FromContext(ctx).WithField("throttleNSClient", "Register") + + go func() { + select { + case <-time.After(n.cfg.dialLimit): + logger.Warn("Reached dial limit, closing conneciton...") + _ = closer.Close() + case <-doneCh: + return + } + }() + return next.NetworkServiceRegistryClient(ctx).Register(ctx, in, opts...) +} + +func (n *limitNSClient) Find(ctx context.Context, in *registry.NetworkServiceQuery, opts ...grpc.CallOption) (registry.NetworkServiceRegistry_FindClient, error) { + cc, ok := clientconn.Load(ctx) + if !ok { + return next.NetworkServiceRegistryClient(ctx).Find(ctx, in, opts...) + } + + closer, ok := cc.(interface{ Close() error }) + if !ok { + return next.NetworkServiceRegistryClient(ctx).Find(ctx, in, opts...) + } + + logger := log.FromContext(ctx).WithField("throttleNSClient", "Find") + + go func() { + select { + case <-time.After(n.cfg.dialLimit): + logger.Warn("Reached dial limit, closing conneciton...") + _ = closer.Close() + case <-ctx.Done(): + return + } + }() + + return next.NetworkServiceRegistryClient(ctx).Find(ctx, in, opts...) +} + +func (n *limitNSClient) Unregister(ctx context.Context, in *registry.NetworkService, opts ...grpc.CallOption) (*empty.Empty, error) { + cc, ok := clientconn.Load(ctx) + if !ok { + return next.NetworkServiceRegistryClient(ctx).Unregister(ctx, in, opts...) + } + + closer, ok := cc.(interface{ Close() error }) + if !ok { + return next.NetworkServiceRegistryClient(ctx).Unregister(ctx, in, opts...) + } + + doneCh := make(chan struct{}) + defer close(doneCh) + + logger := log.FromContext(ctx).WithField("throttleNSClient", "Unregister") + + go func() { + select { + case <-time.After(n.cfg.dialLimit): + logger.Warn("Reached dial limit, closing conneciton...") + _ = closer.Close() + case <-doneCh: + return + } + }() + + return next.NetworkServiceRegistryClient(ctx).Unregister(ctx, in, opts...) +} + +// NewNetworkServiceRegistryClient - returns a new null client that does nothing but call next.NetworkServiceRegistryClient(ctx). +func NewNetworkServiceRegistryClient(opts ...Option) registry.NetworkServiceRegistryClient { + cfg := &limitConfig{ + dialLimit: time.Minute, + } + for _, opt := range opts[:] { + opt(cfg) + } + return &limitNSClient{ + cfg: cfg, + } +} diff --git a/pkg/registry/common/limit/ns_client_test.go b/pkg/registry/common/limit/ns_client_test.go new file mode 100644 index 000000000..4f15d8514 --- /dev/null +++ b/pkg/registry/common/limit/ns_client_test.go @@ -0,0 +1,124 @@ +// Copyright (c) 2024 Cisco and/or its affiliates. +// +// SPDX-License-Identifier: Apache-2.0 +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at: +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package limit_test + +import ( + "context" + "sync/atomic" + "testing" + "time" + + "github.com/networkservicemesh/api/pkg/api/registry" + "github.com/networkservicemesh/sdk/pkg/registry/common/clientconn" + "github.com/networkservicemesh/sdk/pkg/registry/common/limit" + "github.com/networkservicemesh/sdk/pkg/registry/utils/checks/checkcontext" + + "github.com/networkservicemesh/sdk/pkg/registry/core/chain" + "github.com/networkservicemesh/sdk/pkg/registry/utils/metadata" + "github.com/stretchr/testify/require" + "go.uber.org/goleak" + "google.golang.org/grpc" +) + +type myConnection struct { + closed atomic.Bool + grpc.ClientConnInterface +} + +func (cc *myConnection) Close() error { + cc.closed.Store(true) + return nil +} + +func Test_DialLimitShouldCalled_OnLimitReached(t *testing.T) { + t.Cleanup(func() { goleak.VerifyNone(t) }) + + var cc = new(myConnection) + var myChain = chain.NewNetworkServiceRegistryClient( + metadata.NewNetworkServiceClient(), + clientconn.NewNetworkServiceRegistryClient(), + checkcontext.NewNSClient(t, func(t *testing.T, ctx context.Context) { + clientconn.Store(ctx, cc) + + }), + limit.NewNetworkServiceRegistryClient(limit.WithDialLimit(time.Second/5)), + checkcontext.NewNSClient(t, func(t *testing.T, ctx context.Context) { + time.Sleep(time.Second / 5) + }), + ) + + _, _ = myChain.Register(context.Background(), ®istry.NetworkService{Name: t.Name()}) + + require.Eventually(t, func() bool { + return cc.closed.Load() + }, time.Second/2, time.Millisecond*75) + + cc.closed.Store(false) + + _, _ = myChain.Find(context.Background(), ®istry.NetworkServiceQuery{NetworkService: ®istry.NetworkService{Name: t.Name()}}) + + require.Eventually(t, func() bool { + return cc.closed.Load() + }, time.Second/2, time.Millisecond*75) + + cc.closed.Store(false) + + _, _ = myChain.Unregister(context.Background(), ®istry.NetworkService{Name: t.Name()}) + + require.Eventually(t, func() bool { + return cc.closed.Load() + }, time.Second/2, time.Millisecond*75) +} + +func Test_DialLimitShouldNotBeCalled_OnSuccess(t *testing.T) { + t.Cleanup(func() { goleak.VerifyNone(t) }) + + var cc = new(myConnection) + var myChain = chain.NewNetworkServiceRegistryClient( + metadata.NewNetworkServiceClient(), + clientconn.NewNetworkServiceRegistryClient(), + checkcontext.NewNSClient(t, func(t *testing.T, ctx context.Context) { + clientconn.Store(ctx, cc) + + }), + limit.NewNetworkServiceRegistryClient(limit.WithDialLimit(time.Second/5)), + ) + + ctx, cancel := context.WithCancel(context.Background()) + _, _ = myChain.Register(ctx, ®istry.NetworkService{Name: t.Name()}) + cancel() + + require.Never(t, func() bool { + return cc.closed.Load() + }, time.Second/2, time.Millisecond*75) + + ctx, cancel = context.WithCancel(context.Background()) + _, _ = myChain.Find(ctx, ®istry.NetworkServiceQuery{NetworkService: ®istry.NetworkService{Name: t.Name()}}) + cancel() + + require.Never(t, func() bool { + return cc.closed.Load() + }, time.Second/2, time.Millisecond*75) + + ctx, cancel = context.WithCancel(context.Background()) + _, _ = myChain.Unregister(ctx, ®istry.NetworkService{Name: t.Name()}) + cancel() + + require.Never(t, func() bool { + return cc.closed.Load() + }, time.Second/2, time.Millisecond*75) +} diff --git a/pkg/registry/common/limit/nse_client.go b/pkg/registry/common/limit/nse_client.go new file mode 100644 index 000000000..5d3dc0fad --- /dev/null +++ b/pkg/registry/common/limit/nse_client.go @@ -0,0 +1,131 @@ +// Copyright (c) 2024 Cisco and/or its affiliates. +// +// SPDX-License-Identifier: Apache-2.0 +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at: +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package limit + +import ( + "context" + "time" + + "github.com/golang/protobuf/ptypes/empty" + "github.com/networkservicemesh/api/pkg/api/registry" + "github.com/networkservicemesh/sdk/pkg/registry/common/clientconn" + "github.com/networkservicemesh/sdk/pkg/registry/core/next" + "github.com/networkservicemesh/sdk/pkg/tools/log" + "google.golang.org/grpc" +) + +type limitNSEClient struct { + cfg *limitConfig +} + +func (n *limitNSEClient) Register(ctx context.Context, in *registry.NetworkServiceEndpoint, opts ...grpc.CallOption) (*registry.NetworkServiceEndpoint, error) { + cc, ok := clientconn.Load(ctx) + if !ok { + return next.NetworkServiceEndpointRegistryClient(ctx).Register(ctx, in, opts...) + } + + closer, ok := cc.(interface{ Close() error }) + if !ok { + return next.NetworkServiceEndpointRegistryClient(ctx).Register(ctx, in, opts...) + } + + doneCh := make(chan struct{}) + defer close(doneCh) + + logger := log.FromContext(ctx).WithField("throttleNSEClient", "Register") + + go func() { + select { + case <-time.After(n.cfg.dialLimit): + logger.Warn("Reached dial limit, closing conneciton...") + _ = closer.Close() + case <-doneCh: + return + } + }() + return next.NetworkServiceEndpointRegistryClient(ctx).Register(ctx, in, opts...) +} + +func (n *limitNSEClient) Find(ctx context.Context, in *registry.NetworkServiceEndpointQuery, opts ...grpc.CallOption) (registry.NetworkServiceEndpointRegistry_FindClient, error) { + cc, ok := clientconn.Load(ctx) + if !ok { + return next.NetworkServiceEndpointRegistryClient(ctx).Find(ctx, in, opts...) + } + + closer, ok := cc.(interface{ Close() error }) + if !ok { + return next.NetworkServiceEndpointRegistryClient(ctx).Find(ctx, in, opts...) + } + + logger := log.FromContext(ctx).WithField("throttleNSEClient", "Find") + + go func() { + select { + case <-time.After(n.cfg.dialLimit): + logger.Warn("Reached dial limit, closing conneciton...") + _ = closer.Close() + case <-ctx.Done(): + return + } + }() + + return next.NetworkServiceEndpointRegistryClient(ctx).Find(ctx, in, opts...) +} + +func (n *limitNSEClient) Unregister(ctx context.Context, in *registry.NetworkServiceEndpoint, opts ...grpc.CallOption) (*empty.Empty, error) { + cc, ok := clientconn.Load(ctx) + if !ok { + return next.NetworkServiceEndpointRegistryClient(ctx).Unregister(ctx, in, opts...) + } + + closer, ok := cc.(interface{ Close() error }) + if !ok { + return next.NetworkServiceEndpointRegistryClient(ctx).Unregister(ctx, in, opts...) + } + + doneCh := make(chan struct{}) + defer close(doneCh) + + logger := log.FromContext(ctx).WithField("throttleNSEClient", "Unregister") + + go func() { + select { + case <-time.After(n.cfg.dialLimit): + logger.Warn("Reached dial limit, closing conneciton...") + _ = closer.Close() + case <-doneCh: + return + } + }() + + return next.NetworkServiceEndpointRegistryClient(ctx).Unregister(ctx, in, opts...) +} + +// NewNetworkServiceEndpointRegistryClient - returns a new null client that does nothing but call next.NetworkServiceEndpointRegistryClient(ctx). +func NewNetworkServiceEndpointRegistryClient(opts ...Option) registry.NetworkServiceEndpointRegistryClient { + cfg := &limitConfig{ + dialLimit: time.Minute, + } + + for _, opt := range opts[:] { + opt(cfg) + } + + return &limitNSEClient{ + cfg: cfg, + } +} diff --git a/pkg/registry/common/limit/nse_client_test.go b/pkg/registry/common/limit/nse_client_test.go new file mode 100644 index 000000000..997416227 --- /dev/null +++ b/pkg/registry/common/limit/nse_client_test.go @@ -0,0 +1,112 @@ +// Copyright (c) 2024 Cisco and/or its affiliates. +// +// SPDX-License-Identifier: Apache-2.0 +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at: +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package limit_test + +import ( + "context" + "testing" + "time" + + "github.com/networkservicemesh/api/pkg/api/registry" + "github.com/networkservicemesh/sdk/pkg/registry/common/clientconn" + "github.com/networkservicemesh/sdk/pkg/registry/common/limit" + "github.com/networkservicemesh/sdk/pkg/registry/utils/checks/checkcontext" + + "github.com/networkservicemesh/sdk/pkg/registry/core/chain" + "github.com/networkservicemesh/sdk/pkg/registry/utils/metadata" + "github.com/stretchr/testify/require" + "go.uber.org/goleak" +) + +func Test_NSEDialLimitShouldCalled_OnLimitReached(t *testing.T) { + t.Cleanup(func() { goleak.VerifyNone(t) }) + + var cc = new(myConnection) + var myChain = chain.NewNetworkServiceEndpointRegistryClient( + metadata.NewNetworkServiceEndpointClient(), + clientconn.NewNetworkServiceEndpointRegistryClient(), + checkcontext.NewNSEClient(t, func(t *testing.T, ctx context.Context) { + clientconn.Store(ctx, cc) + + }), + limit.NewNetworkServiceEndpointRegistryClient(limit.WithDialLimit(time.Second/5)), + checkcontext.NewNSEClient(t, func(t *testing.T, ctx context.Context) { + time.Sleep(time.Second / 5) + }), + ) + + _, _ = myChain.Register(context.Background(), ®istry.NetworkServiceEndpoint{Name: t.Name()}) + + require.Eventually(t, func() bool { + return cc.closed.Load() + }, time.Second/2, time.Millisecond*75) + + cc.closed.Store(false) + + _, _ = myChain.Find(context.Background(), ®istry.NetworkServiceEndpointQuery{NetworkServiceEndpoint: ®istry.NetworkServiceEndpoint{Name: t.Name()}}) + + require.Eventually(t, func() bool { + return cc.closed.Load() + }, time.Second/2, time.Millisecond*75) + + cc.closed.Store(false) + + _, _ = myChain.Unregister(context.Background(), ®istry.NetworkServiceEndpoint{Name: t.Name()}) + + require.Eventually(t, func() bool { + return cc.closed.Load() + }, time.Second/2, time.Millisecond*75) +} + +func Test_NSEDialLimitShouldNotBeCalled_OnSuccess(t *testing.T) { + t.Cleanup(func() { goleak.VerifyNone(t) }) + + var cc = new(myConnection) + var myChain = chain.NewNetworkServiceEndpointRegistryClient( + metadata.NewNetworkServiceEndpointClient(), + clientconn.NewNetworkServiceEndpointRegistryClient(), + checkcontext.NewNSEClient(t, func(t *testing.T, ctx context.Context) { + clientconn.Store(ctx, cc) + + }), + limit.NewNetworkServiceEndpointRegistryClient(limit.WithDialLimit(time.Second/5)), + ) + + ctx, cancel := context.WithCancel(context.Background()) + _, _ = myChain.Register(ctx, ®istry.NetworkServiceEndpoint{Name: t.Name()}) + cancel() + + require.Never(t, func() bool { + return cc.closed.Load() + }, time.Second/2, time.Millisecond*75) + + ctx, cancel = context.WithCancel(context.Background()) + _, _ = myChain.Find(ctx, ®istry.NetworkServiceEndpointQuery{NetworkServiceEndpoint: ®istry.NetworkServiceEndpoint{Name: t.Name()}}) + cancel() + + require.Never(t, func() bool { + return cc.closed.Load() + }, time.Second/2, time.Millisecond*75) + + ctx, cancel = context.WithCancel(context.Background()) + _, _ = myChain.Unregister(ctx, ®istry.NetworkServiceEndpoint{Name: t.Name()}) + cancel() + + require.Never(t, func() bool { + return cc.closed.Load() + }, time.Second/2, time.Millisecond*75) +} From 5e94dfb33b8178caa8bff53d20bc1dd5eb0d23ca Mon Sep 17 00:00:00 2001 From: denis-tingaikin Date: Tue, 10 Sep 2024 16:17:45 +0300 Subject: [PATCH 2/4] fix unstable tests Signed-off-by: denis-tingaikin --- pkg/registry/common/limit/ns_client.go | 19 +++++++++++++++++-- pkg/registry/common/limit/nse_client.go | 20 ++++++++++++++++++-- 2 files changed, 35 insertions(+), 4 deletions(-) diff --git a/pkg/registry/common/limit/ns_client.go b/pkg/registry/common/limit/ns_client.go index 53d063f9e..0ca54cd48 100644 --- a/pkg/registry/common/limit/ns_client.go +++ b/pkg/registry/common/limit/ns_client.go @@ -74,17 +74,32 @@ func (n *limitNSClient) Find(ctx context.Context, in *registry.NetworkServiceQue logger := log.FromContext(ctx).WithField("throttleNSClient", "Find") + doneCh := make(chan struct{}) + defer close(doneCh) + go func() { select { case <-time.After(n.cfg.dialLimit): logger.Warn("Reached dial limit, closing conneciton...") _ = closer.Close() - case <-ctx.Done(): + case <-doneCh: return } }() - return next.NetworkServiceRegistryClient(ctx).Find(ctx, in, opts...) + resp, err := next.NetworkServiceRegistryClient(ctx).Find(ctx, in, opts...) + if err == nil { + go func() { + select { + case <-time.After(n.cfg.dialLimit): + logger.Warn("Reached dial limit, closing conneciton...") + _ = closer.Close() + case <-resp.Context().Done(): + return + } + }() + } + return resp, err } func (n *limitNSClient) Unregister(ctx context.Context, in *registry.NetworkService, opts ...grpc.CallOption) (*empty.Empty, error) { diff --git a/pkg/registry/common/limit/nse_client.go b/pkg/registry/common/limit/nse_client.go index 5d3dc0fad..e5360117d 100644 --- a/pkg/registry/common/limit/nse_client.go +++ b/pkg/registry/common/limit/nse_client.go @@ -72,18 +72,34 @@ func (n *limitNSEClient) Find(ctx context.Context, in *registry.NetworkServiceEn } logger := log.FromContext(ctx).WithField("throttleNSEClient", "Find") + doneCh := make(chan struct{}) + defer close(doneCh) go func() { select { case <-time.After(n.cfg.dialLimit): logger.Warn("Reached dial limit, closing conneciton...") _ = closer.Close() - case <-ctx.Done(): + case <-doneCh: return } }() - return next.NetworkServiceEndpointRegistryClient(ctx).Find(ctx, in, opts...) + resp, err := next.NetworkServiceEndpointRegistryClient(ctx).Find(ctx, in, opts...) + + if err == nil { + go func() { + select { + case <-time.After(n.cfg.dialLimit): + logger.Warn("Reached dial limit, closing conneciton...") + _ = closer.Close() + case <-resp.Context().Done(): + return + } + }() + } + + return resp, err } func (n *limitNSEClient) Unregister(ctx context.Context, in *registry.NetworkServiceEndpoint, opts ...grpc.CallOption) (*empty.Empty, error) { From bad093e9ef383e88fe60fa9a65ca4b9814232b62 Mon Sep 17 00:00:00 2001 From: denis-tingaikin Date: Tue, 10 Sep 2024 16:21:27 +0300 Subject: [PATCH 3/4] fix linter Signed-off-by: denis-tingaikin --- .golangci.yml | 1 - pkg/networkservice/common/limit/client.go | 9 +++++---- pkg/networkservice/common/limit/client_test.go | 7 ++++--- pkg/registry/common/limit/ns_client.go | 13 +++++++------ pkg/registry/common/limit/ns_client_test.go | 8 ++++---- pkg/registry/common/limit/nse_client.go | 13 +++++++------ pkg/registry/common/limit/nse_client_test.go | 8 ++++---- 7 files changed, 31 insertions(+), 28 deletions(-) diff --git a/.golangci.yml b/.golangci.yml index 3d9c8b0fe..6fe97685f 100644 --- a/.golangci.yml +++ b/.golangci.yml @@ -132,7 +132,6 @@ linters: - deadcode - depguard - dogsled - - dupl - errcheck - funlen - gochecknoinits diff --git a/pkg/networkservice/common/limit/client.go b/pkg/networkservice/common/limit/client.go index c2b998a7f..c103c18a4 100644 --- a/pkg/networkservice/common/limit/client.go +++ b/pkg/networkservice/common/limit/client.go @@ -23,10 +23,11 @@ import ( "github.com/golang/protobuf/ptypes/empty" "github.com/networkservicemesh/api/pkg/api/networkservice" + "google.golang.org/grpc" + "github.com/networkservicemesh/sdk/pkg/networkservice/common/clientconn" "github.com/networkservicemesh/sdk/pkg/networkservice/core/next" "github.com/networkservicemesh/sdk/pkg/tools/log" - "google.golang.org/grpc" ) type Option func(c *limitClient) @@ -47,7 +48,7 @@ func NewClient(opts ...Option) networkservice.NetworkServiceClient { dialLimit: time.Minute, } - for _, opt := range opts[:] { + for _, opt := range opts { opt(ret) } @@ -73,7 +74,7 @@ func (n *limitClient) Request(ctx context.Context, request *networkservice.Netwo go func() { select { case <-time.After(n.dialLimit): - logger.Warn("Reached dial limit, closing conneciton...") + logger.Warn("Reached dial limit, closing connection...") _ = closer.Close() case <-doneCh: return @@ -101,7 +102,7 @@ func (n *limitClient) Close(ctx context.Context, conn *networkservice.Connection go func() { select { case <-time.After(n.dialLimit): - logger.Warn("Reached dial limit, closing conneciton...") + logger.Warn("Reached dial limit, closing connection...") _ = closer.Close() case <-doneCh: return diff --git a/pkg/networkservice/common/limit/client_test.go b/pkg/networkservice/common/limit/client_test.go index 2483c7646..ca2956697 100644 --- a/pkg/networkservice/common/limit/client_test.go +++ b/pkg/networkservice/common/limit/client_test.go @@ -7,15 +7,16 @@ import ( "time" "github.com/networkservicemesh/api/pkg/api/networkservice" + "github.com/stretchr/testify/require" + "go.uber.org/goleak" + "google.golang.org/grpc" + "github.com/networkservicemesh/sdk/pkg/networkservice/common/clientconn" "github.com/networkservicemesh/sdk/pkg/networkservice/common/limit" "github.com/networkservicemesh/sdk/pkg/networkservice/core/chain" "github.com/networkservicemesh/sdk/pkg/networkservice/utils/checks/checkclose" "github.com/networkservicemesh/sdk/pkg/networkservice/utils/checks/checkrequest" "github.com/networkservicemesh/sdk/pkg/networkservice/utils/metadata" - "github.com/stretchr/testify/require" - "go.uber.org/goleak" - "google.golang.org/grpc" ) type myConnection struct { diff --git a/pkg/registry/common/limit/ns_client.go b/pkg/registry/common/limit/ns_client.go index 0ca54cd48..d9079d81c 100644 --- a/pkg/registry/common/limit/ns_client.go +++ b/pkg/registry/common/limit/ns_client.go @@ -23,10 +23,11 @@ import ( "github.com/golang/protobuf/ptypes/empty" "github.com/networkservicemesh/api/pkg/api/registry" + "google.golang.org/grpc" + "github.com/networkservicemesh/sdk/pkg/registry/common/clientconn" "github.com/networkservicemesh/sdk/pkg/registry/core/next" "github.com/networkservicemesh/sdk/pkg/tools/log" - "google.golang.org/grpc" ) type limitNSClient struct { @@ -52,7 +53,7 @@ func (n *limitNSClient) Register(ctx context.Context, in *registry.NetworkServic go func() { select { case <-time.After(n.cfg.dialLimit): - logger.Warn("Reached dial limit, closing conneciton...") + logger.Warn("Reached dial limit, closing connection...") _ = closer.Close() case <-doneCh: return @@ -80,7 +81,7 @@ func (n *limitNSClient) Find(ctx context.Context, in *registry.NetworkServiceQue go func() { select { case <-time.After(n.cfg.dialLimit): - logger.Warn("Reached dial limit, closing conneciton...") + logger.Warn("Reached dial limit, closing connection...") _ = closer.Close() case <-doneCh: return @@ -92,7 +93,7 @@ func (n *limitNSClient) Find(ctx context.Context, in *registry.NetworkServiceQue go func() { select { case <-time.After(n.cfg.dialLimit): - logger.Warn("Reached dial limit, closing conneciton...") + logger.Warn("Reached dial limit, closing connection...") _ = closer.Close() case <-resp.Context().Done(): return @@ -121,7 +122,7 @@ func (n *limitNSClient) Unregister(ctx context.Context, in *registry.NetworkServ go func() { select { case <-time.After(n.cfg.dialLimit): - logger.Warn("Reached dial limit, closing conneciton...") + logger.Warn("Reached dial limit, closing connection...") _ = closer.Close() case <-doneCh: return @@ -136,7 +137,7 @@ func NewNetworkServiceRegistryClient(opts ...Option) registry.NetworkServiceRegi cfg := &limitConfig{ dialLimit: time.Minute, } - for _, opt := range opts[:] { + for _, opt := range opts { opt(cfg) } return &limitNSClient{ diff --git a/pkg/registry/common/limit/ns_client_test.go b/pkg/registry/common/limit/ns_client_test.go index 4f15d8514..cd7e647f0 100644 --- a/pkg/registry/common/limit/ns_client_test.go +++ b/pkg/registry/common/limit/ns_client_test.go @@ -23,15 +23,17 @@ import ( "time" "github.com/networkservicemesh/api/pkg/api/registry" + "github.com/networkservicemesh/sdk/pkg/registry/common/clientconn" "github.com/networkservicemesh/sdk/pkg/registry/common/limit" "github.com/networkservicemesh/sdk/pkg/registry/utils/checks/checkcontext" - "github.com/networkservicemesh/sdk/pkg/registry/core/chain" - "github.com/networkservicemesh/sdk/pkg/registry/utils/metadata" "github.com/stretchr/testify/require" "go.uber.org/goleak" "google.golang.org/grpc" + + "github.com/networkservicemesh/sdk/pkg/registry/core/chain" + "github.com/networkservicemesh/sdk/pkg/registry/utils/metadata" ) type myConnection struct { @@ -53,7 +55,6 @@ func Test_DialLimitShouldCalled_OnLimitReached(t *testing.T) { clientconn.NewNetworkServiceRegistryClient(), checkcontext.NewNSClient(t, func(t *testing.T, ctx context.Context) { clientconn.Store(ctx, cc) - }), limit.NewNetworkServiceRegistryClient(limit.WithDialLimit(time.Second/5)), checkcontext.NewNSClient(t, func(t *testing.T, ctx context.Context) { @@ -93,7 +94,6 @@ func Test_DialLimitShouldNotBeCalled_OnSuccess(t *testing.T) { clientconn.NewNetworkServiceRegistryClient(), checkcontext.NewNSClient(t, func(t *testing.T, ctx context.Context) { clientconn.Store(ctx, cc) - }), limit.NewNetworkServiceRegistryClient(limit.WithDialLimit(time.Second/5)), ) diff --git a/pkg/registry/common/limit/nse_client.go b/pkg/registry/common/limit/nse_client.go index e5360117d..4d98a4d54 100644 --- a/pkg/registry/common/limit/nse_client.go +++ b/pkg/registry/common/limit/nse_client.go @@ -22,10 +22,11 @@ import ( "github.com/golang/protobuf/ptypes/empty" "github.com/networkservicemesh/api/pkg/api/registry" + "google.golang.org/grpc" + "github.com/networkservicemesh/sdk/pkg/registry/common/clientconn" "github.com/networkservicemesh/sdk/pkg/registry/core/next" "github.com/networkservicemesh/sdk/pkg/tools/log" - "google.golang.org/grpc" ) type limitNSEClient struct { @@ -51,7 +52,7 @@ func (n *limitNSEClient) Register(ctx context.Context, in *registry.NetworkServi go func() { select { case <-time.After(n.cfg.dialLimit): - logger.Warn("Reached dial limit, closing conneciton...") + logger.Warn("Reached dial limit, closing connection...") _ = closer.Close() case <-doneCh: return @@ -78,7 +79,7 @@ func (n *limitNSEClient) Find(ctx context.Context, in *registry.NetworkServiceEn go func() { select { case <-time.After(n.cfg.dialLimit): - logger.Warn("Reached dial limit, closing conneciton...") + logger.Warn("Reached dial limit, closing connection...") _ = closer.Close() case <-doneCh: return @@ -91,7 +92,7 @@ func (n *limitNSEClient) Find(ctx context.Context, in *registry.NetworkServiceEn go func() { select { case <-time.After(n.cfg.dialLimit): - logger.Warn("Reached dial limit, closing conneciton...") + logger.Warn("Reached dial limit, closing connection...") _ = closer.Close() case <-resp.Context().Done(): return @@ -121,7 +122,7 @@ func (n *limitNSEClient) Unregister(ctx context.Context, in *registry.NetworkSer go func() { select { case <-time.After(n.cfg.dialLimit): - logger.Warn("Reached dial limit, closing conneciton...") + logger.Warn("Reached dial limit, closing connection...") _ = closer.Close() case <-doneCh: return @@ -137,7 +138,7 @@ func NewNetworkServiceEndpointRegistryClient(opts ...Option) registry.NetworkSer dialLimit: time.Minute, } - for _, opt := range opts[:] { + for _, opt := range opts { opt(cfg) } diff --git a/pkg/registry/common/limit/nse_client_test.go b/pkg/registry/common/limit/nse_client_test.go index 997416227..6939ac26f 100644 --- a/pkg/registry/common/limit/nse_client_test.go +++ b/pkg/registry/common/limit/nse_client_test.go @@ -22,14 +22,16 @@ import ( "time" "github.com/networkservicemesh/api/pkg/api/registry" + "github.com/networkservicemesh/sdk/pkg/registry/common/clientconn" "github.com/networkservicemesh/sdk/pkg/registry/common/limit" "github.com/networkservicemesh/sdk/pkg/registry/utils/checks/checkcontext" - "github.com/networkservicemesh/sdk/pkg/registry/core/chain" - "github.com/networkservicemesh/sdk/pkg/registry/utils/metadata" "github.com/stretchr/testify/require" "go.uber.org/goleak" + + "github.com/networkservicemesh/sdk/pkg/registry/core/chain" + "github.com/networkservicemesh/sdk/pkg/registry/utils/metadata" ) func Test_NSEDialLimitShouldCalled_OnLimitReached(t *testing.T) { @@ -41,7 +43,6 @@ func Test_NSEDialLimitShouldCalled_OnLimitReached(t *testing.T) { clientconn.NewNetworkServiceEndpointRegistryClient(), checkcontext.NewNSEClient(t, func(t *testing.T, ctx context.Context) { clientconn.Store(ctx, cc) - }), limit.NewNetworkServiceEndpointRegistryClient(limit.WithDialLimit(time.Second/5)), checkcontext.NewNSEClient(t, func(t *testing.T, ctx context.Context) { @@ -81,7 +82,6 @@ func Test_NSEDialLimitShouldNotBeCalled_OnSuccess(t *testing.T) { clientconn.NewNetworkServiceEndpointRegistryClient(), checkcontext.NewNSEClient(t, func(t *testing.T, ctx context.Context) { clientconn.Store(ctx, cc) - }), limit.NewNetworkServiceEndpointRegistryClient(limit.WithDialLimit(time.Second/5)), ) From 1bb7da342019e31c47336a2b9b500e032ca45d0d Mon Sep 17 00:00:00 2001 From: denis-tingaikin Date: Tue, 10 Sep 2024 16:25:07 +0300 Subject: [PATCH 4/4] fix linter Signed-off-by: denis-tingaikin --- pkg/networkservice/chains/client/client.go | 2 +- pkg/networkservice/common/limit/client.go | 2 ++ pkg/registry/common/limit/common.go | 1 + 3 files changed, 4 insertions(+), 1 deletion(-) diff --git a/pkg/networkservice/chains/client/client.go b/pkg/networkservice/chains/client/client.go index 3c0173e90..aca90ef0c 100644 --- a/pkg/networkservice/chains/client/client.go +++ b/pkg/networkservice/chains/client/client.go @@ -1,4 +1,4 @@ -// Copyright (c) 2021-2022 Cisco and/or its affiliates. +// Copyright (c) 2021-2024 Cisco and/or its affiliates. // // SPDX-License-Identifier: Apache-2.0 // diff --git a/pkg/networkservice/common/limit/client.go b/pkg/networkservice/common/limit/client.go index c103c18a4..29160143c 100644 --- a/pkg/networkservice/common/limit/client.go +++ b/pkg/networkservice/common/limit/client.go @@ -30,6 +30,7 @@ import ( "github.com/networkservicemesh/sdk/pkg/tools/log" ) +// Option overrides default values type Option func(c *limitClient) // WithDialLimit sets dial limit @@ -43,6 +44,7 @@ type limitClient struct { dialLimit time.Duration } +// NewClient returns new NetworkServiceClient that limits rpc func NewClient(opts ...Option) networkservice.NetworkServiceClient { ret := &limitClient{ dialLimit: time.Minute, diff --git a/pkg/registry/common/limit/common.go b/pkg/registry/common/limit/common.go index c745daba7..2f3b572c1 100644 --- a/pkg/registry/common/limit/common.go +++ b/pkg/registry/common/limit/common.go @@ -22,6 +22,7 @@ type limitConfig struct { dialLimit time.Duration } +// Option overrides default values type Option func(cfg *limitConfig) // WithDialLimit sets dial time limit