github.com/DxChainNetwork/dxc@v0.8.1-0.20220824085222-1162e304b6e7/les/vflux/client/wrsiterator.go (about)

     1  // Copyright 2020 The go-ethereum Authors
     2  // This file is part of the go-ethereum library.
     3  //
     4  // The go-ethereum library is free software: you can redistribute it and/or modify
     5  // it under the terms of the GNU Lesser General Public License as published by
     6  // the Free Software Foundation, either version 3 of the License, or
     7  // (at your option) any later version.
     8  //
     9  // The go-ethereum library is distributed in the hope that it will be useful,
    10  // but WITHOUT ANY WARRANTY; without even the implied warranty of
    11  // MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
    12  // GNU Lesser General Public License for more details.
    13  //
    14  // You should have received a copy of the GNU Lesser General Public License
    15  // along with the go-ethereum library. If not, see <http://www.gnu.org/licenses/>.
    16  
    17  package client
    18  
    19  import (
    20  	"sync"
    21  
    22  	"github.com/DxChainNetwork/dxc/les/utils"
    23  	"github.com/DxChainNetwork/dxc/p2p/enode"
    24  	"github.com/DxChainNetwork/dxc/p2p/nodestate"
    25  )
    26  
    27  // WrsIterator returns nodes from the specified selectable set with a weighted random
    28  // selection. Selection weights are provided by a callback function.
    29  type WrsIterator struct {
    30  	lock sync.Mutex
    31  	cond *sync.Cond
    32  
    33  	ns       *nodestate.NodeStateMachine
    34  	wrs      *utils.WeightedRandomSelect
    35  	nextNode *enode.Node
    36  	closed   bool
    37  }
    38  
    39  // NewWrsIterator creates a new WrsIterator. Nodes are selectable if they have all the required
    40  // and none of the disabled flags set. When a node is selected the selectedFlag is set which also
    41  // disables further selectability until it is removed or times out.
    42  func NewWrsIterator(ns *nodestate.NodeStateMachine, requireFlags, disableFlags nodestate.Flags, weightField nodestate.Field) *WrsIterator {
    43  	wfn := func(i interface{}) uint64 {
    44  		n := ns.GetNode(i.(enode.ID))
    45  		if n == nil {
    46  			return 0
    47  		}
    48  		wt, _ := ns.GetField(n, weightField).(uint64)
    49  		return wt
    50  	}
    51  
    52  	w := &WrsIterator{
    53  		ns:  ns,
    54  		wrs: utils.NewWeightedRandomSelect(wfn),
    55  	}
    56  	w.cond = sync.NewCond(&w.lock)
    57  
    58  	ns.SubscribeField(weightField, func(n *enode.Node, state nodestate.Flags, oldValue, newValue interface{}) {
    59  		if state.HasAll(requireFlags) && state.HasNone(disableFlags) {
    60  			w.lock.Lock()
    61  			w.wrs.Update(n.ID())
    62  			w.lock.Unlock()
    63  			w.cond.Signal()
    64  		}
    65  	})
    66  
    67  	ns.SubscribeState(requireFlags.Or(disableFlags), func(n *enode.Node, oldState, newState nodestate.Flags) {
    68  		oldMatch := oldState.HasAll(requireFlags) && oldState.HasNone(disableFlags)
    69  		newMatch := newState.HasAll(requireFlags) && newState.HasNone(disableFlags)
    70  		if newMatch == oldMatch {
    71  			return
    72  		}
    73  
    74  		w.lock.Lock()
    75  		if newMatch {
    76  			w.wrs.Update(n.ID())
    77  		} else {
    78  			w.wrs.Remove(n.ID())
    79  		}
    80  		w.lock.Unlock()
    81  		w.cond.Signal()
    82  	})
    83  	return w
    84  }
    85  
    86  // Next selects the next node.
    87  func (w *WrsIterator) Next() bool {
    88  	w.nextNode = w.chooseNode()
    89  	return w.nextNode != nil
    90  }
    91  
    92  func (w *WrsIterator) chooseNode() *enode.Node {
    93  	w.lock.Lock()
    94  	defer w.lock.Unlock()
    95  
    96  	for {
    97  		for !w.closed && w.wrs.IsEmpty() {
    98  			w.cond.Wait()
    99  		}
   100  		if w.closed {
   101  			return nil
   102  		}
   103  		// Choose the next node at random. Even though w.wrs is guaranteed
   104  		// non-empty here, Choose might return nil if all items have weight
   105  		// zero.
   106  		if c := w.wrs.Choose(); c != nil {
   107  			id := c.(enode.ID)
   108  			w.wrs.Remove(id)
   109  			return w.ns.GetNode(id)
   110  		}
   111  	}
   112  
   113  }
   114  
   115  // Close ends the iterator.
   116  func (w *WrsIterator) Close() {
   117  	w.lock.Lock()
   118  	w.closed = true
   119  	w.lock.Unlock()
   120  	w.cond.Signal()
   121  }
   122  
   123  // Node returns the current node.
   124  func (w *WrsIterator) Node() *enode.Node {
   125  	w.lock.Lock()
   126  	defer w.lock.Unlock()
   127  	return w.nextNode
   128  }