@@ -27,7 +27,7 @@ import (
2727)
2828
2929var (
30- ErrSourceClosed = fmt .Errorf ("source closed: %w" , errEndpointClosed )
30+ ErrSourceClosed = fmt .Errorf ("source closed: %w: %w " , errEndpointClosed , net . ErrClosed )
3131 ErrSourceConnect = errors .New ("cannot connect destination" )
3232 ErrSourceConnectDestinations = fmt .Errorf ("%w: all peer connections were unreachable" , ErrSourceConnect )
3333 ErrSourceNoActiveDestinations = fmt .Errorf ("%w: no active peer connections" , ErrSourceConnect )
@@ -88,6 +88,10 @@ func (s *Source) Dial(network, address string) (net.Conn, error) {
8888// DialContext dials into any available destination. Both network and address are ignored.
8989// Blocks until connection can be established.
9090func (s * Source ) DialContext (ctx context.Context , network , address string ) (net.Conn , error ) {
91+ if err := s .ep .ctx .Err (); err != nil {
92+ return nil , fmt .Errorf ("dial context: %w" , context .Cause (s .ep .ctx ))
93+ }
94+
9195 if s .cfg .DestinationPolicy == NoPolicy {
9296 conns , err := s .findActive ()
9397 if err != nil {
@@ -157,6 +161,20 @@ func (s *Source) runActive(ctx context.Context) {
157161}
158162
159163func (s * Source ) runActiveErr (ctx context.Context ) error {
164+ defer func () {
165+ s .conns .Store (new ([]sourceConn ))
166+ }()
167+
168+ isConnsTracking := s .cfg .DestinationPolicy == LeastConnsPolicy
169+ if isConnsTracking {
170+ defer func () {
171+ s .connsTrackingMu .Lock ()
172+ defer s .connsTrackingMu .Unlock ()
173+
174+ s .connsTracking = map [peerID ]* atomic.Int32 {}
175+ }()
176+ }
177+
160178 return s .ep .peer .activeConnsListen (ctx , func (active map [peerConnKey ]* quic.Conn ) error {
161179 s .logger .Debug ("active conns" , "len" , len (active ))
162180
@@ -166,7 +184,7 @@ func (s *Source) runActiveErr(ctx context.Context) error {
166184 }
167185 s .conns .Store (& conns )
168186
169- if s . connsTracking != nil {
187+ if isConnsTracking {
170188 activePeers := map [peerID ]struct {}{}
171189 for peer := range active {
172190 activePeers [peer .id ] = struct {}{}
0 commit comments