cancel sequence on sucess

Signed-off-by: denis-tingaikin <[email protected]>

Author: denis-tingaikin

pkg/networkservice/common/heal/client.go

@@ -20,6 +20,7 @@ import (
 	"context"
 	"time"
 
+	"github.com/edwarnicke/genericsync"
 	"github.com/networkservicemesh/api/pkg/api/networkservice"
 	"google.golang.org/grpc"
 	"google.golang.org/protobuf/types/known/emptypb"
@@ -32,8 +33,13 @@ import (
 	"github.com/networkservicemesh/sdk/pkg/tools/postpone"
 )
 
+type contextAndCancel struct {
+	context.Context
+	cancel func()
+}
 type healClient struct {
 	chainCtx              context.Context
+	contextMap            genericsync.Map[string, contextAndCancel]
 	livenessCheck         LivenessCheck
 	livenessCheckInterval time.Duration
 	livenessCheckTimeout  time.Duration
@@ -69,8 +75,12 @@ func (h *healClient) Request(ctx context.Context, request *networkservice.Networ
 
 	conn, err := next.Client(ctx).Request(ctx, request, opts...)
 	if err != nil {
+
 		if len(request.GetConnection().GetPath().GetPathSegments()) > 1 && h.livenessCheck == nil && !loopHandle.healingStarted {
-			_ = begin.FromContext(ctx).Request(begin.WithReselect())
+			var healCtx, cancel = context.WithCancel(h.chainCtx)
+			var ctxAndCancel = contextAndCancel{Context: healCtx, cancel: cancel}
+			ctxAndCancel, _ = h.contextMap.LoadOrStore(request.GetConnection().GetId(), ctxAndCancel)
+			_ = begin.FromContext(ctx).Request(begin.WithReselect(), begin.CancelContext(ctxAndCancel.Context))
 		}
 		if loaded && !loopHandle.healingStarted {
 			eventLoopErr := h.startEventLoop(ctx, request.GetConnection())
@@ -90,6 +100,9 @@ func (h *healClient) Request(ctx context.Context, request *networkservice.Networ
 		_, _ = next.Client(closeCtx).Close(closeCtx, conn)
 		return nil, eventLoopErr
 	}
+	if v, ok := h.contextMap.LoadAndDelete(request.GetConnection().GetId()); ok {
+		v.cancel()
+	}
 	return conn, nil
 }