From e314f328c4806bf1b34b6e3c31abdc4afdfaaac4 Mon Sep 17 00:00:00 2001
From: Dmitrii Stepanov <d.stepanov@yadro.com>
Date: Wed, 9 Oct 2024 11:18:24 +0300
Subject: [PATCH] [#1422] tree: Use dialer source for tree service connections

Signed-off-by: Dmitrii Stepanov <d.stepanov@yadro.com>
---
 cmd/frostfs-node/tree.go     | 1 +
 pkg/services/tree/cache.go   | 6 +++++-
 pkg/services/tree/options.go | 8 ++++++++
 pkg/services/tree/service.go | 2 +-
 4 files changed, 15 insertions(+), 2 deletions(-)

diff --git a/cmd/frostfs-node/tree.go b/cmd/frostfs-node/tree.go
index 192f08471..f188e2fbc 100644
--- a/cmd/frostfs-node/tree.go
+++ b/cmd/frostfs-node/tree.go
@@ -67,6 +67,7 @@ func initTreeService(c *cfg) {
 		tree.WithAPELocalOverrideStorage(c.cfgObject.cfgAccessPolicyEngine.accessPolicyEngine.LocalStorage()),
 		tree.WithAPEMorphRuleStorage(c.cfgObject.cfgAccessPolicyEngine.accessPolicyEngine.MorphRuleChainStorage()),
 		tree.WithNetmapState(c.cfgNetmap.state),
+		tree.WithDialerSource(c.dialerSource),
 	)
 
 	c.cfgGRPC.performAndSave(func(_ string, _ net.Listener, s *grpc.Server) {
diff --git a/pkg/services/tree/cache.go b/pkg/services/tree/cache.go
index 38501b852..e490cb855 100644
--- a/pkg/services/tree/cache.go
+++ b/pkg/services/tree/cache.go
@@ -8,6 +8,7 @@ import (
 	"sync"
 	"time"
 
+	internalNet "git.frostfs.info/TrueCloudLab/frostfs-node/internal/net"
 	"git.frostfs.info/TrueCloudLab/frostfs-node/pkg/network"
 	metrics "git.frostfs.info/TrueCloudLab/frostfs-observability/metrics/grpc"
 	tracing "git.frostfs.info/TrueCloudLab/frostfs-observability/tracing/grpc"
@@ -21,6 +22,7 @@ type clientCache struct {
 	sync.Mutex
 	simplelru.LRU[string, cacheItem]
 	key *ecdsa.PrivateKey
+	ds  *internalNet.DialerSource
 }
 
 type cacheItem struct {
@@ -36,7 +38,7 @@ const (
 
 var errRecentlyFailed = errors.New("client has recently failed")
 
-func (c *clientCache) init(pk *ecdsa.PrivateKey) {
+func (c *clientCache) init(pk *ecdsa.PrivateKey, ds *internalNet.DialerSource) {
 	l, _ := simplelru.NewLRU(defaultClientCacheSize, func(_ string, value cacheItem) {
 		if conn := value.cc; conn != nil {
 			_ = conn.Close()
@@ -44,6 +46,7 @@ func (c *clientCache) init(pk *ecdsa.PrivateKey) {
 	})
 	c.LRU = *l
 	c.key = pk
+	c.ds = ds
 }
 
 func (c *clientCache) get(ctx context.Context, netmapAddr string) (TreeServiceClient, error) {
@@ -99,6 +102,7 @@ func (c *clientCache) dialTreeService(ctx context.Context, netmapAddr string) (*
 			metrics.NewStreamClientInterceptor(),
 			tracing.NewStreamClientInterceptor(),
 		),
+		grpc.WithContextDialer(c.ds.GrpcContextDialer()),
 	}
 
 	if !netAddr.IsTLSEnabled() {
diff --git a/pkg/services/tree/options.go b/pkg/services/tree/options.go
index 1db5607f6..1633ae557 100644
--- a/pkg/services/tree/options.go
+++ b/pkg/services/tree/options.go
@@ -4,6 +4,7 @@ import (
 	"crypto/ecdsa"
 	"time"
 
+	"git.frostfs.info/TrueCloudLab/frostfs-node/internal/net"
 	"git.frostfs.info/TrueCloudLab/frostfs-node/pkg/core/container"
 	frostfsidcore "git.frostfs.info/TrueCloudLab/frostfs-node/pkg/core/frostfsid"
 	"git.frostfs.info/TrueCloudLab/frostfs-node/pkg/core/netmap"
@@ -45,6 +46,7 @@ type cfg struct {
 	morphChainStorage    policyengine.MorphRuleChainStorageReader
 
 	metrics MetricsRegister
+	ds      *net.DialerSource
 }
 
 // Option represents configuration option for a tree service.
@@ -161,3 +163,9 @@ func WithNetmapState(state netmap.State) Option {
 		c.state = state
 	}
 }
+
+func WithDialerSource(ds *net.DialerSource) Option {
+	return func(c *cfg) {
+		c.ds = ds
+	}
+}
diff --git a/pkg/services/tree/service.go b/pkg/services/tree/service.go
index 60bb1a6ad..2cb2af294 100644
--- a/pkg/services/tree/service.go
+++ b/pkg/services/tree/service.go
@@ -65,7 +65,7 @@ func New(opts ...Option) *Service {
 		s.log = &logger.Logger{Logger: zap.NewNop()}
 	}
 
-	s.cache.init(s.key)
+	s.cache.init(s.key, s.ds)
 	s.closeCh = make(chan struct{})
 	s.replicateCh = make(chan movePair, s.replicatorChannelCapacity)
 	s.replicateLocalCh = make(chan applyOp)