les: general code refactoring

les/backend.go

@@ -19,6 +19,7 @@ package les
 
 import (
 	"fmt"
+	"sync"
 	"time"
 
 	"github.com/ethereum/go-ethereum/accounts"
@@ -38,6 +39,7 @@ import (
 	"github.com/ethereum/go-ethereum/log"
 	"github.com/ethereum/go-ethereum/node"
 	"github.com/ethereum/go-ethereum/p2p"
+	"github.com/ethereum/go-ethereum/p2p/discv5"
 	"github.com/ethereum/go-ethereum/params"
 	rpc "github.com/ethereum/go-ethereum/rpc"
 )
@@ -49,9 +51,13 @@ type LightEthereum struct {
 	// Channel for shutting down the service
 	shutdownChan chan bool
 	// Handlers
+	peers           *peerSet
 	txPool          *light.TxPool
 	blockchain      *light.LightChain
 	protocolManager *ProtocolManager
+	serverPool      *serverPool
+	reqDist         *requestDistributor
+	retriever       *retrieveManager
 	// DB interfaces
 	chainDb ethdb.Database // Block chain database
 
@@ -63,6 +69,9 @@ type LightEthereum struct {
 
 	networkId     uint64
 	netRPCService *ethapi.PublicNetAPI
+
+	quitSync chan struct{}
+	wg       sync.WaitGroup
 }
 
 func New(ctx *node.ServiceContext, config *eth.Config) (*LightEthereum, error) {
@@ -76,20 +85,26 @@ func New(ctx *node.ServiceContext, config *eth.Config) (*LightEthereum, error) {
 	}
 	log.Info("Initialised chain configuration", "config", chainConfig)
 
-	odr := NewLesOdr(chainDb)
+	peers := newPeerSet()
-	relay := NewLesTxRelay()
+	quitSync := make(chan struct{})
+
 	eth := &LightEthereum{
-		odr:            odr,
-		relay:          relay,
-		chainDb:        chainDb,
 		chainConfig:    chainConfig,
+		chainDb:        chainDb,
 		eventMux:       ctx.EventMux,
+		peers:          peers,
+		reqDist:        newRequestDistributor(peers, quitSync),
 		accountManager: ctx.AccountManager,
 		engine:         eth.CreateConsensusEngine(ctx, config, chainConfig, chainDb),
 		shutdownChan:   make(chan bool),
 		networkId:      config.NetworkId,
 	}
-	if eth.blockchain, err = light.NewLightChain(odr, eth.chainConfig, eth.engine, eth.eventMux); err != nil {
+
+	eth.relay = NewLesTxRelay(peers, eth.reqDist)
+	eth.serverPool = newServerPool(chainDb, quitSync, &eth.wg)
+	eth.retriever = newRetrieveManager(peers, eth.reqDist, eth.serverPool)
+	eth.odr = NewLesOdr(chainDb, eth.retriever)
+	if eth.blockchain, err = light.NewLightChain(eth.odr, eth.chainConfig, eth.engine, eth.eventMux); err != nil {
 		return nil, err
 	}
 	// Rewind the chain in case of an incompatible config upgrade.
@@ -100,13 +115,9 @@ func New(ctx *node.ServiceContext, config *eth.Config) (*LightEthereum, error) {
 	}
 
 	eth.txPool = light.NewTxPool(eth.chainConfig, eth.eventMux, eth.blockchain, eth.relay)
-	lightSync := config.SyncMode == downloader.LightSync
+	if eth.protocolManager, err = NewProtocolManager(eth.chainConfig, true, config.NetworkId, eth.eventMux, eth.engine, eth.peers, eth.blockchain, nil, chainDb, eth.odr, eth.relay, quitSync, &eth.wg); err != nil {
-	if eth.protocolManager, err = NewProtocolManager(eth.chainConfig, lightSync, config.NetworkId, eth.eventMux, eth.engine, eth.blockchain, nil, chainDb, odr, relay); err != nil {
 		return nil, err
 	}
-	relay.ps = eth.protocolManager.peers
-	relay.reqDist = eth.protocolManager.reqDist
-
 	eth.ApiBackend = &LesApiBackend{eth, nil}
 	gpoParams := config.GPO
 	if gpoParams.Default == nil {
@@ -116,6 +127,10 @@ func New(ctx *node.ServiceContext, config *eth.Config) (*LightEthereum, error) {
 	return eth, nil
 }
 
+func lesTopic(genesisHash common.Hash) discv5.Topic {
+	return discv5.Topic("LES@" + common.Bytes2Hex(genesisHash.Bytes()[0:8]))
+}
+
 type LightDummyAPI struct{}
 
 // Etherbase is the address that mining rewards will be send to
@@ -188,7 +203,8 @@ func (s *LightEthereum) Protocols() []p2p.Protocol {
 func (s *LightEthereum) Start(srvr *p2p.Server) error {
 	log.Warn("Light client mode is an experimental feature")
 	s.netRPCService = ethapi.NewPublicNetAPI(srvr, s.networkId)
-	s.protocolManager.Start(srvr)
+	s.serverPool.start(srvr, lesTopic(s.blockchain.Genesis().Hash()))
+	s.protocolManager.Start()
 	return nil
 }
 

les/distributor.go

@@ -34,11 +34,11 @@ var ErrNoPeers = errors.New("no suitable peers available")
 type requestDistributor struct {
 	reqQueue         *list.List
 	lastReqOrder     uint64
+	peers            map[distPeer]struct{}
+	peerLock         sync.RWMutex
 	stopChn, loopChn chan struct{}
 	loopNextSent     bool
 	lock             sync.Mutex
-
-	getAllPeers func() map[distPeer]struct{}
 }
 
 // distPeer is an LES server peer interface for the request distributor.
@@ -71,15 +71,39 @@ type distReq struct {
 }
 
 // newRequestDistributor creates a new request distributor
-func newRequestDistributor(getAllPeers func() map[distPeer]struct{}, stopChn chan struct{}) *requestDistributor {
+func newRequestDistributor(peers *peerSet, stopChn chan struct{}) *requestDistributor {
-	r := &requestDistributor{
+	d := &requestDistributor{
 		reqQueue: list.New(),
 		loopChn:  make(chan struct{}, 2),
 		stopChn:  stopChn,
-		getAllPeers: getAllPeers,
+		peers:    make(map[distPeer]struct{}),
 	}
-	go r.loop()
+	if peers != nil {
-	return r
+		peers.notify(d)
+	}
+	go d.loop()
+	return d
+}
+
+// registerPeer implements peerSetNotify
+func (d *requestDistributor) registerPeer(p *peer) {
+	d.peerLock.Lock()
+	d.peers[p] = struct{}{}
+	d.peerLock.Unlock()
+}
+
+// unregisterPeer implements peerSetNotify
+func (d *requestDistributor) unregisterPeer(p *peer) {
+	d.peerLock.Lock()
+	delete(d.peers, p)
+	d.peerLock.Unlock()
+}
+
+// registerTestPeer adds a new test peer
+func (d *requestDistributor) registerTestPeer(p distPeer) {
+	d.peerLock.Lock()
+	d.peers[p] = struct{}{}
+	d.peerLock.Unlock()
 }
 
 // distMaxWait is the maximum waiting time after which further necessary waiting
@@ -152,8 +176,7 @@ func (sp selectPeerItem) Weight() int64 {
 // nextRequest returns the next possible request from any peer, along with the
 // associated peer and necessary waiting time
 func (d *requestDistributor) nextRequest() (distPeer, *distReq, time.Duration) {
-	peers := d.getAllPeers()
+	checkedPeers := make(map[distPeer]struct{})
-
 	elem := d.reqQueue.Front()
 	var (
 		bestPeer distPeer
@@ -162,11 +185,14 @@ func (d *requestDistributor) nextRequest() (distPeer, *distReq, time.Duration) {
 		sel      *weightedRandomSelect
 	)
 
-	for (len(peers) > 0 || elem == d.reqQueue.Front()) && elem != nil {
+	d.peerLock.RLock()
+	defer d.peerLock.RUnlock()
+
+	for (len(d.peers) > 0 || elem == d.reqQueue.Front()) && elem != nil {
 		req := elem.Value.(*distReq)
 		canSend := false
-		for peer, _ := range peers {
+		for peer, _ := range d.peers {
-			if peer.canQueue() && req.canSend(peer) {
+			if _, ok := checkedPeers[peer]; !ok && peer.canQueue() && req.canSend(peer) {
 				canSend = true
 				cost := req.getCost(peer)
 				wait, bufRemain := peer.waitBefore(cost)
@@ -182,7 +208,7 @@ func (d *requestDistributor) nextRequest() (distPeer, *distReq, time.Duration) {
 						bestWait = wait
 					}
 				}
-				delete(peers, peer)
+				checkedPeers[peer] = struct{}{}
 			}
 		}
 		next := elem.Next()

les/distributor_test.go

@@ -122,20 +122,14 @@ func testRequestDistributor(t *testing.T, resend bool) {
 	stop := make(chan struct{})
 	defer close(stop)
 
+	dist := newRequestDistributor(nil, stop)
 	var peers [testDistPeerCount]*testDistPeer
 	for i, _ := range peers {
 		peers[i] = &testDistPeer{}
 		go peers[i].worker(t, !resend, stop)
+		dist.registerTestPeer(peers[i])
 	}
 
-	dist := newRequestDistributor(func() map[distPeer]struct{} {
-		m := make(map[distPeer]struct{})
-		for _, peer := range peers {
-			m[peer] = struct{}{}
-		}
-		return m
-	}, stop)
-
 	var wg sync.WaitGroup
 
 	for i := 1; i <= testDistReqCount; i++ {

les/fetcher.go

@@ -116,6 +116,7 @@ func newLightFetcher(pm *ProtocolManager) *lightFetcher {
 		syncDone:       make(chan *peer),
 		maxConfirmedTd: big.NewInt(0),
 	}
+	pm.peers.notify(f)
 	go f.syncLoop()
 	return f
 }
@@ -209,8 +210,8 @@ func (f *lightFetcher) syncLoop() {
 	}
 }
 
-// addPeer adds a new peer to the fetcher's peer set
+// registerPeer adds a new peer to the fetcher's peer set
-func (f *lightFetcher) addPeer(p *peer) {
+func (f *lightFetcher) registerPeer(p *peer) {
 	p.lock.Lock()
 	p.hasBlock = func(hash common.Hash, number uint64) bool {
 		return f.peerHasBlock(p, hash, number)
@@ -223,8 +224,8 @@ func (f *lightFetcher) addPeer(p *peer) {
 	f.peers[p] = &fetcherPeerInfo{nodeByHash: make(map[common.Hash]*fetcherTreeNode)}
 }
 
-// removePeer removes a new peer from the fetcher's peer set
+// unregisterPeer removes a new peer from the fetcher's peer set
-func (f *lightFetcher) removePeer(p *peer) {
+func (f *lightFetcher) unregisterPeer(p *peer) {
 	p.lock.Lock()
 	p.hasBlock = nil
 	p.lock.Unlock()

les/handler.go

@@ -125,12 +125,12 @@ type ProtocolManager struct {
 
 	// wait group is used for graceful shutdowns during downloading
 	// and processing
-	wg sync.WaitGroup
+	wg *sync.WaitGroup
 }
 
 // NewProtocolManager returns a new ethereum sub protocol manager. The Ethereum sub protocol manages peers capable
 // with the ethereum network.
-func NewProtocolManager(chainConfig *params.ChainConfig, lightSync bool, networkId uint64, mux *event.TypeMux, engine consensus.Engine, blockchain BlockChain, txpool txPool, chainDb ethdb.Database, odr *LesOdr, txrelay *LesTxRelay) (*ProtocolManager, error) {
+func NewProtocolManager(chainConfig *params.ChainConfig, lightSync bool, networkId uint64, mux *event.TypeMux, engine consensus.Engine, peers *peerSet, blockchain BlockChain, txpool txPool, chainDb ethdb.Database, odr *LesOdr, txrelay *LesTxRelay, quitSync chan struct{}, wg *sync.WaitGroup) (*ProtocolManager, error) {
 	// Create the protocol manager with the base fields
 	manager := &ProtocolManager{
 		lightSync:   lightSync,
@@ -138,15 +138,20 @@ func NewProtocolManager(chainConfig *params.ChainConfig, lightSync bool, network
 		blockchain:  blockchain,
 		chainConfig: chainConfig,
 		chainDb:     chainDb,
+		odr:         odr,
 		networkId:   networkId,
 		txpool:      txpool,
 		txrelay:     txrelay,
-		odr:         odr,
+		peers:       peers,
-		peers:       newPeerSet(),
 		newPeerCh:   make(chan *peer),
-		quitSync:    make(chan struct{}),
+		quitSync:    quitSync,
+		wg:          wg,
 		noMorePeers: make(chan struct{}),
 	}
+	if odr != nil {
+		manager.retriever = odr.retriever
+		manager.reqDist = odr.retriever.dist
+	}
 	// Initiate a sub-protocol for every implemented version we can handle
 	manager.SubProtocols = make([]p2p.Protocol, 0, len(ProtocolVersions))
 	for i, version := range ProtocolVersions {
@@ -204,77 +209,22 @@ func NewProtocolManager(chainConfig *params.ChainConfig, lightSync bool, network
 		manager.downloader = downloader.New(downloader.LightSync, chainDb, manager.eventMux, blockchain.HasHeader, nil, blockchain.GetHeaderByHash,
 			nil, blockchain.CurrentHeader, nil, nil, nil, blockchain.GetTdByHash,
 			blockchain.InsertHeaderChain, nil, nil, blockchain.Rollback, removePeer)
-
+		manager.peers.notify((*downloaderPeerNotify)(manager))
-		manager.reqDist = newRequestDistributor(func() map[distPeer]struct{} {
-			m := make(map[distPeer]struct{})
-			peers := manager.peers.AllPeers()
-			for _, peer := range peers {
-				m[peer] = struct{}{}
-			}
-			return m
-		}, manager.quitSync)
-
-		manager.lesTopic = discv5.Topic("LES@" + common.Bytes2Hex(manager.blockchain.Genesis().Hash().Bytes()[0:8]))
-		manager.serverPool = newServerPool(manager.chainDb, []byte("serverPool/"), manager.lesTopic, manager.quitSync, &manager.wg)
-		manager.retriever = newRetrieveManager(manager.reqDist, manager.serverPool, removePeer)
-		if odr != nil {
-			odr.retriever = manager.retriever
-		}
 		manager.fetcher = newLightFetcher(manager)
 	}
 
 	return manager, nil
 }
 
+// removePeer initiates disconnection from a peer by removing it from the peer set
 func (pm *ProtocolManager) removePeer(id string) {
-	// Short circuit if the peer was already removed
+	pm.peers.Unregister(id)
-	peer := pm.peers.Peer(id)
-	if peer == nil {
-		return
-	}
-	log.Debug("Removing light Ethereum peer", "peer", id)
-	if err := pm.peers.Unregister(id); err != nil {
-		if err == errNotRegistered {
-			return
-		}
-	}
-	// Unregister the peer from the downloader and Ethereum peer set
-	if pm.lightSync {
-		pm.downloader.UnregisterPeer(id)
-		if pm.txrelay != nil {
-			pm.txrelay.removePeer(id)
-		}
-		if pm.fetcher != nil {
-			pm.fetcher.removePeer(peer)
-		}
-	}
-	// Hard disconnect at the networking layer
-	if peer != nil {
-		peer.Peer.Disconnect(p2p.DiscUselessPeer)
-	}
 }
 
-func (pm *ProtocolManager) Start(srvr *p2p.Server) {
+func (pm *ProtocolManager) Start() {
-	var topicDisc *discv5.Network
-	if srvr != nil {
-		topicDisc = srvr.DiscV5
-	}
 	if pm.lightSync {
-		// start sync handler
-		if srvr != nil { // srvr is nil during testing
-			pm.serverPool.setServer(srvr)
-		}
 		go pm.syncer()
 	} else {
-		if topicDisc != nil {
-			go func() {
-				logger := log.New("topic", pm.lesTopic)
-				logger.Info("Starting topic registration")
-				defer logger.Info("Terminated topic registration")
-
-				topicDisc.RegisterTopic(pm.lesTopic, pm.quitSync)
-			}()
-		}
 		go func() {
 			for range pm.newPeerCh {
 			}
@@ -337,65 +287,10 @@ func (pm *ProtocolManager) handle(p *peer) error {
 	}()
 	// Register the peer in the downloader. If the downloader considers it banned, we disconnect
 	if pm.lightSync {
-		requestHeadersByHash := func(origin common.Hash, amount int, skip int, reverse bool) error {
-			reqID := getNextReqID()
-			rq := &distReq{
-				getCost: func(dp distPeer) uint64 {
-					peer := dp.(*peer)
-					return peer.GetRequestCost(GetBlockHeadersMsg, amount)
-				},
-				canSend: func(dp distPeer) bool {
-					return dp.(*peer) == p
-				},
-				request: func(dp distPeer) func() {
-					peer := dp.(*peer)
-					cost := peer.GetRequestCost(GetBlockHeadersMsg, amount)
-					peer.fcServer.QueueRequest(reqID, cost)
-					return func() { peer.RequestHeadersByHash(reqID, cost, origin, amount, skip, reverse) }
-				},
-			}
-			_, ok := <-pm.reqDist.queue(rq)
-			if !ok {
-				return ErrNoPeers
-			}
-			return nil
-		}
-		requestHeadersByNumber := func(origin uint64, amount int, skip int, reverse bool) error {
-			reqID := getNextReqID()
-			rq := &distReq{
-				getCost: func(dp distPeer) uint64 {
-					peer := dp.(*peer)
-					return peer.GetRequestCost(GetBlockHeadersMsg, amount)
-				},
-				canSend: func(dp distPeer) bool {
-					return dp.(*peer) == p
-				},
-				request: func(dp distPeer) func() {
-					peer := dp.(*peer)
-					cost := peer.GetRequestCost(GetBlockHeadersMsg, amount)
-					peer.fcServer.QueueRequest(reqID, cost)
-					return func() { peer.RequestHeadersByNumber(reqID, cost, origin, amount, skip, reverse) }
-				},
-			}
-			_, ok := <-pm.reqDist.queue(rq)
-			if !ok {
-				return ErrNoPeers
-			}
-			return nil
-		}
-		if err := pm.downloader.RegisterPeer(p.id, ethVersion, p.HeadAndTd,
-			requestHeadersByHash, requestHeadersByNumber, nil, nil, nil); err != nil {
-			return err
-		}
-		if pm.txrelay != nil {
-			pm.txrelay.addPeer(p)
-		}
-
 		p.lock.Lock()
 		head := p.headInfo
 		p.lock.Unlock()
 		if pm.fetcher != nil {
-			pm.fetcher.addPeer(p)
 			pm.fetcher.announce(p, head)
 		}
 
@@ -941,3 +836,64 @@ func (self *ProtocolManager) NodeInfo() *eth.EthNodeInfo {
 		Head:       self.blockchain.LastBlockHash(),
 	}
 }
+
+// downloaderPeerNotify implements peerSetNotify
+type downloaderPeerNotify ProtocolManager
+
+func (d *downloaderPeerNotify) registerPeer(p *peer) {
+	pm := (*ProtocolManager)(d)
+
+	requestHeadersByHash := func(origin common.Hash, amount int, skip int, reverse bool) error {
+		reqID := getNextReqID()
+		rq := &distReq{
+			getCost: func(dp distPeer) uint64 {
+				peer := dp.(*peer)
+				return peer.GetRequestCost(GetBlockHeadersMsg, amount)
+			},
+			canSend: func(dp distPeer) bool {
+				return dp.(*peer) == p
+			},
+			request: func(dp distPeer) func() {
+				peer := dp.(*peer)
+				cost := peer.GetRequestCost(GetBlockHeadersMsg, amount)
+				peer.fcServer.QueueRequest(reqID, cost)
+				return func() { peer.RequestHeadersByHash(reqID, cost, origin, amount, skip, reverse) }
+			},
+		}
+		_, ok := <-pm.reqDist.queue(rq)
+		if !ok {
+			return ErrNoPeers
+		}
+		return nil
+	}
+	requestHeadersByNumber := func(origin uint64, amount int, skip int, reverse bool) error {
+		reqID := getNextReqID()
+		rq := &distReq{
+			getCost: func(dp distPeer) uint64 {
+				peer := dp.(*peer)
+				return peer.GetRequestCost(GetBlockHeadersMsg, amount)
+			},
+			canSend: func(dp distPeer) bool {
+				return dp.(*peer) == p
+			},
+			request: func(dp distPeer) func() {
+				peer := dp.(*peer)
+				cost := peer.GetRequestCost(GetBlockHeadersMsg, amount)
+				peer.fcServer.QueueRequest(reqID, cost)
+				return func() { peer.RequestHeadersByNumber(reqID, cost, origin, amount, skip, reverse) }
+			},
+		}
+		_, ok := <-pm.reqDist.queue(rq)
+		if !ok {
+			return ErrNoPeers
+		}
+		return nil
+	}
+
+	pm.downloader.RegisterPeer(p.id, ethVersion, p.HeadAndTd, requestHeadersByHash, requestHeadersByNumber, nil, nil, nil)
+}
+
+func (d *downloaderPeerNotify) unregisterPeer(p *peer) {
+	pm := (*ProtocolManager)(d)
+	pm.downloader.UnregisterPeer(p.id)
+}

les/handler_test.go

@@ -25,6 +25,7 @@ import (
 	"github.com/ethereum/go-ethereum/core/types"
 	"github.com/ethereum/go-ethereum/crypto"
 	"github.com/ethereum/go-ethereum/eth/downloader"
+	"github.com/ethereum/go-ethereum/ethdb"
 	"github.com/ethereum/go-ethereum/p2p"
 	"github.com/ethereum/go-ethereum/rlp"
 	"github.com/ethereum/go-ethereum/trie"
@@ -42,7 +43,8 @@ func expectResponse(r p2p.MsgReader, msgcode, reqID, bv uint64, data interface{}
 func TestGetBlockHeadersLes1(t *testing.T) { testGetBlockHeaders(t, 1) }
 
 func testGetBlockHeaders(t *testing.T, protocol int) {
-	pm, _, _ := newTestProtocolManagerMust(t, false, downloader.MaxHashFetch+15, nil)
+	db, _ := ethdb.NewMemDatabase()
+	pm := newTestProtocolManagerMust(t, false, downloader.MaxHashFetch+15, nil, nil, nil, db)
 	bc := pm.blockchain.(*core.BlockChain)
 	peer, _ := newTestPeer(t, "peer", protocol, pm, true)
 	defer peer.close()
@@ -170,7 +172,8 @@ func testGetBlockHeaders(t *testing.T, protocol int) {
 func TestGetBlockBodiesLes1(t *testing.T) { testGetBlockBodies(t, 1) }
 
 func testGetBlockBodies(t *testing.T, protocol int) {
-	pm, _, _ := newTestProtocolManagerMust(t, false, downloader.MaxBlockFetch+15, nil)
+	db, _ := ethdb.NewMemDatabase()
+	pm := newTestProtocolManagerMust(t, false, downloader.MaxBlockFetch+15, nil, nil, nil, db)
 	bc := pm.blockchain.(*core.BlockChain)
 	peer, _ := newTestPeer(t, "peer", protocol, pm, true)
 	defer peer.close()
@@ -246,7 +249,8 @@ func TestGetCodeLes1(t *testing.T) { testGetCode(t, 1) }
 
 func testGetCode(t *testing.T, protocol int) {
 	// Assemble the test environment
-	pm, _, _ := newTestProtocolManagerMust(t, false, 4, testChainGen)
+	db, _ := ethdb.NewMemDatabase()
+	pm := newTestProtocolManagerMust(t, false, 4, testChainGen, nil, nil, db)
 	bc := pm.blockchain.(*core.BlockChain)
 	peer, _ := newTestPeer(t, "peer", protocol, pm, true)
 	defer peer.close()
@@ -278,7 +282,8 @@ func TestGetReceiptLes1(t *testing.T) { testGetReceipt(t, 1) }
 
 func testGetReceipt(t *testing.T, protocol int) {
 	// Assemble the test environment
-	pm, db, _ := newTestProtocolManagerMust(t, false, 4, testChainGen)
+	db, _ := ethdb.NewMemDatabase()
+	pm := newTestProtocolManagerMust(t, false, 4, testChainGen, nil, nil, db)
 	bc := pm.blockchain.(*core.BlockChain)
 	peer, _ := newTestPeer(t, "peer", protocol, pm, true)
 	defer peer.close()
@@ -304,7 +309,8 @@ func TestGetProofsLes1(t *testing.T) { testGetReceipt(t, 1) }
 
 func testGetProofs(t *testing.T, protocol int) {
 	// Assemble the test environment
-	pm, db, _ := newTestProtocolManagerMust(t, false, 4, testChainGen)
+	db, _ := ethdb.NewMemDatabase()
+	pm := newTestProtocolManagerMust(t, false, 4, testChainGen, nil, nil, db)
 	bc := pm.blockchain.(*core.BlockChain)
 	peer, _ := newTestPeer(t, "peer", protocol, pm, true)
 	defer peer.close()

les/helper_test.go

@@ -25,7 +25,6 @@ import (
 	"math/big"
 	"sync"
 	"testing"
-	"time"
 
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/consensus/ethash"
@@ -132,22 +131,22 @@ func testRCL() RequestCostList {
 // newTestProtocolManager creates a new protocol manager for testing purposes,
 // with the given number of blocks already known, and potential notification
 // channels for different events.
-func newTestProtocolManager(lightSync bool, blocks int, generator func(int, *core.BlockGen)) (*ProtocolManager, ethdb.Database, *LesOdr, error) {
+func newTestProtocolManager(lightSync bool, blocks int, generator func(int, *core.BlockGen), peers *peerSet, odr *LesOdr, db ethdb.Database) (*ProtocolManager, error) {
 	var (
 		evmux  = new(event.TypeMux)
 		engine = ethash.NewFaker()
-		db, _  = ethdb.NewMemDatabase()
 		gspec  = core.Genesis{
 			Config: params.TestChainConfig,
 			Alloc:  core.GenesisAlloc{testBankAddress: {Balance: testBankFunds}},
 		}
 		genesis = gspec.MustCommit(db)
-		odr     *LesOdr
 		chain       BlockChain
 	)
+	if peers == nil {
+		peers = newPeerSet()
+	}
 
 	if lightSync {
-		odr = NewLesOdr(db)
 		chain, _ = light.NewLightChain(odr, gspec.Config, engine, evmux)
 	} else {
 		blockchain, _ := core.NewBlockChain(db, gspec.Config, engine, evmux, vm.Config{})
@@ -158,9 +157,9 @@ func newTestProtocolManager(lightSync bool, blocks int, generator func(int, *cor
 		chain = blockchain
 	}
 
-	pm, err := NewProtocolManager(gspec.Config, lightSync, NetworkId, evmux, engine, chain, nil, db, odr, nil)
+	pm, err := NewProtocolManager(gspec.Config, lightSync, NetworkId, evmux, engine, peers, chain, nil, db, odr, nil, make(chan struct{}), new(sync.WaitGroup))
 	if err != nil {
-		return nil, nil, nil, err
+		return nil, err
 	}
 	if !lightSync {
 		srv := &LesServer{protocolManager: pm}
@@ -174,20 +173,20 @@ func newTestProtocolManager(lightSync bool, blocks int, generator func(int, *cor
 		srv.fcManager = flowcontrol.NewClientManager(50, 10, 1000000000)
 		srv.fcCostStats = newCostStats(nil)
 	}
-	pm.Start(nil)
+	pm.Start()
-	return pm, db, odr, nil
+	return pm, nil
 }
 
 // newTestProtocolManagerMust creates a new protocol manager for testing purposes,
 // with the given number of blocks already known, and potential notification
 // channels for different events. In case of an error, the constructor force-
 // fails the test.
-func newTestProtocolManagerMust(t *testing.T, lightSync bool, blocks int, generator func(int, *core.BlockGen)) (*ProtocolManager, ethdb.Database, *LesOdr) {
+func newTestProtocolManagerMust(t *testing.T, lightSync bool, blocks int, generator func(int, *core.BlockGen), peers *peerSet, odr *LesOdr, db ethdb.Database) *ProtocolManager {
-	pm, db, odr, err := newTestProtocolManager(lightSync, blocks, generator)
+	pm, err := newTestProtocolManager(lightSync, blocks, generator, peers, odr, db)
 	if err != nil {
 		t.Fatalf("Failed to create protocol manager: %v", err)
 	}
-	return pm, db, odr
+	return pm
 }
 
 // testTxPool is a fake, helper transaction pool for testing purposes
@@ -342,30 +341,3 @@ func (p *testPeer) handshake(t *testing.T, td *big.Int, head common.Hash, headNu
 func (p *testPeer) close() {
 	p.app.Close()
 }
-
-type testServerPool struct {
-	peer *peer
-	lock sync.RWMutex
-}
-
-func (p *testServerPool) setPeer(peer *peer) {
-	p.lock.Lock()
-	defer p.lock.Unlock()
-
-	p.peer = peer
-}
-
-func (p *testServerPool) getAllPeers() map[distPeer]struct{} {
-	p.lock.RLock()
-	defer p.lock.RUnlock()
-
-	m := make(map[distPeer]struct{})
-	if p.peer != nil {
-		m[p.peer] = struct{}{}
-	}
-	return m
-}
-
-func (p *testServerPool) adjustResponseTime(*poolEntry, time.Duration, bool) {
-
-}

les/odr.go

@@ -31,9 +31,10 @@ type LesOdr struct {
 	retriever *retrieveManager
 }
 
-func NewLesOdr(db ethdb.Database) *LesOdr {
+func NewLesOdr(db ethdb.Database, retriever *retrieveManager) *LesOdr {
 	return &LesOdr{
 		db:        db,
+		retriever: retriever,
 		stop:      make(chan struct{}),
 	}
 }

les/odr_test.go

@@ -158,15 +158,15 @@ func odrContractCall(ctx context.Context, db ethdb.Database, config *params.Chai
 
 func testOdr(t *testing.T, protocol int, expFail uint64, fn odrTestFn) {
 	// Assemble the test environment
-	pm, db, odr := newTestProtocolManagerMust(t, false, 4, testChainGen)
+	peers := newPeerSet()
-	lpm, ldb, odr := newTestProtocolManagerMust(t, true, 0, nil)
+	dist := newRequestDistributor(peers, make(chan struct{}))
+	rm := newRetrieveManager(peers, dist, nil)
+	db, _ := ethdb.NewMemDatabase()
+	ldb, _ := ethdb.NewMemDatabase()
+	odr := NewLesOdr(ldb, rm)
+	pm := newTestProtocolManagerMust(t, false, 4, testChainGen, nil, nil, db)
+	lpm := newTestProtocolManagerMust(t, true, 0, nil, peers, odr, ldb)
 	_, err1, lpeer, err2 := newTestPeerPair("peer", protocol, pm, lpm)
-	pool := &testServerPool{}
-	lpm.reqDist = newRequestDistributor(pool.getAllPeers, lpm.quitSync)
-	odr.reqDist = lpm.reqDist
-	pool.setPeer(lpeer)
-	odr.serverPool = pool
-	lpeer.hasBlock = func(common.Hash, uint64) bool { return true }
 	select {
 	case <-time.After(time.Millisecond * 100):
 	case err := <-err1:
@@ -198,13 +198,19 @@ func testOdr(t *testing.T, protocol int, expFail uint64, fn odrTestFn) {
 	}
 
 	// temporarily remove peer to test odr fails
-	pool.setPeer(nil)
 	// expect retrievals to fail (except genesis block) without a les peer
+	peers.Unregister(lpeer.id)
+	time.Sleep(time.Millisecond * 10) // ensure that all peerSetNotify callbacks are executed
 	test(expFail)
-	pool.setPeer(lpeer)
 	// expect all retrievals to pass
+	peers.Register(lpeer)
+	time.Sleep(time.Millisecond * 10) // ensure that all peerSetNotify callbacks are executed
+	lpeer.lock.Lock()
+	lpeer.hasBlock = func(common.Hash, uint64) bool { return true }
+	lpeer.lock.Unlock()
 	test(5)
-	pool.setPeer(nil)
 	// still expect all retrievals to pass, now data should be cached locally
+	peers.Unregister(lpeer.id)
+	time.Sleep(time.Millisecond * 10) // ensure that all peerSetNotify callbacks are executed
 	test(5)
 }

les/peer.go

@@ -166,9 +166,9 @@ func (p *peer) GetRequestCost(msgcode uint64, amount int) uint64 {
 // HasBlock checks if the peer has a given block
 func (p *peer) HasBlock(hash common.Hash, number uint64) bool {
 	p.lock.RLock()
-	hashBlock := p.hasBlock
+	hasBlock := p.hasBlock
 	p.lock.RUnlock()
-	return hashBlock != nil && hashBlock(hash, number)
+	return hasBlock != nil && hasBlock(hash, number)
 }
 
 // SendAnnounce announces the availability of a number of blocks through
@@ -433,11 +433,19 @@ func (p *peer) String() string {
 	)
 }
 
+// peerSetNotify is a callback interface to notify services about added or
+// removed peers
+type peerSetNotify interface {
+	registerPeer(*peer)
+	unregisterPeer(*peer)
+}
+
 // peerSet represents the collection of active peers currently participating in
 // the Light Ethereum sub-protocol.
 type peerSet struct {
 	peers      map[string]*peer
 	lock       sync.RWMutex
+	notifyList []peerSetNotify
 	closed     bool
 }
 
@@ -448,6 +456,17 @@ func newPeerSet() *peerSet {
 	}
 }
 
+// notify adds a service to be notified about added or removed peers
+func (ps *peerSet) notify(n peerSetNotify) {
+	ps.lock.Lock()
+	defer ps.lock.Unlock()
+
+	ps.notifyList = append(ps.notifyList, n)
+	for _, p := range ps.peers {
+		go n.registerPeer(p)
+	}
+}
+
 // Register injects a new peer into the working set, or returns an error if the
 // peer is already known.
 func (ps *peerSet) Register(p *peer) error {
@@ -462,11 +481,14 @@ func (ps *peerSet) Register(p *peer) error {
 	}
 	ps.peers[p.id] = p
 	p.sendQueue = newExecQueue(100)
+	for _, n := range ps.notifyList {
+		go n.registerPeer(p)
+	}
 	return nil
 }
 
 // Unregister removes a remote peer from the active set, disabling any further
-// actions to/from that particular entity.
+// actions to/from that particular entity. It also initiates disconnection at the networking layer.
 func (ps *peerSet) Unregister(id string) error {
 	ps.lock.Lock()
 	defer ps.lock.Unlock()
@@ -474,7 +496,11 @@ func (ps *peerSet) Unregister(id string) error {
 	if p, ok := ps.peers[id]; !ok {
 		return errNotRegistered
 	} else {
+		for _, n := range ps.notifyList {
+			go n.unregisterPeer(p)
+		}
 		p.sendQueue.quit()
+		p.Peer.Disconnect(p2p.DiscUselessPeer)
 	}
 	delete(ps.peers, id)
 	return nil

les/request_test.go

@@ -68,15 +68,16 @@ func tfCodeAccess(db ethdb.Database, bhash common.Hash, number uint64) light.Odr
 
 func testAccess(t *testing.T, protocol int, fn accessTestFn) {
 	// Assemble the test environment
-	pm, db, _ := newTestProtocolManagerMust(t, false, 4, testChainGen)
+	peers := newPeerSet()
-	lpm, ldb, odr := newTestProtocolManagerMust(t, true, 0, nil)
+	dist := newRequestDistributor(peers, make(chan struct{}))
+	rm := newRetrieveManager(peers, dist, nil)
+	db, _ := ethdb.NewMemDatabase()
+	ldb, _ := ethdb.NewMemDatabase()
+	odr := NewLesOdr(ldb, rm)
+
+	pm := newTestProtocolManagerMust(t, false, 4, testChainGen, nil, nil, db)
+	lpm := newTestProtocolManagerMust(t, true, 0, nil, peers, odr, ldb)
 	_, err1, lpeer, err2 := newTestPeerPair("peer", protocol, pm, lpm)
-	pool := &testServerPool{}
-	lpm.reqDist = newRequestDistributor(pool.getAllPeers, lpm.quitSync)
-	odr.reqDist = lpm.reqDist
-	pool.setPeer(lpeer)
-	odr.serverPool = pool
-	lpeer.hasBlock = func(common.Hash, uint64) bool { return true }
 	select {
 	case <-time.After(time.Millisecond * 100):
 	case err := <-err1:
@@ -108,10 +109,16 @@ func testAccess(t *testing.T, protocol int, fn accessTestFn) {
 	}
 
 	// temporarily remove peer to test odr fails
-	pool.setPeer(nil)
+	peers.Unregister(lpeer.id)
+	time.Sleep(time.Millisecond * 10) // ensure that all peerSetNotify callbacks are executed
 	// expect retrievals to fail (except genesis block) without a les peer
 	test(0)
-	pool.setPeer(lpeer)
+
+	peers.Register(lpeer)
+	time.Sleep(time.Millisecond * 10) // ensure that all peerSetNotify callbacks are executed
+	lpeer.lock.Lock()
+	lpeer.hasBlock = func(common.Hash, uint64) bool { return true }
+	lpeer.lock.Unlock()
 	// expect all retrievals to pass
 	test(5)
 }

les/retrieve.go

@@ -38,8 +38,8 @@ var (
 // matching replies by request ID and handles timeouts and resends if necessary.
 type retrieveManager struct {
 	dist       *requestDistributor
+	peers      *peerSet
 	serverPool peerSelector
-	removePeer peerDropFn
 
 	lock     sync.RWMutex
 	sentReqs map[uint64]*sentReq
@@ -48,9 +48,6 @@ type retrieveManager struct {
 // validatorFunc is a function that processes a reply message
 type validatorFunc func(distPeer, *Msg) error
 
-// peerDropFn is a callback type for dropping a peer detected as malicious.
-type peerDropFn func(id string)
-
 // peerSelector receives feedback info about response times and timeouts
 type peerSelector interface {
 	adjustResponseTime(*poolEntry, time.Duration, bool)
@@ -69,16 +66,12 @@ type sentReq struct {
 	sentCnt, softTimeoutCnt int
 }
 
-// reqPeerCallback is called after a request sent to a certain peer has either been
-// answered or timed out hard
-type reqPeerCallback func(p distPeer, respTime time.Duration, srto, hrto bool)
-
 // newRetrieveManager creates the retrieve manager
-func newRetrieveManager(dist *requestDistributor, serverPool peerSelector, removePeer peerDropFn) *retrieveManager {
+func newRetrieveManager(peers *peerSet, dist *requestDistributor, serverPool peerSelector) *retrieveManager {
 	return &retrieveManager{
+		peers:      peers,
 		dist:       dist,
 		serverPool: serverPool,
-		removePeer: removePeer,
 		sentReqs:   make(map[uint64]*sentReq),
 	}
 }
@@ -199,8 +192,8 @@ func (r *sentReq) tryRetrieve() {
 				}
 				if hrto {
 					pp.Log().Debug("Request timed out hard")
-					if r.rm.removePeer != nil {
+					if r.rm.peers != nil {
-						r.rm.removePeer(pp.id)
+						r.rm.peers.Unregister(pp.id)
 					}
 				}
 			}

les/server.go

@@ -32,6 +32,7 @@ import (
 	"github.com/ethereum/go-ethereum/light"
 	"github.com/ethereum/go-ethereum/log"
 	"github.com/ethereum/go-ethereum/p2p"
+	"github.com/ethereum/go-ethereum/p2p/discv5"
 	"github.com/ethereum/go-ethereum/rlp"
 	"github.com/ethereum/go-ethereum/trie"
 )
@@ -41,17 +42,24 @@ type LesServer struct {
 	fcManager       *flowcontrol.ClientManager // nil if our node is client only
 	fcCostStats     *requestCostStats
 	defParams       *flowcontrol.ServerParams
+	lesTopic        discv5.Topic
+	quitSync        chan struct{}
 	stopped         bool
 }
 
 func NewLesServer(eth *eth.Ethereum, config *eth.Config) (*LesServer, error) {
-	pm, err := NewProtocolManager(eth.BlockChain().Config(), false, config.NetworkId, eth.EventMux(), eth.Engine(), eth.BlockChain(), eth.TxPool(), eth.ChainDb(), nil, nil)
+	quitSync := make(chan struct{})
+	pm, err := NewProtocolManager(eth.BlockChain().Config(), false, config.NetworkId, eth.EventMux(), eth.Engine(), newPeerSet(), eth.BlockChain(), eth.TxPool(), eth.ChainDb(), nil, nil, quitSync, new(sync.WaitGroup))
 	if err != nil {
 		return nil, err
 	}
 	pm.blockLoop()
 
-	srv := &LesServer{protocolManager: pm}
+	srv := &LesServer{
+		protocolManager: pm,
+		quitSync:        quitSync,
+		lesTopic:        lesTopic(eth.BlockChain().Genesis().Hash()),
+	}
 	pm.server = srv
 
 	srv.defParams = &flowcontrol.ServerParams{
@@ -69,7 +77,14 @@ func (s *LesServer) Protocols() []p2p.Protocol {
 
 // Start starts the LES server
 func (s *LesServer) Start(srvr *p2p.Server) {
-	s.protocolManager.Start(srvr)
+	s.protocolManager.Start()
+	go func() {
+		logger := log.New("topic", s.lesTopic)
+		logger.Info("Starting topic registration")
+		defer logger.Info("Terminated topic registration")
+
+		srvr.DiscV5.RegisterTopic(s.lesTopic, s.quitSync)
+	}()
 }
 
 // Stop stops the LES service

les/serverpool.go

@@ -120,13 +120,11 @@ type serverPool struct {
 }
 
 // newServerPool creates a new serverPool instance
-func newServerPool(db ethdb.Database, dbPrefix []byte, topic discv5.Topic, quit chan struct{}, wg *sync.WaitGroup) *serverPool {
+func newServerPool(db ethdb.Database, quit chan struct{}, wg *sync.WaitGroup) *serverPool {
 	pool := &serverPool{
 		db:           db,
-		dbKey:        append(dbPrefix, []byte(topic)...),
 		quit:         quit,
 		wg:           wg,
-		topic:        topic,
 		entries:      make(map[discover.NodeID]*poolEntry),
 		timeout:      make(chan *poolEntry, 1),
 		adjustStats:  make(chan poolStatAdjust, 100),
@@ -137,15 +135,18 @@ func newServerPool(db ethdb.Database, dbPrefix []byte, topic discv5.Topic, quit
 	}
 	pool.knownQueue = newPoolEntryQueue(maxKnownEntries, pool.removeEntry)
 	pool.newQueue = newPoolEntryQueue(maxNewEntries, pool.removeEntry)
-	wg.Add(1)
-	pool.loadNodes()
-
-	go pool.eventLoop()
 	return pool
 }
 
-func (pool *serverPool) setServer(server *p2p.Server) {
+func (pool *serverPool) start(server *p2p.Server, topic discv5.Topic) {
 	pool.server = server
+	pool.topic = topic
+	pool.dbKey = append([]byte("serverPool/"), []byte(topic)...)
+	pool.wg.Add(1)
+	pool.loadNodes()
+
+	go pool.eventLoop()
+
 	pool.checkDial()
 	if pool.server.DiscV5 != nil {
 		pool.discSetPeriod = make(chan time.Duration, 1)

les/txrelay.go

@@ -39,26 +39,28 @@ type LesTxRelay struct {
 	reqDist *requestDistributor
 }
 
-func NewLesTxRelay() *LesTxRelay {
+func NewLesTxRelay(ps *peerSet, reqDist *requestDistributor) *LesTxRelay {
-	return &LesTxRelay{
+	r := &LesTxRelay{
 		txSent:    make(map[common.Hash]*ltrInfo),
 		txPending: make(map[common.Hash]struct{}),
+		ps:        ps,
+		reqDist:   reqDist,
 	}
+	ps.notify(r)
+	return r
 }
 
-func (self *LesTxRelay) addPeer(p *peer) {
+func (self *LesTxRelay) registerPeer(p *peer) {
 	self.lock.Lock()
 	defer self.lock.Unlock()
 
-	self.ps.Register(p)
 	self.peerList = self.ps.AllPeers()
 }
 
-func (self *LesTxRelay) removePeer(id string) {
+func (self *LesTxRelay) unregisterPeer(p *peer) {
 	self.lock.Lock()
 	defer self.lock.Unlock()
 
-	self.ps.Unregister(id)
 	self.peerList = self.ps.AllPeers()
 }