github.com/celestiaorg/celestia-node@v0.15.0-beta.1/share/p2p/discovery/set.go (about)

     1  package discovery
     2  
     3  import (
     4  	"context"
     5  	"sync"
     6  
     7  	"github.com/libp2p/go-libp2p/core/peer"
     8  )
     9  
    10  // limitedSet is a thread safe set of peers with given limit.
    11  // Inspired by libp2p peer.Set but extended with Remove method.
    12  type limitedSet struct {
    13  	lk sync.RWMutex
    14  	ps map[peer.ID]struct{}
    15  
    16  	limit    uint
    17  	waitPeer chan peer.ID
    18  }
    19  
    20  // newLimitedSet constructs a set with the maximum peers amount.
    21  func newLimitedSet(limit uint) *limitedSet {
    22  	ps := new(limitedSet)
    23  	ps.ps = make(map[peer.ID]struct{})
    24  	ps.limit = limit
    25  	ps.waitPeer = make(chan peer.ID)
    26  	return ps
    27  }
    28  
    29  func (ps *limitedSet) Contains(p peer.ID) bool {
    30  	ps.lk.RLock()
    31  	_, ok := ps.ps[p]
    32  	ps.lk.RUnlock()
    33  	return ok
    34  }
    35  
    36  func (ps *limitedSet) Limit() uint {
    37  	return ps.limit
    38  }
    39  
    40  func (ps *limitedSet) Size() uint {
    41  	ps.lk.RLock()
    42  	defer ps.lk.RUnlock()
    43  	return uint(len(ps.ps))
    44  }
    45  
    46  // Add attempts to add the given peer into the set.
    47  func (ps *limitedSet) Add(p peer.ID) (added bool) {
    48  	ps.lk.Lock()
    49  	if _, ok := ps.ps[p]; ok {
    50  		ps.lk.Unlock()
    51  		return false
    52  	}
    53  	ps.ps[p] = struct{}{}
    54  	ps.lk.Unlock()
    55  
    56  	for {
    57  		// peer will be pushed to the channel only when somebody is reading from it.
    58  		// this is done to handle case when Peers() was called on empty set.
    59  		select {
    60  		case ps.waitPeer <- p:
    61  		default:
    62  			return true
    63  		}
    64  	}
    65  }
    66  
    67  func (ps *limitedSet) Remove(id peer.ID) {
    68  	ps.lk.Lock()
    69  	delete(ps.ps, id)
    70  	ps.lk.Unlock()
    71  }
    72  
    73  // Peers returns all discovered peers from the set.
    74  func (ps *limitedSet) Peers(ctx context.Context) ([]peer.ID, error) {
    75  	ps.lk.RLock()
    76  	if len(ps.ps) > 0 {
    77  		out := make([]peer.ID, 0, len(ps.ps))
    78  		for p := range ps.ps {
    79  			out = append(out, p)
    80  		}
    81  		ps.lk.RUnlock()
    82  		return out, nil
    83  	}
    84  	ps.lk.RUnlock()
    85  
    86  	// block until a new peer will be discovered
    87  	select {
    88  	case <-ctx.Done():
    89  		return nil, ctx.Err()
    90  	case p := <-ps.waitPeer:
    91  		return []peer.ID{p}, nil
    92  	}
    93  }