Fixed some networking issues (#68)

* Faster persist timer

* fixed networking issues.

pkg/network/tcp_peer.go

@@ -18,31 +18,27 @@ type TCPPeer struct {
 	// The version of the peer.
 	version *payload.Version
 
-	done   chan error
-	closed chan struct{}
-	disc   chan error
+	done chan error
 
 	wg sync.WaitGroup
 }
 
-func NewTCPPeer(conn net.Conn, proto chan protoTuple) *TCPPeer {
+func NewTCPPeer(conn net.Conn) *TCPPeer {
 	return &TCPPeer{
 		conn:     conn,
-		done:     make(chan error),
-		closed:   make(chan struct{}),
-		disc:     make(chan error),
+		done:     make(chan error, 1),
 		endpoint: util.NewEndpoint(conn.RemoteAddr().String()),
 	}
 }
 
-// Send implements the Peer interface. This will encode the message
+// WriteMsg implements the Peer interface. This will write/encode the message
 // to the underlying connection.
-func (p *TCPPeer) Send(msg *Message) {
-	if err := msg.Encode(p.conn); err != nil {
-		select {
-		case p.disc <- err:
-		case <-p.closed:
-		}
+func (p *TCPPeer) WriteMsg(msg *Message) error {
+	select {
+	case err := <-p.done:
+		return err
+	default:
+		return msg.Encode(p.conn)
 	}
 }
 
@@ -58,71 +54,17 @@ func (p *TCPPeer) Done() chan error {
 	return p.done
 }
 
+// Disconnect will fill the peer's done channel with the given error.
+func (p *TCPPeer) Disconnect(err error) {
+	p.done <- err
+}
+
 // Version implements the Peer interface.
 func (p *TCPPeer) Version() *payload.Version {
 	return p.version
 }
 
-func (p *TCPPeer) readLoop(proto chan protoTuple, readErr chan error) {
-	defer p.wg.Done()
-	for {
-		select {
-		case <-p.closed:
-			return
-		default:
-			msg := &Message{}
-			if err := msg.Decode(p.conn); err != nil {
-				readErr <- err
-				return
-			}
-			p.handleMessage(msg, proto)
-		}
-	}
-}
-
-func (p *TCPPeer) handleMessage(msg *Message, proto chan protoTuple) {
-	switch payload := msg.Payload.(type) {
-	case *payload.Version:
-		p.version = payload
-	}
-	proto <- protoTuple{
-		msg:  msg,
-		peer: p,
-	}
-}
-
-func (p *TCPPeer) run(proto chan protoTuple) {
-	var (
-		readErr = make(chan error, 1)
-		err     error
-	)
-	p.wg.Add(1)
-	go p.readLoop(proto, readErr)
-
-run:
-	for {
-		select {
-		case err = <-p.disc:
-			break run
-		case err = <-readErr:
-			break run
-		}
-	}
-
-	// If the peer has not started the protocol with the server
-	// there will be noone reading from this channel.
-	select {
-	case p.done <- err:
-	default:
-	}
-
-	close(p.closed)
-	p.conn.Close()
-	p.wg.Wait()
-	return
-}
-
-// Disconnect implements the Peer interface.
-func (p *TCPPeer) Disconnect(reason error) {
-	p.disc <- reason
+// SetVersion implements the Peer interface.
+func (p *TCPPeer) SetVersion(v *payload.Version) {
+	p.version = v
 }