github.com/ethersphere/bee/v2@v2.2.0/pkg/topology/lightnode/container.go (about)

     1  // Copyright 2021 The Swarm Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  package lightnode
     6  
     7  import (
     8  	"context"
     9  	"crypto/rand"
    10  	"math/big"
    11  	"sync"
    12  
    13  	"github.com/ethersphere/bee/v2/pkg/p2p"
    14  	"github.com/ethersphere/bee/v2/pkg/swarm"
    15  	"github.com/ethersphere/bee/v2/pkg/topology"
    16  	"github.com/ethersphere/bee/v2/pkg/topology/pslice"
    17  )
    18  
    19  type Container struct {
    20  	base              swarm.Address
    21  	peerMu            sync.Mutex // peerMu guards connectedPeers and disconnectedPeers.
    22  	connectedPeers    *pslice.PSlice
    23  	disconnectedPeers *pslice.PSlice
    24  	metrics           metrics
    25  }
    26  
    27  func NewContainer(base swarm.Address) *Container {
    28  	return &Container{
    29  		base:              base,
    30  		connectedPeers:    pslice.New(1, base),
    31  		disconnectedPeers: pslice.New(1, base),
    32  		metrics:           newMetrics(),
    33  	}
    34  }
    35  
    36  func (c *Container) Connected(ctx context.Context, peer p2p.Peer) {
    37  	c.peerMu.Lock()
    38  	defer c.peerMu.Unlock()
    39  
    40  	addr := peer.Address
    41  	c.connectedPeers.Add(addr)
    42  	c.disconnectedPeers.Remove(addr)
    43  
    44  	c.metrics.CurrentlyConnectedPeers.Set(float64(c.connectedPeers.Length()))
    45  	c.metrics.CurrentlyDisconnectedPeers.Set(float64(c.disconnectedPeers.Length()))
    46  }
    47  
    48  func (c *Container) Disconnected(peer p2p.Peer) {
    49  	c.peerMu.Lock()
    50  	defer c.peerMu.Unlock()
    51  
    52  	addr := peer.Address
    53  	if found := c.connectedPeers.Exists(addr); found {
    54  		c.connectedPeers.Remove(addr)
    55  		c.disconnectedPeers.Add(addr)
    56  	}
    57  
    58  	c.metrics.CurrentlyConnectedPeers.Set(float64(c.connectedPeers.Length()))
    59  	c.metrics.CurrentlyDisconnectedPeers.Set(float64(c.disconnectedPeers.Length()))
    60  }
    61  
    62  func (c *Container) Count() int {
    63  	return c.connectedPeers.Length()
    64  }
    65  
    66  func (c *Container) RandomPeer(not swarm.Address) (swarm.Address, error) {
    67  	c.peerMu.Lock()
    68  	defer c.peerMu.Unlock()
    69  	var (
    70  		cnt   = big.NewInt(int64(c.Count()))
    71  		addr  = swarm.ZeroAddress
    72  		count = int64(0)
    73  	)
    74  
    75  PICKPEER:
    76  	i, e := rand.Int(rand.Reader, cnt)
    77  	if e != nil {
    78  		return swarm.ZeroAddress, e
    79  	}
    80  	i64 := i.Int64()
    81  
    82  	count = 0
    83  	_ = c.connectedPeers.EachBinRev(func(peer swarm.Address, _ uint8) (bool, bool, error) {
    84  		if count == i64 {
    85  			addr = peer
    86  			return true, false, nil
    87  		}
    88  		count++
    89  		return false, false, nil
    90  	})
    91  
    92  	if addr.Equal(not) {
    93  		goto PICKPEER
    94  	}
    95  
    96  	return addr, nil
    97  }
    98  
    99  func (c *Container) EachPeer(pf topology.EachPeerFunc) error {
   100  	return c.connectedPeers.EachBin(pf)
   101  }
   102  
   103  func (c *Container) PeerInfo() topology.BinInfo {
   104  	return topology.BinInfo{
   105  		BinPopulation:     uint(c.connectedPeers.Length()),
   106  		BinConnected:      uint(c.connectedPeers.Length()),
   107  		DisconnectedPeers: peersInfo(c.disconnectedPeers),
   108  		ConnectedPeers:    peersInfo(c.connectedPeers),
   109  	}
   110  }
   111  
   112  func peersInfo(s *pslice.PSlice) []*topology.PeerInfo {
   113  	if s.Length() == 0 {
   114  		return nil
   115  	}
   116  	peers := make([]*topology.PeerInfo, 0, s.Length())
   117  	_ = s.EachBin(func(addr swarm.Address, po uint8) (bool, bool, error) {
   118  		peers = append(peers, &topology.PeerInfo{Address: addr})
   119  		return false, false, nil
   120  	})
   121  	return peers
   122  }