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