Add more mutexes in dial chain element to fix race conditions (#1670)

* some minor change

Signed-off-by: NikitaSkrynnik <nikita.skrynnik@xored.com>

* add more locks

Signed-off-by: NikitaSkrynnik <nikita.skrynnik@xored.com>

---------

Signed-off-by: NikitaSkrynnik <nikita.skrynnik@xored.com>

pkg/networkservice/common/dial/dialer.go

@@ -31,10 +31,9 @@ import (
 )
 
 type dialer struct {
-	ctx            context.Context
-	cleanupContext context.Context
-	clientURL      *url.URL
-	cleanupCancel  context.CancelFunc
+	ctx           context.Context
+	clientURL     *url.URL
+	cleanupCancel context.CancelFunc
 	*grpc.ClientConn
 	dialOptions []grpc.DialOption
 	dialTimeout time.Duration
@@ -70,41 +69,53 @@ func (di *dialer) Dial(ctx context.Context, clientURL *url.URL) error {
 	}
 
 	// Dial
+	di.mu.Lock()
 	target := grpcutils.URLToTarget(di.clientURL)
+	di.mu.Unlock()
+
 	cc, err := grpc.DialContext(dialCtx, target, di.dialOptions...)
 	if err != nil {
 		if cc != nil {
 			_ = cc.Close()
 		}
 		return errors.Wrapf(err, "failed to dial %s", target)
 	}
+	di.mu.Lock()
 	di.ClientConn = cc
-
-	di.cleanupContext, di.cleanupCancel = context.WithCancel(di.ctx)
+	var cleanupContext context.Context
+	cleanupContext, di.cleanupCancel = context.WithCancel(di.ctx)
+	di.mu.Unlock()
 
 	go func(cleanupContext context.Context, cc *grpc.ClientConn) {
 		<-cleanupContext.Done()
 		_ = cc.Close()
-	}(di.cleanupContext, cc)
+	}(cleanupContext, cc)
 	return nil
 }
 
 func (di *dialer) Close() error {
 	if di != nil && di.cleanupCancel != nil {
+		di.mu.Lock()
 		di.cleanupCancel()
+		di.mu.Unlock()
 		runtime.Gosched()
 	}
 	return nil
 }
 
 func (di *dialer) Invoke(ctx context.Context, method string, args, reply interface{}, opts ...grpc.CallOption) error {
+	di.mu.Lock()
+	defer di.mu.Unlock()
 	if di.ClientConn == nil {
 		return errors.New("no dialer.ClientConn found")
 	}
 	return di.ClientConn.Invoke(ctx, method, args, reply, opts...)
 }
 
 func (di *dialer) NewStream(ctx context.Context, desc *grpc.StreamDesc, method string, opts ...grpc.CallOption) (grpc.ClientStream, error) {
+	di.mu.Lock()
+	defer di.mu.Unlock()
+
 	if di.ClientConn == nil {
 		return nil, errors.New("no dialer.ClientConn found")
 	}

pkg/networkservice/common/discoverforwarder/server.go

@@ -1,6 +1,6 @@
 // Copyright (c) 2021-2022 Doc.ai and/or its affiliates.
 //
-// Copyright (c) 2023 Cisco and/or its affiliates.
+// Copyright (c) 2023-2024 Cisco and/or its affiliates.
 //
 // SPDX-License-Identifier: Apache-2.0
 //