begin: update context after a successful request

pkg/networkservice/common/begin/client.go

@@ -65,7 +65,6 @@ func (b *beginClient) Request(ctx context.Context, request *networkservice.Netwo
 			conn, err = b.Request(ctx, request, opts...)
 			return
 		}
-		eventFactoryClient.updateContext(ctx)
 
 		ctx = withEventFactory(ctx, eventFactoryClient)
 		request.Connection = mergeConnection(eventFactoryClient.returnedConnection, request.GetConnection(), eventFactoryClient.request.GetConnection())
@@ -83,6 +82,7 @@ func (b *beginClient) Request(ctx context.Context, request *networkservice.Netwo
 		eventFactoryClient.state = established
 
 		eventFactoryClient.returnedConnection = conn.Clone()
+		eventFactoryClient.updateContext(ctx)
 	})
 	return conn, err
 }

pkg/networkservice/common/begin/event_factory_client_test.go

@@ -21,6 +21,8 @@ import (
 	"testing"
 	"time"
 
+	"github.com/pkg/errors"
+
 	"github.com/stretchr/testify/assert"
 	"go.uber.org/goleak"
 	"google.golang.org/grpc"
@@ -38,13 +40,13 @@ import (
 func TestRefresh_Client(t *testing.T) {
 	t.Cleanup(func() { goleak.VerifyNone(t) })
 
-	syncChan := make(chan struct{})
 	checkCtxCl := &checkContextClient{t: t}
-	eventFactoryCl := &eventFactoryClient{ch: syncChan}
+	eventFactoryCl := &eventFactoryClient{}
 	client := chain.NewNetworkServiceClient(
 		begin.NewClient(),
 		checkCtxCl,
 		eventFactoryCl,
+		&failedNSEClient{},
 	)
 
 	ctx, cancel := context.WithCancel(context.Background())
@@ -62,27 +64,36 @@ func TestRefresh_Client(t *testing.T) {
 
 	// Change context value before refresh Request
 	ctx = context.WithValue(ctx, contextKey{}, "value_2")
-	checkCtxCl.setExpectedValue("value_2")
+
+	// Call refresh that will fail
 	request.Connection = conn.Clone()
+	request.Connection.NetworkServiceEndpointName = failedNSENameClient
+	checkCtxCl.setExpectedValue("value_2")
+	_, err = client.Request(ctx, request.Clone())
+	assert.Error(t, err)
 
-	// Call refresh
+	// Call refresh from eventFactory. We are expecting the previous value in the context
+	checkCtxCl.setExpectedValue("value_1")
+	eventFactoryCl.callRefresh()
+
+	// Call refresh that will successful
+	request.Connection.NetworkServiceEndpointName = ""
+	checkCtxCl.setExpectedValue("value_2")
 	conn, err = client.Request(ctx, request.Clone())
 	assert.NotNil(t, t, conn)
 	assert.NoError(t, err)
 
 	// Call refresh from eventFactory. We are expecting updated value in the context
 	eventFactoryCl.callRefresh()
-	<-syncChan
 }
 
 // This test reproduces the situation when Close and Request were called at the same time
 // nolint:dupl
 func TestRefreshDuringClose_Client(t *testing.T) {
 	t.Cleanup(func() { goleak.VerifyNone(t) })
 
-	syncChan := make(chan struct{})
 	checkCtxCl := &checkContextClient{t: t}
-	eventFactoryCl := &eventFactoryClient{ch: syncChan}
+	eventFactoryCl := &eventFactoryClient{}
 	client := chain.NewNetworkServiceClient(
 		begin.NewClient(),
 		checkCtxCl,
@@ -109,7 +120,6 @@ func TestRefreshDuringClose_Client(t *testing.T) {
 
 	// Call Close from eventFactory
 	eventFactoryCl.callClose()
-	<-syncChan
 
 	// Call refresh  (should be called at the same time as Close)
 	conn, err = client.Request(ctx, request.Clone())
@@ -118,12 +128,10 @@ func TestRefreshDuringClose_Client(t *testing.T) {
 
 	// Call refresh from eventFactory. We are expecting updated value in the context
 	eventFactoryCl.callRefresh()
-	<-syncChan
 }
 
 type eventFactoryClient struct {
 	ctx context.Context
-	ch  chan<- struct{}
 }
 
 func (s *eventFactoryClient) Request(ctx context.Context, request *networkservice.NetworkServiceRequest, opts ...grpc.CallOption) (*networkservice.Connection, error) {
@@ -139,18 +147,12 @@ func (s *eventFactoryClient) Close(ctx context.Context, conn *networkservice.Con
 
 func (s *eventFactoryClient) callClose() {
 	eventFactory := begin.FromContext(s.ctx)
-	go func() {
-		s.ch <- struct{}{}
-		eventFactory.Close()
-	}()
+	eventFactory.Close()
 }
 
 func (s *eventFactoryClient) callRefresh() {
 	eventFactory := begin.FromContext(s.ctx)
-	go func() {
-		s.ch <- struct{}{}
-		eventFactory.Request()
-	}()
+	<-eventFactory.Request()
 }
 
 type contextKey struct{}
@@ -172,3 +174,21 @@ func (c *checkContextClient) Close(ctx context.Context, conn *networkservice.Con
 func (c *checkContextClient) setExpectedValue(value string) {
 	c.expectedValue = value
 }
+
+const failedNSENameClient = "failedNSE"
+
+type failedNSEClient struct{}
+
+func (f *failedNSEClient) Request(ctx context.Context, request *networkservice.NetworkServiceRequest, opts ...grpc.CallOption) (*networkservice.Connection, error) {
+	if request.GetConnection().NetworkServiceEndpointName == failedNSENameClient {
+		return nil, errors.New("failed")
+	}
+	return next.Client(ctx).Request(ctx, request, opts...)
+}
+
+func (f *failedNSEClient) Close(ctx context.Context, conn *networkservice.Connection, opts ...grpc.CallOption) (*emptypb.Empty, error) {
+	if conn.NetworkServiceEndpointName == failedNSENameClient {
+		return nil, errors.New("failed")
+	}
+	return next.Client(ctx).Close(ctx, conn, opts...)
+}

pkg/networkservice/common/begin/event_factory_server_test.go

@@ -21,6 +21,8 @@ import (
 	"testing"
 	"time"
 
+	"github.com/pkg/errors"
+
 	"github.com/stretchr/testify/assert"
 	"go.uber.org/goleak"
 	"google.golang.org/protobuf/types/known/emptypb"
@@ -37,13 +39,13 @@ import (
 func TestRefresh_Server(t *testing.T) {
 	t.Cleanup(func() { goleak.VerifyNone(t) })
 
-	syncChan := make(chan struct{})
 	checkCtxServ := &checkContextServer{t: t}
-	eventFactoryServ := &eventFactoryServer{ch: syncChan}
+	eventFactoryServ := &eventFactoryServer{}
 	server := chain.NewNetworkServiceServer(
 		begin.NewServer(),
 		checkCtxServ,
 		eventFactoryServ,
+		&failedNSEServer{},
 	)
 
 	ctx, cancel := context.WithCancel(context.Background())
@@ -61,27 +63,36 @@ func TestRefresh_Server(t *testing.T) {
 
 	// Change context value before refresh Request
 	ctx = context.WithValue(ctx, contextKey{}, "value_2")
-	checkCtxServ.setExpectedValue("value_2")
+
+	// Call refresh that will fail
 	request.Connection = conn.Clone()
+	request.Connection.NetworkServiceEndpointName = failedNSENameServer
+	checkCtxServ.setExpectedValue("value_2")
+	_, err = server.Request(ctx, request.Clone())
+	assert.Error(t, err)
 
-	// Call refresh
+	// Call refresh from eventFactory. We are expecting the previous value in the context
+	checkCtxServ.setExpectedValue("value_1")
+	eventFactoryServ.callRefresh()
+
+	// Call refresh that will successful
+	request.Connection.NetworkServiceEndpointName = ""
+	checkCtxServ.setExpectedValue("value_2")
 	conn, err = server.Request(ctx, request.Clone())
 	assert.NotNil(t, t, conn)
 	assert.NoError(t, err)
 
 	// Call refresh from eventFactory. We are expecting updated value in the context
 	eventFactoryServ.callRefresh()
-	<-syncChan
 }
 
 // This test reproduces the situation when Close and Request were called at the same time
 // nolint:dupl
 func TestRefreshDuringClose_Server(t *testing.T) {
 	t.Cleanup(func() { goleak.VerifyNone(t) })
 
-	syncChan := make(chan struct{})
 	checkCtxServ := &checkContextServer{t: t}
-	eventFactoryServ := &eventFactoryServer{ch: syncChan}
+	eventFactoryServ := &eventFactoryServer{}
 	server := chain.NewNetworkServiceServer(
 		begin.NewServer(),
 		checkCtxServ,
@@ -108,7 +119,6 @@ func TestRefreshDuringClose_Server(t *testing.T) {
 
 	// Call Close from eventFactory
 	eventFactoryServ.callClose()
-	<-syncChan
 
 	// Call refresh  (should be called at the same time as Close)
 	conn, err = server.Request(ctx, request.Clone())
@@ -117,12 +127,10 @@ func TestRefreshDuringClose_Server(t *testing.T) {
 
 	// Call refresh from eventFactory. We are expecting updated value in the context
 	eventFactoryServ.callRefresh()
-	<-syncChan
 }
 
 type eventFactoryServer struct {
 	ctx context.Context
-	ch  chan<- struct{}
 }
 
 func (e *eventFactoryServer) Request(ctx context.Context, request *networkservice.NetworkServiceRequest) (*networkservice.Connection, error) {
@@ -138,18 +146,12 @@ func (e *eventFactoryServer) Close(ctx context.Context, conn *networkservice.Con
 
 func (e *eventFactoryServer) callClose() {
 	eventFactory := begin.FromContext(e.ctx)
-	go func() {
-		e.ch <- struct{}{}
-		eventFactory.Close()
-	}()
+	eventFactory.Close()
 }
 
 func (e *eventFactoryServer) callRefresh() {
 	eventFactory := begin.FromContext(e.ctx)
-	go func() {
-		e.ch <- struct{}{}
-		eventFactory.Request()
-	}()
+	<-eventFactory.Request()
 }
 
 type checkContextServer struct {
@@ -169,3 +171,21 @@ func (c *checkContextServer) Close(ctx context.Context, conn *networkservice.Con
 func (c *checkContextServer) setExpectedValue(value string) {
 	c.expectedValue = value
 }
+
+const failedNSENameServer = "failedNSE"
+
+type failedNSEServer struct{}
+
+func (f *failedNSEServer) Request(ctx context.Context, request *networkservice.NetworkServiceRequest) (*networkservice.Connection, error) {
+	if request.GetConnection().NetworkServiceEndpointName == failedNSENameServer {
+		return nil, errors.New("failed")
+	}
+	return next.Server(ctx).Request(ctx, request)
+}
+
+func (f *failedNSEServer) Close(ctx context.Context, conn *networkservice.Connection) (*emptypb.Empty, error) {
+	if conn.NetworkServiceEndpointName == failedNSENameServer {
+		return nil, errors.New("failed")
+	}
+	return next.Server(ctx).Close(ctx, conn)
+}

pkg/networkservice/common/begin/server.go

@@ -61,7 +61,6 @@ func (b *beginServer) Request(ctx context.Context, request *networkservice.Netwo
 			conn, err = b.Request(ctx, request)
 			return
 		}
-		eventFactoryServer.updateContext(ctx)
 
 		ctx = withEventFactory(ctx, eventFactoryServer)
 		conn, err = next.Server(ctx).Request(ctx, request)
@@ -77,6 +76,7 @@ func (b *beginServer) Request(ctx context.Context, request *networkservice.Netwo
 		eventFactoryServer.state = established
 
 		eventFactoryServer.returnedConnection = conn.Clone()
+		eventFactoryServer.updateContext(ctx)
 	})
 	return conn, err
 }

pkg/registry/common/begin/ns_client.go

@@ -61,7 +61,6 @@ func (b *beginNSClient) Register(ctx context.Context, in *registry.NetworkServic
 			resp, err = b.Register(ctx, in, opts...)
 			return
 		}
-		eventFactoryClient.updateContext(ctx)
 
 		ctx = withEventFactory(ctx, eventFactoryClient)
 		resp, err = next.NetworkServiceRegistryClient(ctx).Register(ctx, in, opts...)
@@ -76,6 +75,7 @@ func (b *beginNSClient) Register(ctx context.Context, in *registry.NetworkServic
 		eventFactoryClient.state = established
 		eventFactoryClient.registration = mergeNS(in, resp.Clone())
 		eventFactoryClient.response = resp.Clone()
+		eventFactoryClient.updateContext(ctx)
 	})
 	return resp, err
 }

pkg/registry/common/begin/ns_server.go

@@ -60,7 +60,6 @@ func (b *beginNSServer) Register(ctx context.Context, in *registry.NetworkServic
 			resp, err = b.Register(ctx, in)
 			return
 		}
-		eventFactoryServer.updateContext(ctx)
 
 		ctx = withEventFactory(ctx, eventFactoryServer)
 		resp, err = next.NetworkServiceRegistryServer(ctx).Register(ctx, in)
@@ -74,6 +73,7 @@ func (b *beginNSServer) Register(ctx context.Context, in *registry.NetworkServic
 		eventFactoryServer.registration = mergeNS(in, resp)
 		eventFactoryServer.state = established
 		eventFactoryServer.response = resp
+		eventFactoryServer.updateContext(ctx)
 	})
 	return resp, err
 }

pkg/registry/common/begin/nse_client.go

@@ -61,7 +61,6 @@ func (b *beginNSEClient) Register(ctx context.Context, in *registry.NetworkServi
 			resp, err = b.Register(ctx, in, opts...)
 			return
 		}
-		eventFactoryClient.updateContext(ctx)
 
 		ctx = withEventFactory(ctx, eventFactoryClient)
 		resp, err = next.NetworkServiceEndpointRegistryClient(ctx).Register(ctx, in, opts...)
@@ -76,6 +75,7 @@ func (b *beginNSEClient) Register(ctx context.Context, in *registry.NetworkServi
 		eventFactoryClient.state = established
 		eventFactoryClient.registration = mergeNSE(in, resp.Clone())
 		eventFactoryClient.response = resp.Clone()
+		eventFactoryClient.updateContext(ctx)
 	})
 	return resp, err
 }