p2p: properly decrement metered peer count

This commit is contained in:
Kurkó Mihály 2018-10-14 15:39:23 +03:00
parent 070a519aa5
commit 8556748638

View file

@ -159,17 +159,19 @@ func (c *meteredConn) Write(b []byte) (n int, err error) {
// handshakeDone is called when a peer handshake is done. Registers the peer to // handshakeDone is called when a peer handshake is done. Registers the peer to
// the ingress and the egress traffic registries using the peer's IP and node ID, // the ingress and the egress traffic registries using the peer's IP and node ID,
// also emits connect event. // also emits connect event.
func (c *meteredConn) handshakeDone(id enode.ID) { func (c *meteredConn) handshakeDone(nodeID enode.ID) {
id := nodeID.String()
if atomic.AddInt32(&meteredPeerCount, 1) >= MeteredPeerLimit { if atomic.AddInt32(&meteredPeerCount, 1) >= MeteredPeerLimit {
// Don't register the peer in the traffic registries.
atomic.AddInt32(&meteredPeerCount, -1) atomic.AddInt32(&meteredPeerCount, -1)
c.lock.Lock() c.lock.Lock()
c.id, c.trafficMetered = id.String(), false c.id, c.trafficMetered = id, false
c.lock.Unlock() c.lock.Unlock()
log.Warn("Metered peer count reached the limit") log.Warn("Metered peer count reached the limit")
} else { } else {
key := fmt.Sprintf("%s/%s", c.ip, id) key := fmt.Sprintf("%s/%s", c.ip, id)
c.lock.Lock() c.lock.Lock()
c.id, c.trafficMetered = id.String(), true c.id, c.trafficMetered = id, true
c.ingressMeter = metrics.NewRegisteredMeter(key, PeerIngressRegistry) c.ingressMeter = metrics.NewRegisteredMeter(key, PeerIngressRegistry)
c.egressMeter = metrics.NewRegisteredMeter(key, PeerEgressRegistry) c.egressMeter = metrics.NewRegisteredMeter(key, PeerEgressRegistry)
c.lock.Unlock() c.lock.Unlock()
@ -177,7 +179,7 @@ func (c *meteredConn) handshakeDone(id enode.ID) {
meteredPeerFeed.Send(MeteredPeerEvent{ meteredPeerFeed.Send(MeteredPeerEvent{
Type: PeerConnected, Type: PeerConnected,
IP: c.ip, IP: c.ip,
ID: id.String(), ID: id,
Elapsed: time.Since(c.connected), Elapsed: time.Since(c.connected),
}) })
} }
@ -188,7 +190,7 @@ func (c *meteredConn) Close() error {
err := c.Conn.Close() err := c.Conn.Close()
c.lock.RLock() c.lock.RLock()
if c.id == "" { if c.id == "" {
// If the peer disconnects before the handshake // If the peer disconnects before the handshake.
c.lock.RUnlock() c.lock.RUnlock()
meteredPeerFeed.Send(MeteredPeerEvent{ meteredPeerFeed.Send(MeteredPeerEvent{
Type: PeerHandshakeFailed, Type: PeerHandshakeFailed,
@ -197,10 +199,9 @@ func (c *meteredConn) Close() error {
}) })
return err return err
} }
// Decrement the metered peer count
atomic.AddInt32(&meteredPeerCount, -1)
id := c.id id := c.id
if !c.trafficMetered { if !c.trafficMetered {
// If the peer isn't registered in the traffic registries.
c.lock.RUnlock() c.lock.RUnlock()
meteredPeerFeed.Send(MeteredPeerEvent{ meteredPeerFeed.Send(MeteredPeerEvent{
Type: PeerDisconnected, Type: PeerDisconnected,
@ -212,6 +213,9 @@ func (c *meteredConn) Close() error {
ingress, egress := uint64(c.ingressMeter.Count()), uint64(c.egressMeter.Count()) ingress, egress := uint64(c.ingressMeter.Count()), uint64(c.egressMeter.Count())
c.lock.RUnlock() c.lock.RUnlock()
// Decrement the metered peer count
atomic.AddInt32(&meteredPeerCount, -1)
// Unregister the peer from the traffic registries // Unregister the peer from the traffic registries
key := fmt.Sprintf("%s/%s", c.ip, id) key := fmt.Sprintf("%s/%s", c.ip, id)
PeerIngressRegistry.Unregister(key) PeerIngressRegistry.Unregister(key)