From 7f0882767c3cc9f460fe7691e7e8bc321a37a0c8 Mon Sep 17 00:00:00 2001
From: Roman Khimov <roman@nspcc.ru>
Date: Wed, 15 Jan 2020 18:25:58 +0300
Subject: [PATCH] network: remove useless Done() method from the peer

It's internal state of the peer that no one should care about.
---
 pkg/network/helper_test.go |  4 ----
 pkg/network/peer.go        |  1 -
 pkg/network/tcp_peer.go    | 46 +++++++++++++-------------------------
 3 files changed, 16 insertions(+), 35 deletions(-)

diff --git a/pkg/network/helper_test.go b/pkg/network/helper_test.go
index 1d7fe5693..b35290935 100644
--- a/pkg/network/helper_test.go
+++ b/pkg/network/helper_test.go
@@ -183,10 +183,6 @@ func (p *localPeer) WriteMsg(msg *Message) error {
 	p.messageHandler(p.t, msg)
 	return nil
 }
-func (p *localPeer) Done() chan error {
-	done := make(chan error)
-	return done
-}
 func (p *localPeer) Version() *payload.Version {
 	return p.version
 }
diff --git a/pkg/network/peer.go b/pkg/network/peer.go
index a2b55375a..a11f09de5 100644
--- a/pkg/network/peer.go
+++ b/pkg/network/peer.go
@@ -19,7 +19,6 @@ type Peer interface {
 	PeerAddr() net.Addr
 	Disconnect(error)
 	WriteMsg(msg *Message) error
-	Done() chan error
 	Version() *payload.Version
 	LastBlockIndex() uint32
 	UpdateLastBlockIndex(lbIndex uint32)
diff --git a/pkg/network/tcp_peer.go b/pkg/network/tcp_peer.go
index bbd22b51b..5be9c8622 100644
--- a/pkg/network/tcp_peer.go
+++ b/pkg/network/tcp_peer.go
@@ -38,9 +38,10 @@ type TCPPeer struct {
 	lastBlockIndex uint32
 
 	lock      sync.RWMutex
+	finale    sync.Once
 	handShake handShakeStage
 
-	done chan error
+	done chan struct{}
 
 	wg sync.WaitGroup
 
@@ -53,7 +54,7 @@ func NewTCPPeer(conn net.Conn, s *Server) *TCPPeer {
 	return &TCPPeer{
 		conn:   conn,
 		server: s,
-		done:   make(chan error, 1),
+		done:   make(chan struct{}),
 	}
 }
 
@@ -68,19 +69,14 @@ func (p *TCPPeer) WriteMsg(msg *Message) error {
 }
 
 func (p *TCPPeer) writeMsg(msg *Message) error {
-	select {
-	case err := <-p.done:
-		return err
-	default:
-		w := io.NewBufBinWriter()
-		if err := msg.Encode(w.BinWriter); err != nil {
-			return err
-		}
-
-		_, err := p.conn.Write(w.Bytes())
-
+	w := io.NewBufBinWriter()
+	if err := msg.Encode(w.BinWriter); err != nil {
 		return err
 	}
+
+	_, err := p.conn.Write(w.Bytes())
+
+	return err
 }
 
 // handleConn handles the read side of the connection, it should be started as
@@ -137,8 +133,8 @@ func (p *TCPPeer) StartProtocol() {
 	pingTimer := time.NewTimer(p.server.PingTimeout)
 	for {
 		select {
-		case err = <-p.Done():
-			// time to stop
+		case <-p.done:
+			return
 		case m := <-p.server.addrReq:
 			err = p.WriteMsg(m)
 		case <-timer.C:
@@ -273,23 +269,13 @@ func (p *TCPPeer) PeerAddr() net.Addr {
 	return tcpAddr
 }
 
-// Done implements the Peer interface and notifies
-// all other resources operating on it that this peer
-// is no longer running.
-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.server.unregister <- peerDrop{p, err}
-	p.conn.Close()
-	select {
-	case p.done <- err:
-		// one message to the queue
-	default:
-		// the other side may already be gone, it's OK
-	}
+	p.finale.Do(func() {
+		p.server.unregister <- peerDrop{p, err}
+		p.conn.Close()
+		close(p.done)
+	})
 }
 
 // Version implements the Peer interface.