From 23fc4f1965c3fd82c878545d38525b3e4c0d385c Mon Sep 17 00:00:00 2001 From: denis-tingaikin Date: Mon, 11 Dec 2023 18:11:02 +0300 Subject: [PATCH] cancel sequence on success Signed-off-by: denis-tingaikin --- pkg/networkservice/common/heal/client.go | 14 +++++++++++++- 1 file changed, 13 insertions(+), 1 deletion(-) diff --git a/pkg/networkservice/common/heal/client.go b/pkg/networkservice/common/heal/client.go index e7558c390..0181858aa 100644 --- a/pkg/networkservice/common/heal/client.go +++ b/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 @@ -70,7 +76,10 @@ 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 +99,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 }