begin should let NSE unregister it self even if it doesn't have an event factory

pkg/registry/common/begin/close_server_test.go

@@ -1,4 +1,4 @@
-// Copyright (c) 2022 Cisco and/or its affiliates.
+// Copyright (c) 2022-2023 Cisco and/or its affiliates.
 //
 // SPDX-License-Identifier: Apache-2.0
 //
@@ -18,20 +18,16 @@ package begin_test
 
 import (
 	"context"
-	"sync"
 	"testing"
 
 	"github.com/networkservicemesh/api/pkg/api/registry"
 
 	"github.com/networkservicemesh/sdk/pkg/registry/common/begin"
-	"github.com/networkservicemesh/sdk/pkg/registry/common/null"
 	"github.com/networkservicemesh/sdk/pkg/registry/core/adapters"
 	"github.com/networkservicemesh/sdk/pkg/registry/core/chain"
-	"github.com/networkservicemesh/sdk/pkg/registry/core/next"
 
-	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 	"go.uber.org/goleak"
-	"google.golang.org/protobuf/types/known/emptypb"
 )
 
 func TestCloseServer(t *testing.T) {
@@ -46,48 +42,12 @@ func TestCloseServer(t *testing.T) {
 	conn, err := server.Register(ctx, &registry.NetworkServiceEndpoint{
 		Name: id,
 	})
-	assert.NotNil(t, t, conn)
-	assert.NoError(t, err)
-	assert.Equal(t, conn.GetNetworkServiceLabels()[mark].Labels[mark], mark)
+	require.NotNil(t, conn)
+	require.NoError(t, err)
+	require.Equal(t, conn.GetNetworkServiceLabels()[mark].Labels[mark], mark)
 	conn = conn.Clone()
 	delete(conn.GetNetworkServiceLabels()[mark].Labels, mark)
-	assert.Zero(t, conn.GetNetworkServiceLabels()[mark].Labels[mark])
+	require.Zero(t, conn.GetNetworkServiceLabels()[mark].Labels[mark])
 	_, err = server.Unregister(ctx, conn)
-	assert.NoError(t, err)
-}
-
-func TestDoubleCloseServer(t *testing.T) {
-	t.Cleanup(func() { goleak.VerifyNone(t) })
-	server := chain.NewNetworkServiceEndpointRegistryServer(
-		begin.NewNetworkServiceEndpointRegistryServer(),
-		&doubleCloseServer{t: t, NetworkServiceEndpointRegistryServer: null.NewNetworkServiceEndpointRegistryServer()},
-	)
-	id := "1"
-	ctx, cancel := context.WithCancel(context.Background())
-	defer cancel()
-	conn, err := server.Register(ctx, &registry.NetworkServiceEndpoint{
-		Name: id,
-	})
-	assert.NotNil(t, t, conn)
-	assert.NoError(t, err)
-	conn = conn.Clone()
-	_, err = server.Unregister(ctx, conn)
-	assert.NoError(t, err)
-	_, err = server.Unregister(ctx, conn)
-	assert.NoError(t, err)
-}
-
-type doubleCloseServer struct {
-	t *testing.T
-	sync.Once
-	registry.NetworkServiceEndpointRegistryServer
-}
-
-func (s *doubleCloseServer) Unregister(ctx context.Context, in *registry.NetworkServiceEndpoint) (*emptypb.Empty, error) {
-	count := 1
-	s.Do(func() {
-		count++
-	})
-	assert.Equal(s.t, 2, count, "Close has been called more than once")
-	return next.NetworkServiceEndpointRegistryServer(ctx).Unregister(ctx, in)
+	require.NoError(t, err)
 }

pkg/registry/common/begin/nse_server.go

@@ -92,8 +92,7 @@ func (b *beginNSEServer) Unregister(ctx context.Context, in *registry.NetworkSer
 	}
 	eventFactoryServer, ok := b.Load(id)
 	if !ok {
-		// If we don't have a connection to Close, just let it be
-		return &emptypb.Empty{}, nil
+		return next.NetworkServiceEndpointRegistryServer(ctx).Unregister(ctx, in)
 	}
 	var err error
 	<-eventFactoryServer.executor.AsyncExec(func() {

pkg/registry/common/expire/nse_server.go

@@ -31,6 +31,7 @@ import (
 	"github.com/networkservicemesh/sdk/pkg/registry/common/begin"
 	"github.com/networkservicemesh/sdk/pkg/registry/core/next"
 	"github.com/networkservicemesh/sdk/pkg/tools/clock"
+	"github.com/networkservicemesh/sdk/pkg/tools/extend"
 	"github.com/networkservicemesh/sdk/pkg/tools/log"
 )
 
@@ -97,7 +98,7 @@ func (s *expireNSEServer) Register(ctx context.Context, nse *registry.NetworkSer
 		case <-expireContext.Done():
 			return
 		case <-expireCh:
-			factory.Unregister(begin.CancelContext(expireContext))
+			factory.Unregister(begin.CancelContext(extend.WithValuesFromContext(expireContext, ctx)))
 		}
 	}()