[sdk#1026] Add closectx

pkg/networkservice/common/authorize/client.go

@@ -23,12 +23,14 @@ import (
 	"sync/atomic"
 
 	"github.com/golang/protobuf/ptypes/empty"
+	"github.com/pkg/errors"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/peer"
 
 	"github.com/networkservicemesh/api/pkg/api/networkservice"
 
 	"github.com/networkservicemesh/sdk/pkg/networkservice/core/next"
+	"github.com/networkservicemesh/sdk/pkg/tools/postpone"
 )
 
 type authorizeClient struct {
@@ -41,37 +43,53 @@ func NewClient(opts ...Option) networkservice.NetworkServiceClient {
 	var result = &authorizeClient{
 		policies: defaultPolicies(),
 	}
+
 	for _, o := range opts {
 		o.apply(&result.policies)
 	}
+
 	return result
 }
 
 func (a *authorizeClient) Request(ctx context.Context, request *networkservice.NetworkServiceRequest, opts ...grpc.CallOption) (*networkservice.Connection, error) {
 	var p peer.Peer
 	opts = append(opts, grpc.Peer(&p))
-	resp, err := next.Client(ctx).Request(ctx, request, opts...)
+
+	postponeCtxFunc := postpone.ContextWithValues(ctx)
+
+	conn, err := next.Client(ctx).Request(ctx, request, opts...)
 	if err != nil {
 		return nil, err
 	}
+
 	if p != (peer.Peer{}) {
 		a.serverPeer.Store(&p)
 		ctx = peer.NewContext(ctx, &p)
 	}
-	if err = a.policies.check(ctx, resp); err != nil {
-		_, _ = next.Client(ctx).Close(ctx, resp, opts...)
+
+	if err = a.policies.check(ctx, conn); err != nil {
+		closeCtx, cancelClose := postponeCtxFunc()
+		defer cancelClose()
+
+		if _, closeErr := next.Client(ctx).Close(closeCtx, conn, opts...); closeErr != nil {
+			err = errors.Wrapf(err, "connection closed with error: %s", closeErr.Error())
+		}
+
 		return nil, err
 	}
-	return resp, err
+
+	return conn, err
 }
 
 func (a *authorizeClient) Close(ctx context.Context, conn *networkservice.Connection, opts ...grpc.CallOption) (*empty.Empty, error) {
 	p, ok := a.serverPeer.Load().(*peer.Peer)
 	if ok && p != nil {
 		ctx = peer.NewContext(ctx, p)
 	}
+
 	if err := a.policies.check(ctx, conn); err != nil {
 		return nil, err
 	}
+
 	return next.Client(ctx).Close(ctx, conn, opts...)
 }

pkg/networkservice/common/connect/server.go

@@ -35,6 +35,7 @@ import (
 	"github.com/networkservicemesh/sdk/pkg/tools/clienturlctx"
 	"github.com/networkservicemesh/sdk/pkg/tools/log"
 	"github.com/networkservicemesh/sdk/pkg/tools/multiexecutor"
+	"github.com/networkservicemesh/sdk/pkg/tools/postpone"
 )
 
 // ClientFactory is used to created new clients when new connection is created.
@@ -88,13 +89,18 @@ func (s *connectServer) Request(ctx context.Context, request *networkservice.Net
 		return nil, errors.Errorf("clientURL not found for incoming connection: %+v", request.GetConnection())
 	}
 
-	c := s.client(ctx, request.GetConnection())
-	if err := c.client.ctx.Err(); err != nil {
+	c, err := s.client(ctx, request.GetConnection())
+	if err != nil {
+		return nil, err
+	}
+	if err = c.client.ctx.Err(); err != nil {
 		s.deleteClient(c, clientURL.String())
 		s.connInfos.Delete(request.GetConnection().GetId())
 		return nil, err
 	}
 
+	postponeCtxFunc := postpone.ContextWithValues(ctx)
+
 	conn, err := c.client.Request(ctx, request.Clone())
 	if err != nil {
 		if _, ok := s.connInfos.Load(request.GetConnection().GetId()); !ok {
@@ -116,9 +122,12 @@ func (s *connectServer) Request(ctx context.Context, request *networkservice.Net
 	conn, err = next.Server(ctx).Request(ctx, request)
 	// Close connection if next.Server Request finished with error
 	if err != nil && !refreshRequest {
-		_, cErr := s.Close(ctx, request.Connection.Clone())
-		if cErr != nil {
-			err = errors.Wrapf(cErr, "connection closed with error: %v", cErr)
+		closeCtx, cancelClose := postponeCtxFunc()
+		defer cancelClose()
+
+		_, closeErr := s.Close(closeCtx, request.Connection.Clone())
+		if closeErr != nil {
+			err = errors.Wrapf(err, "connection closed with error: %s", closeErr.Error())
 		}
 	}
 	return conn, err
@@ -144,19 +153,23 @@ func (s *connectServer) Close(ctx context.Context, conn *networkservice.Connecti
 	return &empty.Empty{}, err
 }
 
-func (s *connectServer) client(ctx context.Context, conn *networkservice.Connection) *clientInfo {
+func (s *connectServer) client(ctx context.Context, conn *networkservice.Connection) (*clientInfo, error) {
 	logger := log.FromContext(ctx).WithField("connectServer", "client")
 	clientURL := clienturlctx.ClientURL(ctx)
 
 	// First check if we have already requested some clientURL with this conn.GetID().
 	if connInfo, ok := s.connInfos.Load(conn.GetId()); ok {
 		if *connInfo.clientURL == *clientURL {
-			return connInfo.client
+			return connInfo.client, nil
 		}
 
 		// For some reason we have changed the clientURL, so we need to close and delete the existing client.
 		if _, clientErr := connInfo.client.client.Close(ctx, conn); clientErr != nil {
-			logger.Warnf("failed to close client: %+v", clientErr)
+			logger.Warnf("failed to close client: %s", clientErr.Error())
+		}
+
+		if err := ctx.Err(); err != nil {
+			return nil, err
 		}
 
 		s.closeClient(connInfo.client, connInfo.clientURL.String())
@@ -174,7 +187,7 @@ func (s *connectServer) client(ctx context.Context, conn *networkservice.Connect
 		}
 		c.count++
 	})
-	return c
+	return c, nil
 }
 
 func (s *connectServer) newClient(clientURL *url.URL) *clientInfo {

pkg/networkservice/common/connect/server_test.go

@@ -520,6 +520,82 @@ func TestConnectServer_DialTimeout(t *testing.T) {
 	timer.Stop()
 }
 
+func TestConnectServer_ChangeURLWithExpiredContext(t *testing.T) {
+	t.Cleanup(func() { goleak.VerifyNone(t) })
+
+	// 1. Create connectServer
+
+	s := next.NewNetworkServiceServer(
+		connect.NewServer(context.Background(),
+			func(ctx context.Context, cc grpc.ClientConnInterface) networkservice.NetworkServiceClient {
+				return next.NewNetworkServiceClient(
+					newMonitorClient(ctx, cc),
+					networkservice.NewNetworkServiceClient(cc),
+				)
+			},
+			connect.WithDialTimeout(time.Second),
+			connect.WithDialOptions(grpc.WithInsecure()),
+		),
+	)
+
+	// 3. Setup servers
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	urlA := &url.URL{Scheme: "tcp", Host: "127.0.0.1:"}
+	serverA := new(captureServer)
+
+	err := startServer(ctx, urlA, serverA)
+	require.NoError(t, err)
+
+	urlB := &url.URL{Scheme: "tcp", Host: "127.0.0.1:"}
+	serverB := new(captureServer)
+
+	err = startServer(ctx, urlB, serverB)
+	require.NoError(t, err)
+
+	// 4. Create request
+
+	request := &networkservice.NetworkServiceRequest{
+		Connection: &networkservice.Connection{
+			Id: "id",
+		},
+	}
+
+	// 5. Request A
+
+	requestCtx, requestCancel := context.WithCancel(context.Background())
+
+	conn, err := s.Request(clienturlctx.WithClientURL(requestCtx, urlA), request.Clone())
+	require.NoError(t, err)
+
+	require.NotNil(t, serverA.capturedRequest)
+
+	requestCancel()
+
+	// 6. Request B with timed out context
+
+	request.Connection = conn.Clone()
+
+	_, err = s.Request(clienturlctx.WithClientURL(requestCtx, urlB), request.Clone())
+	require.Error(t, err)
+
+	require.Nil(t, serverB.capturedRequest)
+
+	// 7. Request B with valid context
+
+	requestCtx, requestCancel = context.WithCancel(context.Background())
+
+	_, err = s.Request(clienturlctx.WithClientURL(requestCtx, urlB), request.Clone())
+	require.NoError(t, err)
+
+	require.Nil(t, serverA.capturedRequest)
+	require.NotNil(t, serverB.capturedRequest)
+
+	requestCancel()
+}
+
 type editServer struct {
 	key       string
 	value     string

pkg/networkservice/common/heal/client.go

@@ -28,6 +28,7 @@ import (
 
 	"github.com/networkservicemesh/sdk/pkg/networkservice/core/next"
 	"github.com/networkservicemesh/sdk/pkg/tools/cancelctx"
+	"github.com/networkservicemesh/sdk/pkg/tools/postpone"
 )
 
 type connectionInfo struct {
@@ -103,6 +104,8 @@ func (u *healClient) Request(ctx context.Context, request *networkservice.Networ
 		}
 	}()
 
+	postponeCtxFunc := postpone.ContextWithValues(ctx)
+
 	conn, err = next.Client(ctx).Request(ctx, request, opts...)
 	if err != nil {
 		return nil, err
@@ -122,11 +125,20 @@ func (u *healClient) Request(ctx context.Context, request *networkservice.Networ
 	//   than monitor stream update.
 	select {
 	case <-ctx.Done():
-		return nil, ctx.Err()
+		err = ctx.Err()
 	case <-condCh:
 	case <-u.ctx.Done():
-		_, _ = next.Client(ctx).Close(ctx, conn)
-		return nil, errors.Errorf("timeout waiting for connection monitor: %s", conn.GetId())
+		err = errors.Errorf("timeout waiting for connection monitor: %s", conn.GetId())
+	}
+	if err != nil {
+		closeCtx, cancelClose := postponeCtxFunc()
+		defer cancelClose()
+
+		if _, closeErr := next.Client(ctx).Close(closeCtx, conn); closeErr != nil {
+			err = errors.Wrapf(err, "connection closed with error: %s", closeErr.Error())
+		}
+
+		return nil, err
 	}
 
 	connInfo.cond.L.Lock()

pkg/networkservice/common/heal/heal_test.go

@@ -22,11 +22,12 @@ import (
 	"testing"
 	"time"
 
-	"github.com/networkservicemesh/api/pkg/api/networkservice"
 	"github.com/stretchr/testify/require"
 	"go.uber.org/goleak"
 	"google.golang.org/grpc"
 
+	"github.com/networkservicemesh/api/pkg/api/networkservice"
+
 	"github.com/networkservicemesh/sdk/pkg/networkservice/chains/endpoint"
 	"github.com/networkservicemesh/sdk/pkg/networkservice/common/clienturl"
 	"github.com/networkservicemesh/sdk/pkg/networkservice/common/connect"
@@ -40,10 +41,17 @@ import (
 	"github.com/networkservicemesh/sdk/pkg/tools/sandbox"
 )
 
-func startRemoteServer(ctx context.Context, t *testing.T, expireDuration time.Duration) (*url.URL, context.CancelFunc) {
+func startRemoteServer(
+	ctx context.Context,
+	t *testing.T,
+	expireDuration time.Duration,
+	additionalFunctionality ...networkservice.NetworkServiceServer,
+) (*url.URL, context.CancelFunc) {
 	ctx, cancel := context.WithCancel(ctx)
 
-	server := endpoint.NewServer(ctx, sandbox.GenerateExpiringToken(expireDuration), endpoint.WithName("remote"))
+	server := endpoint.NewServer(ctx, sandbox.GenerateExpiringToken(expireDuration),
+		endpoint.WithName("remote"),
+		endpoint.WithAdditionalFunctionality(additionalFunctionality...))
 
 	grpcServer := grpc.NewServer()
 	server.Register(grpcServer)
@@ -103,3 +111,56 @@ func TestHeal_CloseChain(t *testing.T) {
 		return counter.Closes() == 1
 	}, time.Second, 10*time.Millisecond)
 }
+
+func TestHeal_CloseChainOnNoMonitorUpdate(t *testing.T) {
+	t.Cleanup(func() { goleak.VerifyNone(t) })
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	counter := new(count.Server)
+	remoteURL, remoteCancel := startRemoteServer(ctx, t, 10*time.Minute, counter)
+	defer remoteCancel()
+
+	// Create fake monitor connection
+	monitorURL, monitorCancel := startRemoteServer(ctx, t, 0)
+	monitorCC, err := grpc.DialContext(ctx, grpcutils.URLToTarget(monitorURL), grpc.WithBlock(), grpc.WithInsecure())
+	require.NoError(t, err)
+	defer func() {
+		monitorCancel()
+		_ = monitorCC.Close()
+	}()
+
+	serverChain := new(networkservice.NetworkServiceClient)
+	*serverChain = adapters.NewServerToClient(
+		next.NewNetworkServiceServer(
+			updatepath.NewServer("server"),
+			updatetoken.NewServer(sandbox.GenerateTestToken),
+			heal.NewServer(ctx,
+				heal.WithOnHeal(serverChain)),
+			clienturl.NewServer(remoteURL),
+			connect.NewServer(ctx, func(ctx context.Context, cc grpc.ClientConnInterface) networkservice.NetworkServiceClient {
+				return next.NewNetworkServiceClient(
+					heal.NewClient(ctx, networkservice.NewMonitorConnectionClient(monitorCC)),
+					networkservice.NewNetworkServiceClient(cc),
+				)
+			}, connect.WithDialOptions(grpc.WithInsecure())),
+		),
+	)
+
+	client := next.NewNetworkServiceClient(
+		updatepath.NewClient("client"),
+		*serverChain,
+	)
+
+	requestCtx, cancelRequest := context.WithTimeout(ctx, time.Second)
+	defer cancelRequest()
+
+	_, err = client.Request(requestCtx, &networkservice.NetworkServiceRequest{
+		Connection: new(networkservice.Connection),
+	})
+	require.Error(t, err)
+
+	require.Equal(t, 1, counter.Requests())
+	require.Equal(t, 1, counter.Closes())
+}

pkg/registry/common/refresh/nse_registry_client.go

@@ -30,6 +30,7 @@ import (
 	"github.com/networkservicemesh/sdk/pkg/registry/core/next"
 	"github.com/networkservicemesh/sdk/pkg/tools/clock"
 	"github.com/networkservicemesh/sdk/pkg/tools/log"
+	"github.com/networkservicemesh/sdk/pkg/tools/postpone"
 	"github.com/networkservicemesh/sdk/pkg/tools/serializectx"
 )
 
@@ -60,6 +61,8 @@ func (c *refreshNSEClient) Register(ctx context.Context, nse *registry.NetworkSe
 		cancel()
 	}
 
+	postponeCtxFunc := postpone.ContextWithValues(ctx)
+
 	reg, err := next.NetworkServiceEndpointRegistryClient(ctx).Register(ctx, nse, opts...)
 	if err != nil {
 		return nil, err
@@ -71,7 +74,10 @@ func (c *refreshNSEClient) Register(ctx context.Context, nse *registry.NetworkSe
 
 		cancel, err = c.startRefresh(ctx, refreshNSE, expirationDuration)
 		if err != nil {
-			if _, unregisterErr := next.NetworkServiceEndpointRegistryServer(ctx).Unregister(ctx, reg); unregisterErr != nil {
+			unregisterCtx, cancelUnregister := postponeCtxFunc()
+			defer cancelUnregister()
+
+			if _, unregisterErr := next.NetworkServiceEndpointRegistryServer(ctx).Unregister(unregisterCtx, reg); unregisterErr != nil {
 				logger.Errorf("failed to unregister endpoint on error: %s %s", reg.Name, unregisterErr.Error())
 			}
 			return nil, err