les: implement light server pool

les/handler.go

@@ -22,10 +22,12 @@ import (
 	"errors"
 	"fmt"
 	"math/big"
+	"net"
 	"sync"
 	"time"
 
 	"github.com/ethereum/go-ethereum/common"
+	"github.com/ethereum/go-ethereum/common/mclock"
 	"github.com/ethereum/go-ethereum/core"
 	"github.com/ethereum/go-ethereum/core/state"
 	"github.com/ethereum/go-ethereum/core/types"
@@ -101,10 +103,7 @@ type ProtocolManager struct {
 	chainDb     ethdb.Database
 	odr         *LesOdr
 	server      *LesServer
-
-	topicDisc *discv5.Network
-	lesTopic  discv5.Topic
-	p2pServer *p2p.Server
+	serverPool  *serverPool
 
 	downloader *downloader.Downloader
 	fetcher    *lightFetcher
@@ -157,13 +156,46 @@ func NewProtocolManager(chainConfig *params.ChainConfig, lightSync bool, network
 			Version: version,
 			Length:  ProtocolLengths[i],
 			Run: func(p *p2p.Peer, rw p2p.MsgReadWriter) error {
+				var entry *poolEntry
+				if manager.serverPool != nil {
+					addr := p.RemoteAddr().(*net.TCPAddr)
+					entry = manager.serverPool.connect(p.ID(), addr.IP, uint16(addr.Port))
+					if entry == nil {
+						return fmt.Errorf("unwanted connection")
+					}
+				}
 				peer := manager.newPeer(int(version), networkId, p, rw)
+				peer.poolEntry = entry
 				select {
 				case manager.newPeerCh <- peer:
 					manager.wg.Add(1)
 					defer manager.wg.Done()
-					return manager.handle(peer)
+					start := mclock.Now()
+					err := manager.handle(peer)
+					if entry != nil {
+						connTime := time.Duration(mclock.Now() - start)
+						stopped := false
+						select {
+						case <-manager.quitSync:
+							stopped = true
+						default:
+						}
+						//fmt.Println("connTime", peer.id, connTime, stopped, err)
+						quality := float64(1)
+						setQuality := true
+						if connTime < time.Minute*10 {
+							quality = 0
+							if stopped {
+								setQuality = false
+							}
+						}
+						manager.serverPool.disconnect(entry, quality, setQuality)
+					}
+					return err
 				case <-manager.quitSync:
+					if entry != nil {
+						manager.serverPool.disconnect(entry, 0, false)
+					}
 					return p2p.DiscQuitting
 				}
 			},
@@ -236,54 +268,24 @@ func (pm *ProtocolManager) removePeer(id string) {
 	}
 }
 
-func (pm *ProtocolManager) findServers() {
-	if pm.p2pServer == nil || pm.topicDisc == nil {
-		return
-	}
-	glog.V(logger.Debug).Infoln("Looking for topic", string(pm.lesTopic))
-	enodes := make(chan string, 100)
-	stop := make(chan struct{})
-	go pm.topicDisc.SearchTopic(pm.lesTopic, stop, enodes)
-	go func() {
-		added := make(map[string]bool)
-		for {
-			select {
-			case enode := <-enodes:
-				if !added[enode] {
-					glog.V(logger.Info).Infoln("Found LES server:", enode)
-					added[enode] = true
-					if node, err := discover.ParseNode(enode); err == nil {
-						pm.p2pServer.AddPeer(node)
-					}
-				}
-			case <-stop:
-				return
-			}
-		}
-	}()
-	select {
-	case <-time.After(time.Second * 20):
-	case <-pm.quitSync:
-	}
-	close(stop)
-}
-
 func (pm *ProtocolManager) Start(srvr *p2p.Server) {
-	pm.p2pServer = srvr
+	var topicDisc *discv5.Network
 	if srvr != nil {
-		pm.topicDisc = srvr.DiscV5
+		topicDisc = srvr.DiscV5
 	}
-	pm.lesTopic = discv5.Topic("LES@" + common.Bytes2Hex(pm.blockchain.Genesis().Hash().Bytes()[0:8]))
+	lesTopic := discv5.Topic("LES@" + common.Bytes2Hex(pm.blockchain.Genesis().Hash().Bytes()[0:8]))
 	if pm.lightSync {
 		// start sync handler
-		go pm.findServers()
+		if srvr != nil {
+			pm.serverPool = newServerPool(pm.chainDb, []byte("serverPool/"), srvr, lesTopic, pm.quitSync, &pm.wg)
+		}
 		go pm.syncer()
 	} else {
-		if pm.topicDisc != nil {
+		if topicDisc != nil {
 			go func() {
-				glog.V(logger.Debug).Infoln("Starting registering topic", string(pm.lesTopic))
-				pm.topicDisc.RegisterTopic(pm.lesTopic, pm.quitSync)
-				glog.V(logger.Debug).Infoln("Stopped registering topic", string(pm.lesTopic))
+				glog.V(logger.Debug).Infoln("Starting registering topic", string(lesTopic))
+				topicDisc.RegisterTopic(lesTopic, pm.quitSync)
+				glog.V(logger.Debug).Infoln("Stopped registering topic", string(lesTopic))
 			}()
 		}
 		go func() {
@@ -373,6 +375,10 @@ func (pm *ProtocolManager) handle(p *peer) error {
 		}
 
 		pm.fetcher.notify(p, nil)
+
+		if p.poolEntry != nil {
+			pm.serverPool.registered(p.poolEntry)
+		}
 	}
 
 	stop := make(chan struct{})

les/peer.go

@@ -57,6 +57,8 @@ type peer struct {
 
 	announceChn chan announceData
 
+	poolEntry *poolEntry
+
 	fcClient       *flowcontrol.ClientNode // nil if the peer is server only
 	fcServer       *flowcontrol.ServerNode // nil if the peer is client only
 	fcServerParams *flowcontrol.ServerParams

les/randselect.go

+// Copyright 2016 The go-ethereum Authors
+// This file is part of the go-ethereum library.
+//
+// The go-ethereum library is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Lesser General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// The go-ethereum library is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU Lesser General Public License for more details.
+//
+// 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 implements the Light Ethereum Subprotocol.
+package les
+
+import (
+	"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
+// 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
+type weightedRandomSelect struct {
+	root *wrsNode
+	idx  map[wrsItem]int
+}
+
+// newWeightedRandomSelect returns a new weightedRandomSelect structure
+func newWeightedRandomSelect() *weightedRandomSelect {
+	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) {
+	idx, ok := w.idx[item]
+	if ok {
+		w.root.setWeight(idx, weight)
+		if weight == 0 {
+			delete(w.idx, item)
+		}
+	} else {
+		if weight != 0 {
+			if w.root.itemCnt == w.root.maxItems {
+				// add a new level
+				newRoot := &wrsNode{sumWeight: w.root.sumWeight, itemCnt: w.root.itemCnt, level: w.root.level + 1, maxItems: w.root.maxItems * wrsBranches}
+				newRoot.items[0] = w.root
+				newRoot.weights[0] = w.root.sumWeight
+				w.root = newRoot
+			}
+			w.idx[item] = w.root.insert(item, weight)
+		}
+	}
+}
+
+// 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 {
+	for {
+		if w.root.sumWeight == 0 {
+			return nil
+		}
+		val := rand.Int63n(w.root.sumWeight)
+		choice, lastWeight := w.root.choose(val)
+		weight := choice.Weight()
+		if weight != lastWeight {
+			w.setWeight(choice, weight)
+		}
+		if weight >= lastWeight || rand.Int63n(lastWeight) < weight {
+			return choice
+		}
+	}
+}
+
+const wrsBranches = 8 // max number of branches in the wrsNode tree
+
+// wrsNode is a node of a tree structure that can store wrsItems or further wrsNodes.
+type wrsNode struct {
+	items                    [wrsBranches]interface{}
+	weights                  [wrsBranches]int64
+	sumWeight                int64
+	level, itemCnt, maxItems int
+}
+
+// insert recursively inserts a new item to the tree and returns the item index
+func (n *wrsNode) insert(item wrsItem, weight int64) int {
+	branch := 0
+	for n.items[branch] != nil && (n.level == 0 || n.items[branch].(*wrsNode).itemCnt == n.items[branch].(*wrsNode).maxItems) {
+		branch++
+		if branch == wrsBranches {
+			panic(nil)
+		}
+	}
+	n.itemCnt++
+	n.sumWeight += weight
+	n.weights[branch] += weight
+	if n.level == 0 {
+		n.items[branch] = item
+		return branch
+	} else {
+		var subNode *wrsNode
+		if n.items[branch] == nil {
+			subNode = &wrsNode{maxItems: n.maxItems / wrsBranches, level: n.level - 1}
+			n.items[branch] = subNode
+		} else {
+			subNode = n.items[branch].(*wrsNode)
+		}
+		subIdx := subNode.insert(item, weight)
+		return subNode.maxItems*branch + subIdx
+	}
+}
+
+// setWeight updates the weight of a certain item (which should exist) and returns
+// the change of the last weight value stored in the tree
+func (n *wrsNode) setWeight(idx int, weight int64) int64 {
+	if n.level == 0 {
+		oldWeight := n.weights[idx]
+		n.weights[idx] = weight
+		diff := weight - oldWeight
+		n.sumWeight += diff
+		if weight == 0 {
+			n.items[idx] = nil
+			n.itemCnt--
+		}
+		return diff
+	}
+	branchItems := n.maxItems / wrsBranches
+	branch := idx / branchItems
+	diff := n.items[branch].(*wrsNode).setWeight(idx-branch*branchItems, weight)
+	n.weights[branch] += diff
+	n.sumWeight += diff
+	if weight == 0 {
+		n.itemCnt--
+	}
+	return diff
+}
+
+// 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 {
+			if n.level == 0 {
+				return n.items[i].(wrsItem), n.weights[i]
+			} else {
+				return n.items[i].(*wrsNode).choose(val)
+			}
+		} else {
+			val -= w
+		}
+	}
+	panic(nil)
+}

les/randselect_test.go

+// Copyright 2016 The go-ethereum Authors
+// This file is part of the go-ethereum library.
+//
+// The go-ethereum library is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Lesser General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// The go-ethereum library is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU Lesser General Public License for more details.
+//
+// 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
+
+import (
+	"math/rand"
+	"testing"
+)
+
+type testWrsItem struct {
+	idx  int
+	widx *int
+}
+
+func (t *testWrsItem) Weight() int64 {
+	w := *t.widx
+	if w == -1 || w == t.idx {
+		return int64(t.idx + 1)
+	}
+	return 0
+}
+
+func TestWeightedRandomSelect(t *testing.T) {
+	testFn := func(cnt int) {
+		s := newWeightedRandomSelect()
+		w := -1
+		list := make([]testWrsItem, cnt)
+		for i, _ := range list {
+			list[i] = testWrsItem{idx: i, widx: &w}
+			s.update(&list[i])
+		}
+		w = rand.Intn(cnt)
+		c := s.choose()
+		if c == nil {
+			t.Errorf("expected item, got nil")
+		} else {
+			if c.(*testWrsItem).idx != w {
+				t.Errorf("expected another item")
+			}
+		}
+		w = -2
+		if s.choose() != nil {
+			t.Errorf("expected nil, got item")
+		}
+	}
+	testFn(1)
+	testFn(10)
+	testFn(100)
+	testFn(1000)
+	testFn(10000)
+	testFn(100000)
+	testFn(1000000)
+}