use less agressive patch to fix nsmgr

Signed-off-by: Denis Tingaikin <denis.tingajkin@xored.com>
networkservicemesh · Sep 3, 2024 · d9269e1 · d9269e1
1 parent 2e058bf
commit d9269e1
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 4 deletions.
diff --git a/pkg/registry/common/dial/ns_client.go b/pkg/registry/common/dial/ns_client.go
@@ -1,4 +1,4 @@
-// Copyright (c) 2021-2022 Cisco and/or its affiliates.
+// Copyright (c) 2021-2024 Cisco and/or its affiliates.
 //
 // SPDX-License-Identifier: Apache-2.0
 //
@@ -149,8 +149,18 @@ func (c *dialNSClient) Find(ctx context.Context, in *registry.NetworkServiceQuer
 		return nil, err
 	}
 
+	var stopContext, stopCancel = context.WithTimeout(resp.Context(), time.Minute/4)
+	if in.Watch {
+		stopContext = resp.Context()
+	}
+	cleanupFn = func() {
+		stopCancel()
+		clientconn.Delete(ctx)
+		_ = di.Close()
+	}
+
 	go func() {
-		<-resp.Context().Done()
+		<-stopContext.Done()
 		cleanupFn()
 	}()
 

diff --git a/pkg/registry/common/dial/nse_client.go b/pkg/registry/common/dial/nse_client.go
@@ -1,4 +1,4 @@
-// Copyright (c) 2021-2023 Cisco and/or its affiliates.
+// Copyright (c) 2021-2024 Cisco and/or its affiliates.
 //
 // SPDX-License-Identifier: Apache-2.0
 //
@@ -155,8 +155,16 @@ func (c *dialNSEClient) Find(ctx context.Context, in *registry.NetworkServiceEnd
 		return nil, err
 	}
 
+	var stopContext, stopCancel = context.WithTimeout(resp.Context(), time.Minute/4)
+
+	cleanupFn = func() {
+		stopCancel()
+		clientconn.Delete(ctx)
+		_ = di.Close()
+	}
+
 	go func() {
-		<-resp.Context().Done()
+		<-stopContext.Done()
 		cleanupFn()
 	}()