@@ -95,13 +95,21 @@ func DisableReuseport() Option {
9595 return nil
9696 }
9797}
98+
9899func WithConnectionTimeout (d time.Duration ) Option {
99100 return func (tr * TcpTransport ) error {
100101 tr .connectTimeout = d
101102 return nil
102103 }
103104}
104105
106+ func WithResourceManager (rcmgr network.ResourceManager ) Option {
107+ return func (tr * TcpTransport ) error {
108+ tr .rcmgr = rcmgr
109+ return nil
110+ }
111+ }
112+
105113// TcpTransport is the TCP transport.
106114type TcpTransport struct {
107115 // Connection upgrader for upgrading insecure stream connections to
@@ -114,6 +122,8 @@ type TcpTransport struct {
114122 // TCP connect timeout
115123 connectTimeout time.Duration
116124
125+ rcmgr network.ResourceManager
126+
117127 reuse rtpt.Transport
118128}
119129
@@ -159,6 +169,17 @@ func (t *TcpTransport) maDial(ctx context.Context, raddr ma.Multiaddr) (manet.Co
159169
160170// Dial dials the peer at the remote address.
161171func (t * TcpTransport ) Dial (ctx context.Context , raddr ma.Multiaddr , p peer.ID ) (transport.CapableConn , error ) {
172+ var connScope network.ConnectionScope
173+ if t .rcmgr != nil {
174+ var err error
175+ connScope , err = t .rcmgr .OpenConnection (network .DirOutbound , true )
176+ if err != nil {
177+ return nil , err
178+ }
179+ if err := connScope .SetPeer (p ); err != nil {
180+ return nil , err
181+ }
182+ }
162183 conn , err := t .maDial (ctx , raddr )
163184 if err != nil {
164185 return nil , err
@@ -176,7 +197,7 @@ func (t *TcpTransport) Dial(ctx context.Context, raddr ma.Multiaddr, p peer.ID)
176197 if ok , isClient , _ := network .GetSimultaneousConnect (ctx ); ok && ! isClient {
177198 direction = network .DirInbound
178199 }
179- return t .Upgrader .Upgrade (ctx , t , c , direction , p )
200+ return t .Upgrader .Upgrade (ctx , t , c , direction , p , connScope )
180201}
181202
182203// UseReuseport returns true if reuseport is enabled and available.
@@ -198,7 +219,7 @@ func (t *TcpTransport) Listen(laddr ma.Multiaddr) (transport.Listener, error) {
198219 return nil , err
199220 }
200221 list = newTracingListener (& tcpListener {list , 0 })
201- return t .Upgrader .UpgradeListener (t , list ), nil
222+ return t .Upgrader .UpgradeListener (t , list , t . rcmgr ), nil
202223}
203224
204225// Protocols returns the list of terminal protocols this transport can dial.
0 commit comments