fix: expire chain element could work incorrectly with registry-k8s

pkg/registry/common/expire/ns_server.go

@@ -1,4 +1,4 @@
-// Copyright (c) 2020 Doc.ai and/or its affiliates.
+// Copyright (c) 2020-2021 Doc.ai and/or its affiliates.
 //
 // SPDX-License-Identifier: Apache-2.0
 //
@@ -24,6 +24,7 @@ import (
 	"time"
 
 	"github.com/networkservicemesh/sdk/pkg/registry/core/next"
+	"github.com/networkservicemesh/sdk/pkg/tools/extend"
 
 	"github.com/golang/protobuf/ptypes/empty"
 	"github.com/networkservicemesh/api/pkg/api/registry"
@@ -36,11 +37,11 @@ type nsServer struct {
 	nsCounts   intMap
 	contexts   contextMap
 	once       sync.Once
-	ctx        context.Context
+	chainCtx   context.Context
 }
 
 func (n *nsServer) checkUpdates() {
-	c, err := n.nseClient.Find(n.ctx, &registry.NetworkServiceEndpointQuery{
+	c, err := n.nseClient.Find(n.chainCtx, &registry.NetworkServiceEndpointQuery{
 		NetworkServiceEndpoint: &registry.NetworkServiceEndpoint{},
 		Watch:                  true,
 	})
@@ -77,15 +78,15 @@ func (n *nsServer) checkUpdates() {
 func (n *nsServer) Register(ctx context.Context, request *registry.NetworkService) (*registry.NetworkService, error) {
 	n.once.Do(func() {
 		go func() {
-			for n.monitorErr == nil && n.ctx.Err() == nil {
+			for n.monitorErr == nil && n.chainCtx.Err() == nil {
 				n.checkUpdates()
 			}
 		}()
 	})
 	if n.monitorErr != nil {
 		return nil, n.monitorErr
 	}
-	n.contexts.Store(request.Name, ctx)
+	n.contexts.Store(request.Name, extend.WithValuesFromContext(n.chainCtx, ctx))
 	return next.NetworkServiceRegistryServer(ctx).Register(ctx, request)
 }
 
@@ -111,5 +112,5 @@ func (n *nsServer) Unregister(ctx context.Context, request *registry.NetworkServ
 
 // NewNetworkServiceServer wraps passed NetworkServiceRegistryServer and monitor NetworkServiceEndpoints via passed NetworkServiceEndpointRegistryClient
 func NewNetworkServiceServer(ctx context.Context, nseClient registry.NetworkServiceEndpointRegistryClient) registry.NetworkServiceRegistryServer {
-	return &nsServer{nseClient: nseClient, ctx: ctx}
+	return &nsServer{nseClient: nseClient, chainCtx: ctx}
 }

pkg/registry/common/expire/nse_server.go

@@ -41,10 +41,11 @@ func (n *nseServer) Register(ctx context.Context, nse *registry.NetworkServiceEn
 	}
 	expirationTime := time.Now().Add(n.nseExpiration)
 	nse.ExpirationTime = &timestamp.Timestamp{Seconds: expirationTime.Unix(), Nanos: int32(expirationTime.Nanosecond())}
+	unregisterNse := r.Clone()
 	timer := time.AfterFunc(n.nseExpiration, func() {
 		unregisterCtx, cancel := context.WithTimeout(extend.WithValuesFromContext(context.Background(), ctx), n.nseExpiration)
 		defer cancel()
-		_, _ = next.NetworkServiceEndpointRegistryServer(unregisterCtx).Unregister(unregisterCtx, nse)
+		_, _ = next.NetworkServiceEndpointRegistryServer(unregisterCtx).Unregister(unregisterCtx, unregisterNse)
 	})
 	if t, load := n.timers.LoadOrStore(nse.Name, timer); load {
 		timer.Stop()

pkg/registry/common/expire/nse_server_test.go

@@ -1,4 +1,4 @@
-// Copyright (c) 2020 Doc.ai and/or its affiliates.
+// Copyright (c) 2020-2021 Doc.ai and/or its affiliates.
 //
 // SPDX-License-Identifier: Apache-2.0
 //
@@ -24,6 +24,7 @@ import (
 	"go.uber.org/goleak"
 
 	"github.com/networkservicemesh/sdk/pkg/registry/common/memory"
+	"github.com/networkservicemesh/sdk/pkg/registry/common/refresh"
 	"github.com/networkservicemesh/sdk/pkg/registry/core/next"
 
 	"github.com/networkservicemesh/api/pkg/api/registry"
@@ -54,3 +55,37 @@ func TestNewNetworkServiceEndpointRegistryServer(t *testing.T) {
 		return len(list) == 0
 	}, time.Second, time.Millisecond*100)
 }
+
+func Test_ExpireEndpointRegistryServer_ShouldCorrectlyRescheduleTimer(t *testing.T) {
+	defer goleak.VerifyNone(t, goleak.IgnoreCurrent())
+
+	ctx, cancel := context.WithCancel(context.Background())
+
+	s := next.NewNetworkServiceEndpointRegistryServer(expire.NewNetworkServiceEndpointRegistryServer(testPeriod*2), memory.NewNetworkServiceEndpointRegistryServer())
+	c := next.NewNetworkServiceEndpointRegistryClient(refresh.NewNetworkServiceEndpointRegistryClient(refresh.WithChainContext(ctx)), adapters.NetworkServiceEndpointServerToClient(s))
+
+	_, err := c.Register(context.Background(), &registry.NetworkServiceEndpoint{})
+	require.NoError(t, err)
+
+	deadline := time.Now().Add(time.Second)
+
+	for time.Until(deadline) > 0 {
+		stream, err := c.Find(context.Background(), &registry.NetworkServiceEndpointQuery{
+			NetworkServiceEndpoint: &registry.NetworkServiceEndpoint{},
+		})
+		require.NoError(t, err)
+		list := registry.ReadNetworkServiceEndpointList(stream)
+		require.Len(t, list, 1)
+	}
+
+	cancel()
+
+	require.Eventually(t, func() bool {
+		stream, err := c.Find(context.Background(), &registry.NetworkServiceEndpointQuery{
+			NetworkServiceEndpoint: &registry.NetworkServiceEndpoint{},
+		})
+		require.Nil(t, err)
+		list := registry.ReadNetworkServiceEndpointList(stream)
+		return len(list) == 0
+	}, time.Second, time.Millisecond*100)
+}

pkg/registry/common/refresh/nse_registry_client.go

@@ -1,4 +1,4 @@
-// Copyright (c) 2020 Doc.ai and/or its affiliates.
+// Copyright (c) 2020-2021 Doc.ai and/or its affiliates.
 //
 // SPDX-License-Identifier: Apache-2.0
 //
@@ -30,6 +30,7 @@ import (
 )
 
 type refreshNSEClient struct {
+	chainContext          context.Context
 	nseCancels            cancelsMap
 	retryDelay            time.Duration
 	defaultExpiryDuration time.Duration
@@ -75,7 +76,7 @@ func (c *refreshNSEClient) Register(ctx context.Context, in *registry.NetworkSer
 	if cancel, ok := c.nseCancels.Load(resp.Name); ok {
 		cancel()
 	}
-	ctx, cancel := context.WithCancel(context.Background())
+	ctx, cancel := context.WithCancel(c.chainContext)
 	c.nseCancels.Store(resp.Name, cancel)
 	c.startRefresh(ctx, nextClient, resp)
 	return resp, err
@@ -102,6 +103,7 @@ func NewNetworkServiceEndpointRegistryClient(options ...Option) registry.Network
 	c := &refreshNSEClient{
 		retryDelay:            time.Second * 5,
 		defaultExpiryDuration: time.Minute * 30,
+		chainContext:          context.Background(),
 	}
 
 	for _, o := range options {

pkg/registry/common/refresh/options.go

@@ -1,4 +1,4 @@
-// Copyright (c) 2020 Doc.ai and/or its affiliates.
+// Copyright (c) 2020-2021 Doc.ai and/or its affiliates.
 //
 // SPDX-License-Identifier: Apache-2.0
 //
@@ -16,7 +16,10 @@
 
 package refresh
 
-import "time"
+import (
+	"context"
+	"time"
+)
 
 // Option is expire registry configuration option
 type Option interface {
@@ -42,3 +45,10 @@ func WithDefaultExpiryDuration(duration time.Duration) Option {
 		c.defaultExpiryDuration = duration
 	})
 }
+
+// WithChainContext sets a chain context
+func WithChainContext(ctx context.Context) Option {
+	return applierFunc(func(c *refreshNSEClient) {
+		c.chainContext = ctx
+	})
+}