WIP

p2p/internal/dialsched.go

@@ -14,17 +14,13 @@
 // 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 p2p
+package p2pint
 
 import (
 	"container/heap"
-	"crypto/rand"
 	"fmt"
-	"net"
 	"time"
 
-	"github.com/ethereum/go-ethereum/logger"
-	"github.com/ethereum/go-ethereum/logger/glog"
 	"github.com/ethereum/go-ethereum/p2p/discover"
 )
 
@@ -71,13 +67,13 @@ const (
 	lookupInterval = 4 * time.Second
 )
 
-// dialstate schedules dials, discovery lookups and collects
+// DialState schedules dials, discovery lookups and collects
 // peer sugggestions from protocols.
 // It gets a chance to compute new tasks on every iteration
 // of the main loop in Server.run.
-type dialstate struct {
+type DialState struct {
-	ntab        discoverTable
+	peers       *PeerSet
-	protocols   []Protocol
+	ntab        DiscoverTable
 	maxDynDials int // per protocol
 
 	lookupRunning bool
@@ -85,19 +81,11 @@ type dialstate struct {
 	lookupBuf     []*discover.Node // current discovery lookup results
 	randomNodes   []*discover.Node // filled from Table
 
-	dialing map[discover.NodeID]connFlag
+	dialing map[discover.NodeID]Flag
 	static  map[discover.NodeID]*discover.Node
 	hist    *dialHistory
 }
 
-type discoverTable interface {
-	Self() *discover.Node
-	Close()
-	Bootstrap([]*discover.Node)
-	Lookup(target discover.NodeID) []*discover.Node
-	ReadRandomNodes([]*discover.Node) int
-}
-
 // the dial history remembers recent dials.
 type dialHistory []pastDial
 
@@ -107,74 +95,77 @@ type pastDial struct {
 	exp time.Time
 }
 
-type task interface {
+type DiscoverTable interface {
-	Do(*Server)
+	Self() *discover.Node
+	Close()
+	Bootstrap([]*discover.Node)
+	Lookup(target discover.NodeID) []*discover.Node
+	ReadRandomNodes([]*discover.Node) int
 }
 
-// A dialTask is generated for each node that is dialed.
+// A DialTask is generated for each node that is dialed.
-type dialTask struct {
+type DialTask struct {
-	flags connFlag
+	Flags Flag
-	dest  *discover.Node
+	Dest  *discover.Node
 }
 
-// discoverTask runs discovery table operations.
+// DiscoverTask runs discovery table operations.
-// Only one discoverTask is active at any time.
+// Only one DiscoverTask is active at any time.
 //
 // If bootstrap is true, the task runs Table.Bootstrap,
 // otherwise it performs a random lookup and leaves the
 // results in the task.
-type discoverTask struct {
+type DiscoverTask struct {
-	bootstrap bool
+	Bootstrap bool
-	results   []*discover.Node
+	Results   []*discover.Node
 }
 
-// A waitExpireTask is generated if there are no other tasks
+// A WaitExpireTask is generated if there are no other tasks
-// to keep the loop in Server.run ticking.
+// running and the dial history is non-empty. This ensures that
-type waitExpireTask struct {
+// the loop executing the tasks keeps ticking.
+type WaitExpireTask struct {
 	time.Duration
 }
 
-func newDialState(static []*discover.Node, ntab discoverTable, maxdyn int) *dialstate {
+func NewDialState(ps *PeerSet, ntab DiscoverTable, maxdyn int) *DialState {
-	s := &dialstate{
+	s := &DialState{
+		peers:       ps,
 		maxDynDials: maxdyn,
 		ntab:        ntab,
 		static:      make(map[discover.NodeID]*discover.Node),
-		dialing:     make(map[discover.NodeID]connFlag),
+		dialing:     make(map[discover.NodeID]Flag),
 		randomNodes: make([]*discover.Node, maxdyn/2),
 		hist:        new(dialHistory),
 	}
-	for _, n := range static {
-		s.static[n.ID] = n
-	}
 	return s
 }
 
-func (s *dialstate) addStatic(n *discover.Node) {
+func (s *DialState) AddStatic(n *discover.Node) {
 	s.static[n.ID] = n
 }
 
 // newTasks is called from the main loop to new tasks.
-func (s *dialstate) newTasks(nRunning int, peers map[discover.NodeID]*Peer, now time.Time) []task {
+func (s *DialState) NewTasks(nRunning int, now time.Time) []interface{} {
-	var newtasks []task
+	var newtasks []interface{}
-	addDial := func(flag connFlag, n *discover.Node) bool {
+	addDial := func(flag Flag, n *discover.Node) bool {
 		_, dialing := s.dialing[n.ID]
-		if dialing || peers[n.ID] != nil || s.hist.contains(n.ID) {
+		if dialing || s.peers.IsConnected(n.ID) || s.hist.contains(n.ID) {
 			return false
 		}
 		s.dialing[n.ID] = flag
-		newtasks = append(newtasks, &dialTask{flags: flag, dest: n})
+		newtasks = append(newtasks, &DialTask{Flags: flag, Dest: n})
 		return true
 	}
 
 	// Compute number of dynamic dials necessary at this point.
-	needDynDials := s.needDynDials(peers)
+	needDynDials := s.needDynDials()
 
 	// Expire the dial history on every invocation.
 	s.hist.expire(now)
 
 	// Create dials for static nodes if they are not connected.
 	for _, n := range s.static {
-		addDial(staticDialedConn, n)
+		addDial(StaticDialedConn, n)
 	}
 
 	// Use random nodes from the table for half of the necessary
@@ -183,7 +174,7 @@ func (s *dialstate) newTasks(nRunning int, peers map[discover.NodeID]*Peer, now
 	if randomCandidates > 0 && s.bootstrapped {
 		n := s.ntab.ReadRandomNodes(s.randomNodes)
 		for i := 0; i < randomCandidates && i < n; i++ {
-			if addDial(dynDialedConn, s.randomNodes[i]) {
+			if addDial(DynDialedConn, s.randomNodes[i]) {
 				needDynDials--
 			}
 		}
@@ -192,7 +183,7 @@ func (s *dialstate) newTasks(nRunning int, peers map[discover.NodeID]*Peer, now
 	// items from the result buffer.
 	i := 0
 	for ; i < len(s.lookupBuf) && needDynDials > 0; i++ {
-		if addDial(dynDialedConn, s.lookupBuf[i]) {
+		if addDial(DynDialedConn, s.lookupBuf[i]) {
 			needDynDials--
 		}
 	}
@@ -202,7 +193,7 @@ func (s *dialstate) newTasks(nRunning int, peers map[discover.NodeID]*Peer, now
 	// results.
 	if len(s.lookupBuf) < needDynDials && !s.lookupRunning {
 		s.lookupRunning = true
-		newtasks = append(newtasks, &discoverTask{bootstrap: !s.bootstrapped})
+		newtasks = append(newtasks, &DiscoverTask{Bootstrap: !s.bootstrapped})
 	}
 
 	// Launch a timer to wait for the next node to expire if all
@@ -210,81 +201,50 @@ func (s *dialstate) newTasks(nRunning int, peers map[discover.NodeID]*Peer, now
 	// This should prevent cases where the dialer logic is not ticked
 	// because there are no pending events.
 	if nRunning == 0 && len(newtasks) == 0 && s.hist.Len() > 0 {
-		t := &waitExpireTask{s.hist.min().exp.Sub(now)}
+		t := &WaitExpireTask{s.hist.min().exp.Sub(now)}
 		newtasks = append(newtasks, t)
 	}
 	return newtasks
 }
 
-// taskDone is called from the main loop when a task has finished.
+// TaskDone should be called when a task has finished.
-func (s *dialstate) taskDone(t task, now time.Time) {
+func (s *DialState) taskDone(t interface{}, now time.Time) {
 	switch t := t.(type) {
-	case *dialTask:
+	case *DialTask:
-		s.hist.add(t.dest.ID, now.Add(dialHistoryExpiration))
+		s.hist.add(t.Dest.ID, now.Add(dialHistoryExpiration))
-		delete(s.dialing, t.dest.ID)
+		delete(s.dialing, t.Dest.ID)
-	case *discoverTask:
+	case *DiscoverTask:
-		if t.bootstrap {
+		if t.Bootstrap {
 			s.bootstrapped = true
 		}
 		s.lookupRunning = false
-		s.lookupBuf = append(s.lookupBuf, t.results...)
+		s.lookupBuf = append(s.lookupBuf, t.Results...)
+	case *WaitExpireTask:
+		// nothing to do here
+	default:
+		panic("unknown task type")
 	}
 }
 
 // computes number of required dynamic dials
-func (s *dialstate) needDynDials(peers map[discover.NodeID]*Peer) int {
+func (s *DialState) needDynDials() int {
 	need := s.maxDynDials
-	for _, p := range peers {
+	for _, p := range s.peers.NumDynPeers() {
-		if p.rw.is(dynDialedConn) {
-			for _ = range p.running {
 		need--
 	}
-		}
-	}
 	for _, flag := range s.dialing {
-		if flag&dynDialedConn != 0 {
+		if flag&DynDialedConn != 0 {
-			need--
+v			need--
 		}
 	}
 	return need
 }
 
-func (t *dialTask) Do(srv *Server) {
+func (t *DialTask) String() string {
-	addr := &net.TCPAddr{IP: t.dest.IP, Port: int(t.dest.TCP)}
+	return fmt.Sprintf("%v %x %v:%d", t.Flags, t.Dest.ID[:8], t.Dest.IP, t.Dest.TCP)
-	glog.V(logger.Debug).Infof("dialing %v\n", t.dest)
-	fd, err := srv.Dialer.Dial("tcp", addr.String())
-	if err != nil {
-		glog.V(logger.Detail).Infof("dial error: %v", err)
-		return
-	}
-	mfd := newMeteredConn(fd, false)
-
-	srv.setupConn(mfd, t.flags, t.dest)
 }
 
-func (t *dialTask) String() string {
+func (t *DiscoverTask) String() (s string) {
-	return fmt.Sprintf("%v %x %v:%d", t.flags, t.dest.ID[:8], t.dest.IP, t.dest.TCP)
-}
-
-func (t *discoverTask) Do(srv *Server) {
-	if t.bootstrap {
-		srv.ntab.Bootstrap(srv.BootstrapNodes)
-		return
-	}
-	// newTasks generates a lookup task whenever dynamic dials are
-	// necessary. Lookups need to take some time, otherwise the
-	// event loop spins too fast.
-	next := srv.lastLookup.Add(lookupInterval)
-	if now := time.Now(); now.Before(next) {
-		time.Sleep(next.Sub(now))
-	}
-	srv.lastLookup = time.Now()
-	var target discover.NodeID
-	rand.Read(target[:])
-	t.results = srv.ntab.Lookup(target)
-}
-
-func (t *discoverTask) String() (s string) {
 	if t.bootstrap {
 		s = "discovery bootstrap"
 	} else {
@@ -296,10 +256,7 @@ func (t *discoverTask) String() (s string) {
 	return s
 }
 
-func (t waitExpireTask) Do(*Server) {
+func (t *WaitExpireTask) String() string {
-	time.Sleep(t.Duration)
-}
-func (t waitExpireTask) String() string {
 	return fmt.Sprintf("wait for dial hist expire (%v)", t.Duration)
 }
 

p2p/internal/peerset.go

@@ -0,0 +1,117 @@
+// Copyright 2015 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 p2pint
+
+import "github.com/ethereum/go-ethereum/p2p/discover"
+
+type Flag int
+
+const (
+	DynDialedConn Flag = 1 << iota
+	StaticDialedConn
+	InboundConn
+	TrustedConn
+)
+
+// PeerSet tracks connected peers and various indices
+// associcated with them.
+type PeerSet struct {
+	limit     int
+	counts    map[string]int
+	all       map[discover.NodeID]Peer
+	preferred map[discover.NodeID][]string
+	static    map[discover.NodeID]struct{}
+}
+
+type Peer interface {
+	ID() discover.NodeID
+	ActiveProtocols() []string
+	Flag() Flag
+}
+
+func NewPeerSet(limit int) *PeerSet {
+	return &PeerSet{
+		limit:     limit,
+		counts:    make(map[string]int),
+		all:       make(map[discover.NodeID]psPeer),
+		preferred: make(map[discover.NodeID][]string),
+		static:    make(map[discover.NodeID]struct{}),
+	}
+}
+
+func (ps *PeerSet) Add(p psPeer) {
+	for _, protoName := range p.activeProtocols() {
+		ps.counts[protoName]++
+	}
+	ps.all[p.ID()] = p
+}
+
+func (ps *PeerSet) Remove(p psPeer) {
+	for _, protoName := range p.activeProtocols() {
+		ps.setIdle(p, protoName)
+		ps.counts[protoName]--
+	}
+	delete(ps.all, p.ID())
+}
+
+func (ps *PeerSet) NumPeers() int {
+	return len(ps.all)
+}
+
+func (ps *PeerSet) IsAtCapacity() bool {
+	for _, numProtoPeers := range ps.counts {
+		if numProtoPeers < ps.limit {
+			return false
+		}
+	}
+	return true
+}
+
+func (ps *PeerSet) IsConnected(id discover.NodeID) bool {
+	_, ok := ps.all[id]
+	return ok
+}
+
+// disables auto-disconnect for a particular peer
+func (ps *PeerSet) SetPreferred(p psPeer, protoName string) {
+	pl := ps.preferred[p.ID()]
+	for _, name := range pl {
+		if name == protoName {
+			return
+		}
+	}
+	ps.preferred[p.ID()] = append(pl, protoName)
+}
+
+// opposite of setPreferred, removes from ps.preferred
+func (ps *PeerSet) SetIdle(p psPeer, protoName string) {
+	pl := ps.preferred[p.ID()]
+	pos := -1
+	for i, name := range pl {
+		if name == protoName {
+			pos = i
+			break
+		}
+	}
+	if pos >= 0 {
+		if len(pl) == 1 {
+			delete(ps.preferred, p.ID())
+		} else {
+			ps.preferred[p.ID()] = append(pl[:pos], pl[pos:]...)
+		}
+	}
+}

p2p/server.go

@@ -18,6 +18,7 @@ package p2p
 
 import (
 	"crypto/ecdsa"
+	"crypto/rand"
 	"errors"
 	"fmt"
 	"net"
@@ -27,6 +28,7 @@ import (
 	"github.com/ethereum/go-ethereum/logger"
 	"github.com/ethereum/go-ethereum/logger/glog"
 	"github.com/ethereum/go-ethereum/p2p/discover"
+	"github.com/ethereum/go-ethereum/p2p/internal"
 	"github.com/ethereum/go-ethereum/p2p/nat"
 )
 
@@ -147,21 +149,12 @@ type Server struct {
 
 type peerOpFunc func(map[discover.NodeID]*Peer)
 
-type connFlag int
-
-const (
-	dynDialedConn connFlag = 1 << iota
-	staticDialedConn
-	inboundConn
-	trustedConn
-)
-
 // conn wraps a network connection with information gathered
 // during the two handshakes.
 type conn struct {
 	fd net.Conn
 	transport
-	flags connFlag
+	flags p2pint.Flag
 	cont  chan error      // The run loop uses cont to signal errors to setupConn.
 	id    discover.NodeID // valid after the encryption handshake
 	caps  []Cap           // valid after the protocol handshake
@@ -428,92 +421,36 @@ func (srv *Server) run(dialstate dialer) {
 			for _, t := range pt[:start] {
 				t := t
 				glog.V(logger.Detail).Infoln("new task:", t)
-				go func() { t.Do(srv); taskdone <- t }()
+				go func() { srv.runDialTask(t); taskdone <- t }()
-			}
+	case *p2pint.DialTask:
-			copy(pt, pt[start:])
+		addr := &net.TCPAddr{IP: t.dest.IP, Port: int(t.dest.TCP)}
-			pendingTasks = pt[:len(pt)-start]
+		glog.V(logger.Debug).Infof("dialing %v\n", t.dest)
-		}
+		fd, err := srv.Dialer.Dial("tcp", addr.String())
-	}
-
-running:
-	for {
-		// Query the dialer for new tasks and launch them.
-		now := time.Now()
-		nt := dialstate.newTasks(len(pendingTasks)+len(tasks), peers, now)
-		scheduleTasks(nt)
-
-		select {
-		case <-srv.quit:
-			// The server was stopped. Run the cleanup logic.
-			glog.V(logger.Detail).Infoln("<-quit: spinning down")
-			break running
-		case n := <-srv.addstatic:
-			// This channel is used by AddPeer to add to the
-			// ephemeral static peer list. Add it to the dialer,
-			// it will keep the node connected.
-			glog.V(logger.Detail).Infoln("<-addstatic:", n)
-			dialstate.addStatic(n)
-		case op := <-srv.peerOp:
-			// This channel is used by Peers and PeerCount.
-			op(peers)
-			srv.peerOpDone <- struct{}{}
-		case t := <-taskdone:
-			// A task got done. Tell dialstate about it so it
-			// can update its state and remove it from the active
-			// tasks list.
-			glog.V(logger.Detail).Infoln("<-taskdone:", t)
-			dialstate.taskDone(t, now)
-			delTask(t)
-		case c := <-srv.posthandshake:
-			// A connection has passed the encryption handshake so
-			// the remote identity is known (but hasn't been verified yet).
-			if trusted[c.id] {
-				// Ensure that the trusted flag is set before checking against MaxPeers.
-				c.flags |= trustedConn
-			}
-			glog.V(logger.Detail).Infoln("<-posthandshake:", c)
-			// TODO: track in-progress inbound node IDs (pre-Peer) to avoid dialing them.
-			c.cont <- srv.encHandshakeChecks(peers, c)
-		case c := <-srv.addpeer:
-			// At this point the connection is past the protocol handshake.
-			// Its capabilities are known and the remote identity is verified.
-			glog.V(logger.Detail).Infoln("<-addpeer:", c)
-			err := srv.protoHandshakeChecks(peers, c)
 		if err != nil {
-				glog.V(logger.Detail).Infof("Not adding %v as peer: %v", c, err)
+			glog.V(logger.Detail).Infof("dial error: %v", err)
-			} else {
+			return
-				// The handshakes are done and it passed all checks.
-				p := newPeer(c, srv.Protocols)
-				peers[c.id] = p
-				go srv.runPeer(p)
-			}
-			// The dialer logic relies on the assumption that
-			// dial tasks complete after the peer has been added or
-			// discarded. Unblock the task last.
-			c.cont <- err
-		case p := <-srv.delpeer:
-			// A peer disconnected.
-			glog.V(logger.Detail).Infoln("<-delpeer:", p)
-			delete(peers, p.ID())
-		}
 		}
+		mfd := newMeteredConn(fd, false)
+		srv.setupConn(mfd, t.flags, t.dest)
 
-	// Terminate discovery. If there is a running lookup it will terminate soon.
+	case *p2pint.DiscoverTask:
-	if srv.ntab != nil {
+		// newTasks generates a lookup task whenever dynamic dials are
-		srv.ntab.Close()
+		// necessary. Lookups need to take some time, otherwise the
+		// event loop spins too fast.
+		next := srv.lastLookup.Add(lookupInterval)
+		if now := time.Now(); now.Before(next) {
+			time.Sleep(next.Sub(now))
 		}
-	// Disconnect all peers.
+		srv.lastLookup = time.Now()
-	for _, p := range peers {
+		var target discover.NodeID
-		p.Disconnect(DiscQuitting)
+		rand.Read(target[:])
-	}
+		t.results = srv.ntab.Lookup(target)
-	// Wait for peers to shut down. Pending connections and tasks are
+
-	// not handled here and will terminate soon-ish because srv.quit
+	case *p2pint.WaitExpireTask:
-	// is closed.
+		time.Sleep(t.Duration)
-	glog.V(logger.Detail).Infof("ignoring %d pending tasks at spindown", len(tasks))
+
-	for len(peers) > 0 {
+	default:
-		p := <-srv.delpeer
+		panic("unknown task type")
-		glog.V(logger.Detail).Infoln("<-delpeer (spindown):", p)
-		delete(peers, p.ID())
 	}
 }