added peerCount.

This commit is contained in:
anthdm 2018-02-01 08:19:29 +01:00
parent 45ac0d237a
commit 0e22ae09bd
2 changed files with 65 additions and 20 deletions

View file

@ -5,7 +5,6 @@ import (
"log" "log"
"net" "net"
"os" "os"
"strconv"
"time" "time"
"github.com/anthdm/neo-go/pkg/network/payload" "github.com/anthdm/neo-go/pkg/network/payload"
@ -53,6 +52,8 @@ type Server struct {
relay bool relay bool
// TCP listener of the server // TCP listener of the server
listener net.Listener listener net.Listener
// channel for safely responding the number of current connected peers.
peerCountCh chan peerCount
} }
// NewServer returns a pointer to a new server. // NewServer returns a pointer to a new server.
@ -64,28 +65,26 @@ func NewServer(net NetMode) *Server {
} }
s := &Server{ s := &Server{
id: util.RandUint32(1111111, 9999999), id: util.RandUint32(1111111, 9999999),
userAgent: fmt.Sprintf("/NEO:%s/", version), userAgent: fmt.Sprintf("/NEO:%s/", version),
logger: logger, logger: logger,
peers: make(map[Peer]bool), peers: make(map[Peer]bool),
register: make(chan Peer), register: make(chan Peer),
unregister: make(chan Peer), unregister: make(chan Peer),
message: make(chan messageTuple), message: make(chan messageTuple),
relay: true, // currently relay is not handled. relay: true, // currently relay is not handled.
net: net, net: net,
quit: make(chan struct{}), quit: make(chan struct{}),
peerCountCh: make(chan peerCount),
} }
return s return s
} }
// Start run's the server. // Start run's the server.
func (s *Server) Start(port string, seeds []string) { // TODO: server should be initialized with a config.
p, err := strconv.Atoi(port[1:len(port)]) func (s *Server) Start(opts StartOpts) {
if err != nil { s.port = uint16(opts.TCP)
s.logger.Fatalf("could not convert port to integer: %s", err)
}
s.port = uint16(p)
fmt.Println(logo()) fmt.Println(logo())
fmt.Println(string(s.userAgent)) fmt.Println(string(s.userAgent))
@ -93,10 +92,14 @@ func (s *Server) Start(port string, seeds []string) {
s.logger.Printf("NET: %s - TCP: %d - RELAY: %v - ID: %d", s.logger.Printf("NET: %s - TCP: %d - RELAY: %v - ID: %d",
s.net, int(s.port), s.relay, s.id) s.net, int(s.port), s.relay, s.id)
go listenTCP(s, port) go listenTCP(s, opts.TCP)
if len(seeds) > 0 { if opts.RPC > 0 {
connectToSeeds(s, seeds) go listenHTTP(s, opts.RPC)
}
if len(opts.Seeds) > 0 {
connectToSeeds(s, opts.Seeds)
} }
s.loop() s.loop()
@ -137,6 +140,9 @@ func (s *Server) loop() {
s.logger.Printf("peer %s disconnected", peer.addr()) s.logger.Printf("peer %s disconnected", peer.addr())
} }
case t := <-s.peerCountCh:
t.count <- len(s.peers)
case <-s.quit: case <-s.quit:
s.shutdown() s.shutdown()
} }
@ -208,6 +214,7 @@ func (s *Server) peerAlreadyConnected(addr net.Addr) bool {
func (s *Server) sendLoop(peer Peer) { func (s *Server) sendLoop(peer Peer) {
// TODO: check if this peer is still connected. // TODO: check if this peer is still connected.
// dont keep asking (maxPeers and no new nodes)
for { for {
getaddrMsg := newMessage(s.net, cmdGetAddr, nil) getaddrMsg := newMessage(s.net, cmdGetAddr, nil)
peer.callGetaddr(getaddrMsg) peer.callGetaddr(getaddrMsg)
@ -216,6 +223,31 @@ func (s *Server) sendLoop(peer Peer) {
} }
} }
type peerCount struct {
count chan int
}
// peerCount returns the number of connected peers to this server.
func (s *Server) peerCount() int {
ch := peerCount{
count: make(chan int),
}
s.peerCountCh <- ch
return <-ch.count
}
// StartOpts holds the server configuration.
type StartOpts struct {
// tcp port
TCP int
// slice of peer addresses the server will connect to
Seeds []string
// JSON-RPC port. If 0 no RPC handler will be attached.
RPC int
}
func logo() string { func logo() string {
return ` return `
_ ____________ __________ _ ____________ __________

View file

@ -23,6 +23,19 @@ func TestHandleVersion(t *testing.T) {
t.Fatal("verack payload should be nil") t.Fatal("verack payload should be nil")
} }
} }
func TestPeerCount(t *testing.T) {
s := NewServer(ModeDevNet)
go s.loop()
lenPeers := 10
for i := 0; i < lenPeers; i++ {
s.register <- NewLocalPeer(s)
}
if have, want := s.peerCount(), lenPeers; want != have {
t.Fatalf("expected %d connected peers got %d", want, have)
}
}
func TestHandleAddrCmd(t *testing.T) { func TestHandleAddrCmd(t *testing.T) {
// todo // todo