les: move randselect to utilities package

weighted_random_selector is a helpful tool for randomly select
items maintained in a set but based on the item weight.

It's used anywhere is LES package, mainly by les client but will
be used in les server with very high chance. So move it into a
common package as the second step for les separation.

les/distributor.go

@@ -22,6 +22,7 @@ import (
 	"time"
 
 	"github.com/ethereum/go-ethereum/common/mclock"
+	"github.com/ethereum/go-ethereum/les/utilities"
 )
 
 // requestDistributor implements a mechanism that distributes requests to
@@ -194,7 +195,7 @@ func (d *requestDistributor) nextRequest() (distPeer, *distReq, time.Duration) {
 	elem := d.reqQueue.Front()
 	var (
 		bestWait time.Duration
-		sel      *weightedRandomSelect
+		sel      *utilities.WeightedRandomSelect
 	)
 
 	d.peerLock.RLock()
@@ -219,9 +220,9 @@ func (d *requestDistributor) nextRequest() (distPeer, *distReq, time.Duration) {
 				wait, bufRemain := peer.waitBefore(cost)
 				if wait == 0 {
 					if sel == nil {
-						sel = newWeightedRandomSelect()
+						sel = utilities.NewWeightedRandomSelect()
 					}
-					sel.update(selectPeerItem{peer: peer, req: req, weight: int64(bufRemain*1000000) + 1})
+					sel.Update(selectPeerItem{peer: peer, req: req, weight: int64(bufRemain*1000000) + 1})
 				} else {
 					if bestWait == 0 || wait < bestWait {
 						bestWait = wait
@@ -239,7 +240,7 @@ func (d *requestDistributor) nextRequest() (distPeer, *distReq, time.Duration) {
 	}
 
 	if sel != nil {
-		c := sel.choose().(selectPeerItem)
+		c := sel.Choose().(selectPeerItem)
 		return c.peer, c.req, 0
 	}
 	return nil, nil, bestWait

les/serverpool.go

@@ -30,6 +30,7 @@ import (
 	"github.com/ethereum/go-ethereum/common/mclock"
 	"github.com/ethereum/go-ethereum/crypto"
 	"github.com/ethereum/go-ethereum/ethdb"
+	"github.com/ethereum/go-ethereum/les/utilities"
 	"github.com/ethereum/go-ethereum/log"
 	"github.com/ethereum/go-ethereum/p2p"
 	"github.com/ethereum/go-ethereum/p2p/discv5"
@@ -129,7 +130,7 @@ type serverPool struct {
 	adjustStats          chan poolStatAdjust
 
 	knownQueue, newQueue       poolEntryQueue
-	knownSelect, newSelect     *weightedRandomSelect
+	knownSelect, newSelect     *utilities.WeightedRandomSelect
 	knownSelected, newSelected int
 	fastDiscover               bool
 	connCh                     chan *connReq
@@ -152,8 +153,8 @@ func newServerPool(db ethdb.Database, ulcServers []string) *serverPool {
 		disconnCh:    make(chan *disconnReq),
 		registerCh:   make(chan *registerReq),
 		closeCh:      make(chan struct{}),
-		knownSelect:  newWeightedRandomSelect(),
+		knownSelect:  utilities.NewWeightedRandomSelect(),
-		newSelect:    newWeightedRandomSelect(),
+		newSelect:    utilities.NewWeightedRandomSelect(),
 		fastDiscover: true,
 		trustedNodes: parseTrustedNodes(ulcServers),
 	}
@@ -402,8 +403,8 @@ func (pool *serverPool) eventLoop() {
 				entry.lastConnected = addr
 				entry.addr = make(map[string]*poolEntryAddress)
 				entry.addr[addr.strKey()] = addr
-				entry.addrSelect = *newWeightedRandomSelect()
+				entry.addrSelect = *utilities.NewWeightedRandomSelect()
-				entry.addrSelect.update(addr)
+				entry.addrSelect.Update(addr)
 				req.result <- entry
 			}
 
@@ -459,7 +460,7 @@ func (pool *serverPool) findOrNewNode(node *enode.Node) *poolEntry {
 		entry = &poolEntry{
 			node:       node,
 			addr:       make(map[string]*poolEntryAddress),
-			addrSelect: *newWeightedRandomSelect(),
+			addrSelect: *utilities.NewWeightedRandomSelect(),
 			shortRetry: shortRetryCnt,
 		}
 		pool.entries[node.ID()] = entry
@@ -477,7 +478,7 @@ func (pool *serverPool) findOrNewNode(node *enode.Node) *poolEntry {
 		entry.addr[addr.strKey()] = addr
 	}
 	addr.lastSeen = now
-	entry.addrSelect.update(addr)
+	entry.addrSelect.Update(addr)
 	if !entry.known {
 		pool.newQueue.setLatest(entry)
 	}
@@ -505,7 +506,7 @@ func (pool *serverPool) loadNodes() {
 		pool.entries[e.node.ID()] = e
 		if pool.trustedNodes[e.node.ID()] == nil {
 			pool.knownQueue.setLatest(e)
-			pool.knownSelect.update((*knownEntry)(e))
+			pool.knownSelect.Update((*knownEntry)(e))
 		}
 	}
 }
@@ -556,8 +557,8 @@ func (pool *serverPool) saveNodes() {
 // Note that it is called by the new/known queues from which the entry has already
 // been removed so removing it from the queues is not necessary.
 func (pool *serverPool) removeEntry(entry *poolEntry) {
-	pool.newSelect.remove((*discoveredEntry)(entry))
+	pool.newSelect.Remove((*discoveredEntry)(entry))
-	pool.knownSelect.remove((*knownEntry)(entry))
+	pool.knownSelect.Remove((*knownEntry)(entry))
 	entry.removed = true
 	delete(pool.entries, entry.node.ID())
 }
@@ -586,8 +587,8 @@ func (pool *serverPool) setRetryDial(entry *poolEntry) {
 // updateCheckDial is called when an entry can potentially be dialed again. It updates
 // its selection weights and checks if new dials can/should be made.
 func (pool *serverPool) updateCheckDial(entry *poolEntry) {
-	pool.newSelect.update((*discoveredEntry)(entry))
+	pool.newSelect.Update((*discoveredEntry)(entry))
-	pool.knownSelect.update((*knownEntry)(entry))
+	pool.knownSelect.Update((*knownEntry)(entry))
 	pool.checkDial()
 }
 
@@ -596,7 +597,7 @@ func (pool *serverPool) updateCheckDial(entry *poolEntry) {
 func (pool *serverPool) checkDial() {
 	fillWithKnownSelects := !pool.fastDiscover
 	for pool.knownSelected < targetKnownSelect {
-		entry := pool.knownSelect.choose()
+		entry := pool.knownSelect.Choose()
 		if entry == nil {
 			fillWithKnownSelects = false
 			break
@@ -604,7 +605,7 @@ func (pool *serverPool) checkDial() {
 		pool.dial((*poolEntry)(entry.(*knownEntry)), true)
 	}
 	for pool.knownSelected+pool.newSelected < targetServerCount {
-		entry := pool.newSelect.choose()
+		entry := pool.newSelect.Choose()
 		if entry == nil {
 			break
 		}
@@ -615,7 +616,7 @@ func (pool *serverPool) checkDial() {
 		// is over, we probably won't find more in the near future so select more
 		// known entries if possible
 		for pool.knownSelected < targetServerCount {
-			entry := pool.knownSelect.choose()
+			entry := pool.knownSelect.Choose()
 			if entry == nil {
 				break
 			}
@@ -636,7 +637,7 @@ func (pool *serverPool) dial(entry *poolEntry, knownSelected bool) {
 	} else {
 		pool.newSelected++
 	}
-	addr := entry.addrSelect.choose().(*poolEntryAddress)
+	addr := entry.addrSelect.Choose().(*poolEntryAddress)
 	log.Debug("Dialing new peer", "lesaddr", entry.node.ID().String()+"@"+addr.strKey(), "set", len(entry.addr), "known", knownSelected)
 	entry.dialed = addr
 	go func() {
@@ -684,7 +685,7 @@ type poolEntry struct {
 	addr                  map[string]*poolEntryAddress
 	node                  *enode.Node
 	lastConnected, dialed *poolEntryAddress
-	addrSelect            weightedRandomSelect
+	addrSelect            utilities.WeightedRandomSelect
 
 	lastDiscovered                mclock.AbsTime
 	known, knownSelected, trusted bool
@@ -734,8 +735,8 @@ func (e *poolEntry) DecodeRLP(s *rlp.Stream) error {
 	e.node = enode.NewV4(pubkey, entry.IP, int(entry.Port), int(entry.Port))
 	e.addr = make(map[string]*poolEntryAddress)
 	e.addr[addr.strKey()] = addr
-	e.addrSelect = *newWeightedRandomSelect()
+	e.addrSelect = *utilities.NewWeightedRandomSelect()
-	e.addrSelect.update(addr)
+	e.addrSelect.Update(addr)
 	e.lastConnected = addr
 	e.connectStats = entry.CStat
 	e.delayStats = entry.DStat

les/utilities/weighted_select.go

@@ -14,43 +14,30 @@
 // You should have received a copy of the GNU Lesser General Public License
 // along with the go-ethereum library. If not, see <http://www.gnu.org/licenses/>.
 
-package les
+package utilities
 
-import (
+import "math/rand"
-	"math/rand"
-)
 
 // wrsItem interface should be implemented by any entries that are to be selected from
-// a weightedRandomSelect set. Note that recalculating monotonously decreasing item
+// a WeightedRandomSelect set. Note that recalculating monotonously decreasing item
-// weights on-demand (without constantly calling update) is allowed
+// weights on-demand (without constantly calling Update) is allowed
 type wrsItem interface {
 	Weight() int64
 }
 
-// weightedRandomSelect is capable of weighted random selection from a set of items
+// WeightedRandomSelect is capable of weighted random selection from a set of items
-type weightedRandomSelect struct {
+type WeightedRandomSelect struct {
 	root *wrsNode
 	idx  map[wrsItem]int
 }
 
-// newWeightedRandomSelect returns a new weightedRandomSelect structure
+// NewWeightedRandomSelect returns a new WeightedRandomSelect structure
-func newWeightedRandomSelect() *weightedRandomSelect {
+func NewWeightedRandomSelect() *WeightedRandomSelect {
-	return &weightedRandomSelect{root: &wrsNode{maxItems: wrsBranches}, idx: make(map[wrsItem]int)}
+	return &WeightedRandomSelect{root: &wrsNode{maxItems: wrsBranches}, idx: make(map[wrsItem]int)}
-}
-
-// update updates an item's weight, adds it if it was non-existent or removes it if
-// the new weight is zero. Note that explicitly updating decreasing weights is not necessary.
-func (w *weightedRandomSelect) update(item wrsItem) {
-	w.setWeight(item, item.Weight())
-}
-
-// remove removes an item from the set
-func (w *weightedRandomSelect) remove(item wrsItem) {
-	w.setWeight(item, 0)
 }
 
 // setWeight sets an item's weight to a specific value (removes it if zero)
-func (w *weightedRandomSelect) setWeight(item wrsItem, weight int64) {
+func (w *WeightedRandomSelect) setWeight(item wrsItem, weight int64) {
 	idx, ok := w.idx[item]
 	if ok {
 		w.root.setWeight(idx, weight)
@@ -71,11 +58,22 @@ func (w *weightedRandomSelect) setWeight(item wrsItem, weight int64) {
 	}
 }
 
-// choose randomly selects an item from the set, with a chance proportional to its
+// Update updates an item's weight, adds it if it was non-existent or removes it if
+// the new weight is zero. Note that explicitly updating decreasing weights is not necessary.
+func (w *WeightedRandomSelect) Update(item wrsItem) {
+	w.setWeight(item, item.Weight())
+}
+
+// Remove removes an item from the set
+func (w *WeightedRandomSelect) Remove(item wrsItem) {
+	w.setWeight(item, 0)
+}
+
+// Choose randomly selects an item from the set, with a chance proportional to its
 // current weight. If the weight of the chosen element has been decreased since the
 // last stored value, returns it with a newWeight/oldWeight chance, otherwise just
 // updates its weight and selects another one
-func (w *weightedRandomSelect) choose() wrsItem {
+func (w *WeightedRandomSelect) Choose() wrsItem {
 	for {
 		if w.root.sumWeight == 0 {
 			return nil
@@ -154,7 +152,7 @@ func (n *wrsNode) setWeight(idx int, weight int64) int64 {
 	return diff
 }
 
-// choose recursively selects an item from the tree and returns it along with its weight
+// Choose recursively selects an item from the tree and returns it along with its weight
 func (n *wrsNode) choose(val int64) (wrsItem, int64) {
 	for i, w := range n.weights {
 		if val < w {

les/utilities/weighted_select_test.go

@@ -14,7 +14,7 @@
 // You should have received a copy of the GNU Lesser General Public License
 // along with the go-ethereum library. If not, see <http://www.gnu.org/licenses/>.
 
-package les
+package utilities
 
 import (
 	"math/rand"
@@ -36,15 +36,15 @@ func (t *testWrsItem) Weight() int64 {
 
 func TestWeightedRandomSelect(t *testing.T) {
 	testFn := func(cnt int) {
-		s := newWeightedRandomSelect()
+		s := NewWeightedRandomSelect()
 		w := -1
 		list := make([]testWrsItem, cnt)
 		for i := range list {
 			list[i] = testWrsItem{idx: i, widx: &w}
-			s.update(&list[i])
+			s.Update(&list[i])
 		}
 		w = rand.Intn(cnt)
-		c := s.choose()
+		c := s.Choose()
 		if c == nil {
 			t.Errorf("expected item, got nil")
 		} else {
@@ -53,7 +53,7 @@ func TestWeightedRandomSelect(t *testing.T) {
 			}
 		}
 		w = -2
-		if s.choose() != nil {
+		if s.Choose() != nil {
 			t.Errorf("expected nil, got item")
 		}
 	}