create postpne ctx before next.Request() call

Signed-off-by: Periyasamy Palanisamy <periyasamy.palanisamy@est.tech>
networkservicemesh · Aug 31, 2021 · e032334 · e032334
1 parent 506ef66
commit e032334
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 3 deletions.
diff --git a/pkg/networkservice/l2ovsconnect/client.go b/pkg/networkservice/l2ovsconnect/client.go
@@ -44,12 +44,14 @@ func NewClient(bridgeName string) networkservice.NetworkServiceClient {
 
 func (c *l2ConnectClient) Request(ctx context.Context, request *networkservice.NetworkServiceRequest, opts ...grpc.CallOption) (*networkservice.Connection, error) {
 	logger := log.FromContext(ctx).WithField("l2ConnectClient", "Request")
+
+	postponeCtxFunc := postpone.ContextWithValues(ctx)
+
 	conn, err := next.Client(ctx).Request(ctx, request, opts...)
 	if err != nil || request.GetConnection().GetNextPathSegment() != nil {
 		return conn, err
 	}
 
-	postponeCtxFunc := postpone.ContextWithValues(ctx)
 	if err := addDel(ctx, logger, c.bridgeName, true); err != nil {
 		closeCtx, cancelClose := postponeCtxFunc()
 		defer cancelClose()

diff --git a/pkg/networkservice/mechanisms/kernel/client.go b/pkg/networkservice/mechanisms/kernel/client.go
@@ -49,15 +49,19 @@ func NewClient(bridgeName string) networkservice.NetworkServiceClient {
 
 func (c *kernelClient) Request(ctx context.Context, request *networkservice.NetworkServiceRequest, opts ...grpc.CallOption) (*networkservice.Connection, error) {
 	logger := log.FromContext(ctx).WithField("kernelClient", "Request")
+
 	request.MechanismPreferences = append(request.MechanismPreferences, &networkservice.Mechanism{
 		Cls:  cls.LOCAL,
 		Type: kernel.MECHANISM,
 	})
+
+	postponeCtxFunc := postpone.ContextWithValues(ctx)
+
 	conn, err := next.Client(ctx).Request(ctx, request, opts...)
 	if err != nil || request.GetConnection().GetNextPathSegment() != nil {
 		return conn, err
 	}
-	postponeCtxFunc := postpone.ContextWithValues(ctx)
+
 	_, exists := conn.GetMechanism().GetParameters()[common.PCIAddressKey]
 	if exists {
 		if err = setupVF(ctx, logger, conn, c.bridgeName, metadata.IsClient(c)); err != nil {
@@ -76,6 +80,7 @@ func (c *kernelClient) Request(ctx context.Context, request *networkservice.Netw
 			}
 		}
 	}
+
 	return conn, err
 }
 

diff --git a/pkg/networkservice/mechanisms/vxlan/client.go b/pkg/networkservice/mechanisms/vxlan/client.go
@@ -53,22 +53,27 @@ func NewClient(tunnelIP net.IP, bridgeName string, mutex sync.Locker, vxlanRefCo
 
 func (c *vxlanClient) Request(ctx context.Context, request *networkservice.NetworkServiceRequest, opts ...grpc.CallOption) (*networkservice.Connection, error) {
 	logger := log.FromContext(ctx).WithField("vxlanClient", "Request")
+
 	request.MechanismPreferences = append(request.MechanismPreferences, &networkservice.Mechanism{
 		Cls:  cls.REMOTE,
 		Type: vxlan.MECHANISM,
 	})
+
+	postponeCtxFunc := postpone.ContextWithValues(ctx)
+
 	conn, err := next.Client(ctx).Request(ctx, request, opts...)
 	if err != nil || request.GetConnection().GetNextPathSegment() != nil {
 		return conn, err
 	}
-	postponeCtxFunc := postpone.ContextWithValues(ctx)
+
 	if err = add(ctx, logger, conn, c.bridgeName, c.vxlanInterfacesMutex, c.vxlanInterfacesMap, true); err != nil {
 		closeCtx, cancelClose := postponeCtxFunc()
 		defer cancelClose()
 		if _, closeErr := c.Close(closeCtx, conn, opts...); closeErr != nil {
 			logger.Errorf("failed to close failed connection: %s %s", conn.GetId(), closeErr.Error())
 		}
 	}
+
 	return conn, err
 }