use monitor server provider in reselectcleanup

Signed-off-by: Danil Uzlov <DanilUzlov@yandex.ru>
networkservicemesh · Jun 19, 2023 · 895c836 · 895c836
1 parent 6e7a633
commit 895c836
Show file tree

Hide file tree

Showing 2 changed files with 42 additions and 19 deletions.
diff --git a/pkg/networkservice/chains/endpoint/server.go b/pkg/networkservice/chains/endpoint/server.go
@@ -127,7 +127,7 @@ func NewServer(ctx context.Context, tokenGenerator token.GeneratorFunc, options
 			begin.NewServer(),
 			updatetoken.NewServer(tokenGenerator),
 			opts.authorizeServer,
-			reselectcleanup.NewServer(),
+			reselectcleanup.NewServer(connectionProvider),
 			metadata.NewServer(),
 			timeout.NewServer(ctx),
 			monitorServer,

diff --git a/pkg/networkservice/common/reselectcleanup/server.go b/pkg/networkservice/common/reselectcleanup/server.go
@@ -19,43 +19,66 @@ package reselectcleanup
 import (
 	"context"
 
-	"github.com/edwarnicke/genericsync"
 	"github.com/golang/protobuf/ptypes/empty"
 	"github.com/networkservicemesh/api/pkg/api/networkservice"
 
+	"github.com/networkservicemesh/sdk/pkg/networkservice/common/monitor"
 	"github.com/networkservicemesh/sdk/pkg/networkservice/core/next"
 	"github.com/networkservicemesh/sdk/pkg/tools/log"
 )
 
 type reselectcleanupServer struct {
-	connections genericsync.Map[string, *networkservice.Connection]
+	connectionProvider monitor.ConnectionProvider
 }
 
-func NewServer() networkservice.NetworkServiceServer {
-	return &reselectcleanupServer{}
+func NewServer(connectionProvider monitor.ConnectionProvider) networkservice.NetworkServiceServer {
+	return &reselectcleanupServer{
+		connectionProvider: connectionProvider,
+	}
 }
 
 func (c *reselectcleanupServer) Request(ctx context.Context, request *networkservice.NetworkServiceRequest) (*networkservice.Connection, error) {
-	if request.GetConnection().GetState() == networkservice.State_RESELECT_REQUESTED {
-		oldConnection, ok := c.connections.Load(request.GetConnection().GetId())
-		if ok && oldConnection != nil {
-			log.FromContext(ctx).Info("Closing connection due to RESELECT_REQUESTED state")
-			_, err := next.Server(ctx).Close(ctx, oldConnection)
-			if err != nil {
-				log.FromContext(ctx).Errorf("Can't close old connection: %v", err)
-			}
+	if request.GetConnection().GetState() != networkservice.State_RESELECT_REQUESTED {
+		return next.Server(ctx).Request(ctx, request)
+	}
+
+	conns, err := c.connectionProvider.Find(&networkservice.MonitorScopeSelector{
+		PathSegments: []*networkservice.PathSegment{
+			{
+				Id:   request.GetConnection().GetCurrentPathSegment().GetId(),
+				Name: request.GetConnection().GetCurrentPathSegment().GetName(),
+			},
+		},
+	})
+	if err != nil {
+		log.FromContext(ctx).Errorf("Can't check if we have old connection to close: %v", err)
+		conn, err := next.Server(ctx).Request(ctx, request)
+		if err == nil {
+			conn.State = networkservice.State_UP
 		}
+		return conn, err
 	}
-	conn, err := next.Server(ctx).Request(ctx, request)
+	oldConnection, ok := conns[request.GetConnection().GetId()]
+	if !ok || oldConnection == nil {
+		// most likely the connection has already been closed
+		conn, err := next.Server(ctx).Request(ctx, request)
+		if err == nil {
+			conn.State = networkservice.State_UP
+		}
+		return conn, err
+	}
+	log.FromContext(ctx).Info("Closing connection due to RESELECT_REQUESTED state")
+	_, err = next.Server(ctx).Close(ctx, oldConnection)
 	if err != nil {
-		return nil, err
+		log.FromContext(ctx).Errorf("Can't close old connection: %v", err)
+	}
+	conn, err := next.Server(ctx).Request(ctx, request)
+	if err == nil {
+		conn.State = networkservice.State_UP
 	}
-	conn.State = networkservice.State_UP
-	c.connections.Store(conn.GetId(), conn.Clone())
-	return conn, nil
+	return conn, err
 }
 
 func (c *reselectcleanupServer) Close(ctx context.Context, conn *networkservice.Connection) (*empty.Empty, error) {
-	c.connections.Delete(conn.GetId())
 	return next.Server(ctx).Close(ctx, conn)
 }