[sdk#1028] Fix Timeout, Expire to handle failed Close, Unregister

pkg/networkservice/common/timeout/server.go

@@ -86,10 +86,15 @@ func (s *timeoutServer) Request(ctx context.Context, request *networkservice.Net
 func (s *timeoutServer) Close(ctx context.Context, conn *networkservice.Connection) (*empty.Empty, error) {
 	logger := log.FromContext(ctx).WithField("timeoutServer", "Close")
 
-	if !s.expireManager.Delete(conn.GetId()) {
+	if s.expireManager.Stop(conn.GetId()) {
+		if _, err := next.Server(ctx).Close(ctx, conn); err != nil {
+			s.expireManager.Start(conn.GetId())
+			return nil, err
+		}
+		s.expireManager.Delete(conn.GetId())
+	} else {
 		logger.Warnf("connection has been already closed: %s", conn.GetId())
-		return new(empty.Empty), nil
 	}
 
-	return next.Server(ctx).Close(ctx, conn)
+	return new(empty.Empty), nil
 }

pkg/networkservice/common/timeout/server_test.go

@@ -247,6 +247,41 @@ func TestTimeoutServer_RefreshFailure(t *testing.T) {
 	require.Condition(t, connServer.validator(0, 1))
 }
 
+func TestTimeoutServer_CloseFailure(t *testing.T) {
+	t.Cleanup(func() { goleak.VerifyNone(t) })
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	clockMock := clockmock.New(ctx)
+	ctx = clock.WithClock(ctx, clockMock)
+
+	connServer := newConnectionsServer(t)
+
+	client := testClient(
+		ctx,
+		null.NewClient(),
+		next.NewNetworkServiceServer(
+			injecterror.NewServer(
+				injecterror.WithRequestErrorTimes(),
+				injecterror.WithCloseErrorTimes(0)),
+			connServer,
+		),
+		tokenTimeout,
+	)
+
+	conn, err := client.Request(ctx, &networkservice.NetworkServiceRequest{})
+	require.NoError(t, err)
+	require.Condition(t, connServer.validator(1, 0))
+
+	_, err = client.Close(ctx, conn)
+	require.Error(t, err)
+	require.Condition(t, connServer.validator(1, 0))
+
+	clockMock.Add(tokenTimeout)
+	require.Eventually(t, connServer.validator(0, 1), testWait, testTick)
+}
+
 type remoteServer struct{}
 
 func (s *remoteServer) Request(ctx context.Context, request *networkservice.NetworkServiceRequest) (*networkservice.Connection, error) {

pkg/registry/common/expire/nse_server.go

@@ -92,10 +92,15 @@ func (s *expireNSEServer) Find(query *registry.NetworkServiceEndpointQuery, serv
 func (s *expireNSEServer) Unregister(ctx context.Context, nse *registry.NetworkServiceEndpoint) (*empty.Empty, error) {
 	logger := log.FromContext(ctx).WithField("expireNSEServer", "Unregister")
 
-	if !s.expireManager.Delete(nse.Name) {
+	if s.expireManager.Stop(nse.Name) {
+		if _, err := next.NetworkServiceEndpointRegistryServer(ctx).Unregister(ctx, nse); err != nil {
+			s.expireManager.Start(nse.Name)
+			return nil, err
+		}
+		s.expireManager.Delete(nse.Name)
+	} else {
 		logger.Warnf("endpoint has been already unregistered: %s", nse.Name)
-		return new(empty.Empty), nil
 	}
 
-	return next.NetworkServiceEndpointRegistryServer(ctx).Unregister(ctx, nse)
+	return new(empty.Empty), nil
 }

pkg/registry/common/expire/nse_server_test.go

@@ -51,6 +51,26 @@ const (
 	testTick      = testWait / 100
 )
 
+func find(ctx context.Context, c registry.NetworkServiceEndpointRegistryClient) (nses []*registry.NetworkServiceEndpoint, err error) {
+	stream, err := c.Find(ctx, &registry.NetworkServiceEndpointQuery{
+		NetworkServiceEndpoint: new(registry.NetworkServiceEndpoint),
+	})
+	if err != nil {
+		return nil, err
+	}
+
+	var nse *registry.NetworkServiceEndpoint
+	for nse, err = stream.Recv(); err == nil; nse, err = stream.Recv() {
+		nses = append(nses, nse)
+	}
+
+	if err != io.EOF {
+		return nil, err
+	}
+
+	return nses, nil
+}
+
 func TestExpireNSEServer_ShouldCorrectlySetExpirationTime_InRemoteCase(t *testing.T) {
 	t.Cleanup(func() { goleak.VerifyNone(t) })
 
@@ -99,17 +119,10 @@ func TestExpireNSEServer_ShouldUseLessExpirationTimeFromInput_AndWork(t *testing
 
 	require.Equal(t, clockMock.Until(resp.ExpirationTime.AsTime()), expireTimeout/2)
 
-	c := adapters.NetworkServiceEndpointServerToClient(mem)
-
 	clockMock.Add(expireTimeout / 2)
 	require.Eventually(t, func() bool {
-		stream, err := c.Find(ctx, &registry.NetworkServiceEndpointQuery{
-			NetworkServiceEndpoint: new(registry.NetworkServiceEndpoint),
-		})
-		require.NoError(t, err)
-
-		_, err = stream.Recv()
-		return err == io.EOF
+		nses, err := find(ctx, adapters.NetworkServiceEndpointServerToClient(mem))
+		return err == nil && len(nses) == 0
 	}, testWait, testTick)
 }
 
@@ -161,35 +174,29 @@ func TestExpireNSEServer_ShouldRemoveNSEAfterExpirationTime(t *testing.T) {
 
 	c := adapters.NetworkServiceEndpointServerToClient(mem)
 
-	stream, err := c.Find(ctx, &registry.NetworkServiceEndpointQuery{
-		NetworkServiceEndpoint: new(registry.NetworkServiceEndpoint),
-	})
-	require.NoError(t, err)
-
-	nse, err := stream.Recv()
+	nses, err := find(ctx, c)
 	require.NoError(t, err)
-	require.Equal(t, nseName, nse.Name)
+	require.Len(t, nses, 1)
+	require.Equal(t, nseName, nses[0].Name)
 
 	clockMock.Add(expireTimeout)
 	require.Eventually(t, func() bool {
-		stream, err = c.Find(ctx, &registry.NetworkServiceEndpointQuery{
-			NetworkServiceEndpoint: new(registry.NetworkServiceEndpoint),
-		})
-		require.NoError(t, err)
-
-		_, err = stream.Recv()
-		return err == io.EOF
+		nses, err = find(ctx, c)
+		return err == nil && len(nses) == 0
 	}, testWait, testTick)
 }
 
 func TestExpireNSEServer_DataRace(t *testing.T) {
 	t.Cleanup(func() { goleak.VerifyNone(t) })
 
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
 	mem := memory.NewNetworkServiceEndpointRegistryServer()
 
 	s := next.NewNetworkServiceEndpointRegistryServer(
 		serialize.NewNetworkServiceEndpointRegistryServer(),
-		expire.NewNetworkServiceEndpointRegistryServer(context.Background(), 0),
+		expire.NewNetworkServiceEndpointRegistryServer(ctx, 0),
 		localbypass.NewNetworkServiceEndpointRegistryServer("tcp://0.0.0.0"),
 		mem,
 	)
@@ -202,16 +209,9 @@ func TestExpireNSEServer_DataRace(t *testing.T) {
 		require.NoError(t, err)
 	}
 
-	c := adapters.NetworkServiceEndpointServerToClient(mem)
-
 	require.Eventually(t, func() bool {
-		stream, err := c.Find(context.Background(), &registry.NetworkServiceEndpointQuery{
-			NetworkServiceEndpoint: new(registry.NetworkServiceEndpoint),
-		})
-		require.NoError(t, err)
-
-		_, err = stream.Recv()
-		return err == io.EOF
+		nses, err := find(context.Background(), adapters.NetworkServiceEndpointServerToClient(mem))
+		return err == nil && len(nses) == 0
 	}, testWait, testTick)
 }
 
@@ -241,13 +241,56 @@ func TestExpireNSEServer_RefreshFailure(t *testing.T) {
 
 	clockMock.Add(expireTimeout)
 	require.Eventually(t, func() bool {
-		stream, err := c.Find(ctx, &registry.NetworkServiceEndpointQuery{
-			NetworkServiceEndpoint: new(registry.NetworkServiceEndpoint),
-		})
-		require.NoError(t, err)
+		nses, err := find(ctx, c)
+		return err == nil && len(nses) == 0
+	}, testWait, testTick)
+}
+
+func TestExpireNSEServer_UnregisterFailure(t *testing.T) {
+	t.Cleanup(func() { goleak.VerifyNone(t) })
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	clockMock := clockmock.New(ctx)
+	ctx = clock.WithClock(ctx, clockMock)
+
+	mem := memory.NewNetworkServiceEndpointRegistryServer()
+
+	s := next.NewNetworkServiceEndpointRegistryServer(
+		serialize.NewNetworkServiceEndpointRegistryServer(),
+		expire.NewNetworkServiceEndpointRegistryServer(ctx, expireTimeout),
+		new(remoteNSEServer), // <-- GRPC invocation
+		mem,
+	)
 
-		_, err = stream.Recv()
-		return err == io.EOF
+	nse, err := s.Register(ctx, &registry.NetworkServiceEndpoint{
+		Name: nseName,
+	})
+	require.NoError(t, err)
+
+	c := adapters.NetworkServiceEndpointServerToClient(mem)
+
+	nses, err := find(ctx, c)
+	require.NoError(t, err)
+	require.Len(t, nses, 1)
+	require.Equal(t, nseName, nses[0].Name)
+
+	unregisterCtx, cancelUnregister := context.WithCancel(ctx)
+	cancelUnregister()
+
+	_, err = s.Unregister(unregisterCtx, nse)
+	require.Error(t, err)
+
+	nses, err = find(ctx, c)
+	require.NoError(t, err)
+	require.Len(t, nses, 1)
+	require.Equal(t, nseName, nses[0].Name)
+
+	clockMock.Add(expireTimeout)
+	require.Eventually(t, func() bool {
+		nses, err = find(ctx, c)
+		return err == nil && len(nses) == 0
 	}, testWait, testTick)
 }
 
@@ -295,14 +338,23 @@ type remoteNSEServer struct {
 }
 
 func (s *remoteNSEServer) Register(ctx context.Context, nse *registry.NetworkServiceEndpoint) (*registry.NetworkServiceEndpoint, error) {
+	if err := ctx.Err(); err != nil {
+		return nil, err
+	}
 	return next.NetworkServiceEndpointRegistryServer(ctx).Register(ctx, nse.Clone())
 }
 
 func (s *remoteNSEServer) Find(query *registry.NetworkServiceEndpointQuery, server registry.NetworkServiceEndpointRegistry_FindServer) error {
+	if err := server.Context().Err(); err != nil {
+		return err
+	}
 	return next.NetworkServiceEndpointRegistryServer(server.Context()).Find(query, server)
 }
 
 func (s *remoteNSEServer) Unregister(ctx context.Context, nse *registry.NetworkServiceEndpoint) (*empty.Empty, error) {
+	if err := ctx.Err(); err != nil {
+		return nil, err
+	}
 	return next.NetworkServiceEndpointRegistryServer(ctx).Unregister(ctx, nse.Clone())
 }
 

pkg/tools/expire/manager.go

@@ -79,10 +79,12 @@ func (m *Manager) New(executor Executor, id string, expirationTime time.Time, cl
 }
 
 // Stop stops expiration for the `id`
-func (m *Manager) Stop(id string) {
-	if t, ok := m.timers.Load(id); ok {
+func (m *Manager) Stop(id string) bool {
+	t, loaded := m.timers.Load(id)
+	if loaded {
 		t.stopped = t.Stop()
 	}
+	return loaded
 }
 
 // Start starts stopped expiration for the `id` with the same expiration time