diff --git a/pkg/network/discovery.go b/pkg/network/discovery.go index 03340eee9..bb9606240 100644 --- a/pkg/network/discovery.go +++ b/pkg/network/discovery.go @@ -18,6 +18,7 @@ type Discoverer interface { RequestRemote(int) RegisterBadAddr(string) RegisterGoodAddr(string) + RegisterConnectedAddr(string) UnregisterConnectedAddr(string) UnconnectedPeers() []string BadPeers() []string @@ -153,8 +154,8 @@ func (d *DefaultDiscovery) UnregisterConnectedAddr(s string) { d.lock.Unlock() } -// registerConnectedAddr tells discoverer that given address is now connected. -func (d *DefaultDiscovery) registerConnectedAddr(addr string) { +// RegisterConnectedAddr tells discoverer that given address is now connected. +func (d *DefaultDiscovery) RegisterConnectedAddr(addr string) { d.lock.Lock() delete(d.unconnectedAddrs, addr) d.connectedAddrs[addr] = true @@ -166,7 +167,7 @@ func (d *DefaultDiscovery) tryAddress(addr string) { d.RegisterBadAddr(addr) d.RequestRemote(1) } else { - d.registerConnectedAddr(addr) + d.RegisterConnectedAddr(addr) } } diff --git a/pkg/network/helper_test.go b/pkg/network/helper_test.go index da58f3b5c..823d78525 100644 --- a/pkg/network/helper_test.go +++ b/pkg/network/helper_test.go @@ -136,6 +136,7 @@ func (d testDiscovery) BackFill(addrs ...string) {} func (d testDiscovery) PoolCount() int { return 0 } func (d testDiscovery) RegisterBadAddr(string) {} func (d testDiscovery) RegisterGoodAddr(string) {} +func (d testDiscovery) RegisterConnectedAddr(string) {} func (d testDiscovery) UnregisterConnectedAddr(string) {} func (d testDiscovery) UnconnectedPeers() []string { return []string{} } func (d testDiscovery) RequestRemote(n int) {} diff --git a/pkg/network/server.go b/pkg/network/server.go index db5867254..1f1af4958 100644 --- a/pkg/network/server.go +++ b/pkg/network/server.go @@ -348,6 +348,7 @@ func (s *Server) handleVersionCmd(p Peer, version *payload.Version) error { return errIdenticalID } peerAddr := p.PeerAddr().String() + s.discovery.RegisterConnectedAddr(peerAddr) s.lock.RLock() for peer := range s.peers { if p == peer {