Use an extended timeout in case of reselect requests

pkg/networkservice/common/begin/server.go

@@ -88,8 +88,12 @@ func (b *beginServer) Request(ctx context.Context, request *networkservice.Netwo
 			eventFactoryServer.request != nil && eventFactoryServer.request.Connection != nil {
 			log.FromContext(ctx).Info("Closing connection due to RESELECT_REQUESTED state")
 
+			closeCtx, cancel := context.WithTimeout(context.Background(), b.closeTimeout)
+			defer cancel()
+
 			eventFactoryCtx, eventFactoryCtxCancel := eventFactoryServer.ctxFunc()
-			_, closeErr := next.Server(eventFactoryCtx).Close(eventFactoryCtx, eventFactoryServer.request.Connection)
+			closeCtx = extend.WithValuesFromContext(closeCtx, eventFactoryCtx)
+			_, closeErr := next.Server(closeCtx).Close(closeCtx, eventFactoryServer.request.Connection)
 			if closeErr != nil {
 				log.FromContext(ctx).Errorf("Can't close old connection: %v", closeErr)
 			}

pkg/networkservice/common/begin/server_test.go

@@ -26,6 +26,7 @@ import (
 	"github.com/networkservicemesh/api/pkg/api/networkservice"
 	"github.com/stretchr/testify/require"
 	"go.uber.org/goleak"
+	"google.golang.org/protobuf/types/known/emptypb"
 
 	"github.com/networkservicemesh/sdk/pkg/networkservice/common/begin"
 	"github.com/networkservicemesh/sdk/pkg/networkservice/core/next"
@@ -42,13 +43,18 @@ type waitServer struct {
 
 func (s *waitServer) Request(ctx context.Context, request *networkservice.NetworkServiceRequest) (*networkservice.Connection, error) {
 	time.Sleep(waitTime)
-	s.requestDone.Store(1)
+	s.requestDone.Add(1)
 	return next.Server(ctx).Request(ctx, request)
 }
 
 func (s *waitServer) Close(ctx context.Context, connection *networkservice.Connection) (*empty.Empty, error) {
-	time.Sleep(waitTime)
-	s.closeDone.Store(1)
+	afterCh := time.After(time.Second)
+	select {
+	case <-ctx.Done():
+		return &emptypb.Empty{}, nil
+	case <-afterCh:
+		s.closeDone.Add(1)
+	}
 	return next.Server(ctx).Close(ctx, connection)
 }
 
@@ -82,3 +88,39 @@ func TestBeginWorksWithSmallTimeout(t *testing.T) {
 		return waitSrv.closeDone.Load() == 1
 	}, waitTime*2, time.Millisecond*500)
 }
+
+func TestBeginHasExtendedTimeoutOnReselect(t *testing.T) {
+	t.Cleanup(func() {
+		goleak.VerifyNone(t)
+	})
+	requestCtx, cancel := context.WithTimeout(context.Background(), time.Millisecond*200)
+	defer cancel()
+
+	waitSrv := &waitServer{}
+	server := next.NewNetworkServiceServer(
+		begin.NewServer(),
+		waitSrv,
+	)
+
+	// Make a first request to create an event factory. Begin should make Request only
+	request := testRequest("id")
+	_, err := server.Request(requestCtx, request)
+	require.EqualError(t, err, context.DeadlineExceeded.Error())
+	require.Equal(t, int32(0), waitSrv.requestDone.Load())
+	require.Eventually(t, func() bool {
+		return waitSrv.requestDone.Load() == 1
+	}, waitTime*2, time.Millisecond*500)
+
+	// Make a second request with RESELECT_REQUESTED. Begin should make Close with extended context first and then Request
+	requestCtx, cancel = context.WithTimeout(context.Background(), time.Millisecond*200)
+	defer cancel()
+	newRequest := request.Clone()
+	newRequest.Connection.State = networkservice.State_RESELECT_REQUESTED
+
+	_, err = server.Request(requestCtx, newRequest)
+	require.EqualError(t, err, context.DeadlineExceeded.Error())
+	require.Equal(t, int32(0), waitSrv.closeDone.Load())
+	require.Eventually(t, func() bool {
+		return waitSrv.closeDone.Load() == 1 && waitSrv.requestDone.Load() == 2
+	}, waitTime*4, time.Millisecond*500)
+}