github.com/arieschain/arieschain@v0.0.0-20191023063405-37c074544356/les/txrelay.go (about)

     1  package les
     2  
     3  import (
     4  	"sync"
     5  
     6  	"github.com/quickchainproject/quickchain/common"
     7  	"github.com/quickchainproject/quickchain/core/types"
     8  )
     9  
    10  type ltrInfo struct {
    11  	tx     *types.Transaction
    12  	sentTo map[*peer]struct{}
    13  }
    14  
    15  type LesTxRelay struct {
    16  	txSent       map[common.Hash]*ltrInfo
    17  	txPending    map[common.Hash]struct{}
    18  	ps           *peerSet
    19  	peerList     []*peer
    20  	peerStartPos int
    21  	lock         sync.RWMutex
    22  
    23  	reqDist *requestDistributor
    24  }
    25  
    26  func NewLesTxRelay(ps *peerSet, reqDist *requestDistributor) *LesTxRelay {
    27  	r := &LesTxRelay{
    28  		txSent:    make(map[common.Hash]*ltrInfo),
    29  		txPending: make(map[common.Hash]struct{}),
    30  		ps:        ps,
    31  		reqDist:   reqDist,
    32  	}
    33  	ps.notify(r)
    34  	return r
    35  }
    36  
    37  func (self *LesTxRelay) registerPeer(p *peer) {
    38  	self.lock.Lock()
    39  	defer self.lock.Unlock()
    40  
    41  	self.peerList = self.ps.AllPeers()
    42  }
    43  
    44  func (self *LesTxRelay) unregisterPeer(p *peer) {
    45  	self.lock.Lock()
    46  	defer self.lock.Unlock()
    47  
    48  	self.peerList = self.ps.AllPeers()
    49  }
    50  
    51  // send sends a list of transactions to at most a given number of peers at
    52  // once, never resending any particular transaction to the same peer twice
    53  func (self *LesTxRelay) send(txs types.Transactions, count int) {
    54  	sendTo := make(map[*peer]types.Transactions)
    55  
    56  	self.peerStartPos++ // rotate the starting position of the peer list
    57  	if self.peerStartPos >= len(self.peerList) {
    58  		self.peerStartPos = 0
    59  	}
    60  
    61  	for _, tx := range txs {
    62  		hash := tx.Hash()
    63  		ltr, ok := self.txSent[hash]
    64  		if !ok {
    65  			ltr = &ltrInfo{
    66  				tx:     tx,
    67  				sentTo: make(map[*peer]struct{}),
    68  			}
    69  			self.txSent[hash] = ltr
    70  			self.txPending[hash] = struct{}{}
    71  		}
    72  
    73  		if len(self.peerList) > 0 {
    74  			cnt := count
    75  			pos := self.peerStartPos
    76  			for {
    77  				peer := self.peerList[pos]
    78  				if _, ok := ltr.sentTo[peer]; !ok {
    79  					sendTo[peer] = append(sendTo[peer], tx)
    80  					ltr.sentTo[peer] = struct{}{}
    81  					cnt--
    82  				}
    83  				if cnt == 0 {
    84  					break // sent it to the desired number of peers
    85  				}
    86  				pos++
    87  				if pos == len(self.peerList) {
    88  					pos = 0
    89  				}
    90  				if pos == self.peerStartPos {
    91  					break // tried all available peers
    92  				}
    93  			}
    94  		}
    95  	}
    96  
    97  	for p, list := range sendTo {
    98  		pp := p
    99  		ll := list
   100  
   101  		reqID := genReqID()
   102  		rq := &distReq{
   103  			getCost: func(dp distPeer) uint64 {
   104  				peer := dp.(*peer)
   105  				return peer.GetRequestCost(SendTxMsg, len(ll))
   106  			},
   107  			canSend: func(dp distPeer) bool {
   108  				return dp.(*peer) == pp
   109  			},
   110  			request: func(dp distPeer) func() {
   111  				peer := dp.(*peer)
   112  				cost := peer.GetRequestCost(SendTxMsg, len(ll))
   113  				peer.fcServer.QueueRequest(reqID, cost)
   114  				return func() { peer.SendTxs(reqID, cost, ll) }
   115  			},
   116  		}
   117  		self.reqDist.queue(rq)
   118  	}
   119  }
   120  
   121  func (self *LesTxRelay) Send(txs types.Transactions) {
   122  	self.lock.Lock()
   123  	defer self.lock.Unlock()
   124  
   125  	self.send(txs, 3)
   126  }
   127  
   128  func (self *LesTxRelay) NewHead(head common.Hash, mined []common.Hash, rollback []common.Hash) {
   129  	self.lock.Lock()
   130  	defer self.lock.Unlock()
   131  
   132  	for _, hash := range mined {
   133  		delete(self.txPending, hash)
   134  	}
   135  
   136  	for _, hash := range rollback {
   137  		self.txPending[hash] = struct{}{}
   138  	}
   139  
   140  	if len(self.txPending) > 0 {
   141  		txs := make(types.Transactions, len(self.txPending))
   142  		i := 0
   143  		for hash := range self.txPending {
   144  			txs[i] = self.txSent[hash].tx
   145  			i++
   146  		}
   147  		self.send(txs, 1)
   148  	}
   149  }
   150  
   151  func (self *LesTxRelay) Discard(hashes []common.Hash) {
   152  	self.lock.Lock()
   153  	defer self.lock.Unlock()
   154  
   155  	for _, hash := range hashes {
   156  		delete(self.txSent, hash)
   157  		delete(self.txPending, hash)
   158  	}
   159  }