Node network improvements (#45)

* small improvements.

* Fixed datarace + cleanup node and peer

* bumped version.

* removed race flag to pass build

pkg/network/tcp_peer.go

@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"net"
 	"os"
+	"sync"
 	"time"
 
 	"github.com/CityOfZion/neo-go/pkg/network/payload"
@@ -31,10 +32,14 @@ type TCPPeer struct {
 	// incoming message along with its peer.
 	handleProto protoHandleFunc
 
-	// Done is used to broadcast this peer has stopped running
+	// Done is used to broadcast that this peer has stopped running
 	// and should be removed as reference.
 	done chan struct{}
-	send chan *Message
+
+	// Every send to this channel will terminate the Peer.
+	discErr chan error
+	closed  chan struct{}
+	wg      sync.WaitGroup
 
 	logger log.Logger
 }
@@ -49,10 +54,11 @@ func NewTCPPeer(conn net.Conn, fun protoHandleFunc) *TCPPeer {
 		endpoint:    e,
 		conn:        conn,
 		done:        make(chan struct{}),
-		send:        make(chan *Message),
 		logger:      logger,
 		connectedAt: time.Now().UTC(),
 		handleProto: fun,
+		discErr:     make(chan error),
+		closed:      make(chan struct{}),
 	}
 }
 
@@ -68,67 +74,80 @@ func (p *TCPPeer) Endpoint() util.Endpoint {
 
 // Send implements the Peer interface.
 func (p *TCPPeer) Send(msg *Message) {
-	p.send <- msg
+	buf := new(bytes.Buffer)
+	if err := msg.encode(buf); err != nil {
+		p.discErr <- err
+		return
+	}
+	if _, err := p.conn.Write(buf.Bytes()); err != nil {
+		p.discErr <- err
+		return
+	}
 }
 
-// Done implemnets the Peer interface.
+// Done implemnets the Peer interface. It use is to
+// notify the Node that this peer is no longer available
+// for sending messages to.
 func (p *TCPPeer) Done() chan struct{} {
 	return p.done
 }
 
-func (p *TCPPeer) run() error {
-	errCh := make(chan error, 1)
+// Disconnect terminates the peer connection.
+func (p *TCPPeer) Disconnect(err error) {
+	select {
+	case p.discErr <- err:
+	case <-p.closed:
+	}
+}
 
-	go p.readLoop(errCh)
-	go p.writeLoop(errCh)
+func (p *TCPPeer) run() (err error) {
+	p.wg.Add(1)
+	go p.readLoop()
 
-	err := <-errCh
-	p.logger.Log("err", err)
-	p.cleanup()
+run:
+	for {
+		select {
+		case err = <-p.discErr:
+			break run
+		}
+	}
+
+	p.conn.Close()
+	close(p.closed)
+	// Close done instead of sending empty struct.
+	// It could happen that startProtocol in Node never happens
+	// on connection errors for example.
+	close(p.done)
+	p.wg.Wait()
 	return err
 }
 
-func (p *TCPPeer) readLoop(errCh chan error) {
+func (p *TCPPeer) readLoop() {
+	defer p.wg.Done()
 	for {
-		msg := &Message{}
-		if err := msg.decode(p.conn); err != nil {
-			errCh <- err
-			break
+		select {
+		case <-p.closed:
+			return
+		default:
+			msg := &Message{}
+			if err := msg.decode(p.conn); err != nil {
+				p.discErr <- err
+				return
+			}
+			p.handleMessage(msg)
 		}
-		p.handleMessage(msg)
 	}
 }
 
-func (p *TCPPeer) writeLoop(errCh chan error) {
-	buf := new(bytes.Buffer)
-
-	for {
-		msg := <-p.send
-		if err := msg.encode(buf); err != nil {
-			errCh <- err
-			break
-		}
-		if _, err := p.conn.Write(buf.Bytes()); err != nil {
-			errCh <- err
-			break
-		}
-		buf.Reset()
-	}
-}
-
-func (p *TCPPeer) cleanup() {
-	p.conn.Close()
-	close(p.send)
-	p.done <- struct{}{}
-}
-
 func (p *TCPPeer) handleMessage(msg *Message) {
 	switch msg.CommandType() {
 	case CMDVersion:
 		version := msg.Payload.(*payload.Version)
 		p.version = version
-		p.handleProto(msg, p)
+		fallthrough
 	default:
-		p.handleProto(msg, p)
+		if err := p.handleProto(msg, p); err != nil {
+			p.discErr <- err
+		}
 	}
 }