github.com/turingchain2020/turingchain@v1.1.21/system/p2p/dht/manage/peerinfo.go (about)

     1  package manage
     2  
     3  import (
     4  	"context"
     5  	"sync"
     6  	"sync/atomic"
     7  	"time"
     8  
     9  	"github.com/turingchain2020/turingchain/queue"
    10  	"github.com/turingchain2020/turingchain/types"
    11  	"github.com/libp2p/go-libp2p-core/host"
    12  	"github.com/libp2p/go-libp2p-core/peer"
    13  )
    14  
    15  // PeerInfoManager peer info manager
    16  type PeerInfoManager struct {
    17  	ctx       context.Context
    18  	peerInfo  sync.Map
    19  	client    queue.Client
    20  	host      host.Host
    21  	maxHeight int64
    22  }
    23  
    24  type peerStoreInfo struct {
    25  	storeTime time.Time
    26  	peer      *types.Peer
    27  }
    28  
    29  // NewPeerInfoManager new peer info manager
    30  func NewPeerInfoManager(ctx context.Context, host host.Host, cli queue.Client) *PeerInfoManager {
    31  	peerInfoManage := &PeerInfoManager{
    32  		ctx:    ctx,
    33  		client: cli,
    34  		host:   host,
    35  	}
    36  	go peerInfoManage.start()
    37  	return peerInfoManage
    38  }
    39  
    40  // Refresh refreshes peer info
    41  func (p *PeerInfoManager) Refresh(peer *types.Peer) {
    42  	if peer == nil {
    43  		return
    44  	}
    45  	storeInfo := peerStoreInfo{
    46  		storeTime: time.Now(),
    47  		peer:      peer,
    48  	}
    49  	p.peerInfo.Store(peer.Name, &storeInfo)
    50  	if peer.GetHeader().GetHeight() > atomic.LoadInt64(&p.maxHeight) {
    51  		atomic.StoreInt64(&p.maxHeight, peer.GetHeader().GetHeight())
    52  	}
    53  }
    54  
    55  // Fetch returns info of given peer
    56  func (p *PeerInfoManager) Fetch(pid peer.ID) *types.Peer {
    57  	key := pid.Pretty()
    58  	v, ok := p.peerInfo.Load(key)
    59  	if !ok {
    60  		return nil
    61  	}
    62  	if info, ok := v.(*peerStoreInfo); ok {
    63  		if time.Since(info.storeTime) > time.Minute {
    64  			p.peerInfo.Delete(key)
    65  			return nil
    66  		}
    67  		return info.peer
    68  	}
    69  	return nil
    70  }
    71  
    72  // FetchAll returns all peers info
    73  func (p *PeerInfoManager) FetchAll() []*types.Peer {
    74  	var peers []*types.Peer
    75  	var self *types.Peer
    76  	p.peerInfo.Range(func(key, value interface{}) bool {
    77  		info := value.(*peerStoreInfo)
    78  		if time.Since(info.storeTime) > time.Minute {
    79  			p.peerInfo.Delete(key)
    80  			return true
    81  		}
    82  		if key.(string) == p.host.ID().Pretty() {
    83  			self = info.peer
    84  			return true
    85  		}
    86  		peers = append(peers, info.peer)
    87  		return true
    88  	})
    89  	if self != nil {
    90  		peers = append(peers, self)
    91  	}
    92  	return peers
    93  }
    94  
    95  // PeerHeight returns block height of given peer
    96  func (p *PeerInfoManager) PeerHeight(pid peer.ID) int64 {
    97  	v, ok := p.peerInfo.Load(pid.Pretty())
    98  	if !ok {
    99  		return -1
   100  	}
   101  	info, ok := v.(*peerStoreInfo)
   102  	if !ok {
   103  		return -1
   104  	}
   105  	if info.peer.GetHeader() == nil {
   106  		return -1
   107  	}
   108  	return info.peer.GetHeader().Height
   109  }
   110  
   111  // PeerMaxHeight returns max block height of all connected peers.
   112  func (p *PeerInfoManager) PeerMaxHeight() int64 {
   113  	return atomic.LoadInt64(&p.maxHeight)
   114  }
   115  
   116  func (p *PeerInfoManager) start() {
   117  	for {
   118  		select {
   119  		case <-p.ctx.Done():
   120  			return
   121  		case <-time.After(time.Second * 30):
   122  			p.prune()
   123  		}
   124  	}
   125  }
   126  func (p *PeerInfoManager) prune() {
   127  	p.peerInfo.Range(func(key interface{}, value interface{}) bool {
   128  		info := value.(*peerStoreInfo)
   129  		if time.Since(info.storeTime) > time.Minute {
   130  			p.peerInfo.Delete(key)
   131  			return true
   132  		}
   133  		return true
   134  	})
   135  }