From b53936a706905d72711d935180df00df4fd5531c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C5=81ukasz=20Magiera?= <magik6k@gmail.com>
Date: Wed, 12 Sep 2018 23:27:15 +0200
Subject: [PATCH] p2p: Close on Listeners
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

License: MIT
Signed-off-by: Łukasz Magiera <magik6k@gmail.com>
---
 core/commands/p2p.go       | 32 +++-----------------------------
 p2p/listener.go            | 28 ++++++++++++++++++++--------
 p2p/local.go               | 11 ++---------
 p2p/remote.go              |  8 +-------
 test/sharness/t0180-p2p.sh |  6 ++++--
 5 files changed, 30 insertions(+), 55 deletions(-)

diff --git a/core/commands/p2p.go b/core/commands/p2p.go
index fe39d8e90..dda18de4c 100644
--- a/core/commands/p2p.go
+++ b/core/commands/p2p.go
@@ -340,36 +340,10 @@ var p2pCloseCmd = &cmds.Command{
 			return true
 		}
 
-		todo := make([]p2p.Listener, 0)
-		n.P2P.ListenersLocal.Lock()
-		for _, l := range n.P2P.ListenersLocal.Listeners {
-			if !match(l) {
-				continue
-			}
-			todo = append(todo, l)
-		}
-		n.P2P.ListenersLocal.Unlock()
-		n.P2P.ListenersP2P.Lock()
-		for _, l := range n.P2P.ListenersP2P.Listeners {
-			if !match(l) {
-				continue
-			}
-			todo = append(todo, l)
-		}
-		n.P2P.ListenersP2P.Unlock()
+		done := n.P2P.ListenersLocal.Close(match)
+		done += n.P2P.ListenersP2P.Close(match)
 
-		var errs []string
-		for _, l := range todo {
-			if err := l.Close(); err != nil {
-				errs = append(errs, err.Error())
-			}
-		}
-		if len(errs) != 0 {
-			res.SetError(fmt.Errorf("errors when closing streams: %s", strings.Join(errs, "; ")), cmdkit.ErrNormal)
-			return
-		}
-
-		res.SetOutput(len(todo))
+		res.SetOutput(done)
 	},
 	Type: int(0),
 	Marshalers: cmds.MarshalerMap{
diff --git a/p2p/listener.go b/p2p/listener.go
index b0fa6e714..e62ab236b 100644
--- a/p2p/listener.go
+++ b/p2p/listener.go
@@ -18,8 +18,8 @@ type Listener interface {
 
 	key() string
 
-	// Close closes the listener. Does not affect child streams
-	Close() error
+	// close closes the listener. Does not affect child streams
+	close()
 }
 
 // Listeners manages a group of Listener implementations,
@@ -73,12 +73,24 @@ func (r *Listeners) Register(l Listener) error {
 	return nil
 }
 
-// Deregister removes p2p listener from this registry
-func (r *Listeners) Deregister(k string) (bool, error) {
+func (r *Listeners) Close(matchFunc func(listener Listener) bool) int {
+	todo := make([]Listener, 0)
 	r.Lock()
-	defer r.Unlock()
+	for _, l := range r.Listeners {
+		if !matchFunc(l) {
+			continue
+		}
 
-	_, ok := r.Listeners[k]
-	delete(r.Listeners, k)
-	return ok, nil
+		if _, ok := r.Listeners[l.key()]; ok {
+			delete(r.Listeners, l.key())
+			todo = append(todo, l)
+		}
+	}
+	r.Unlock()
+
+	for _, l := range todo {
+		l.close()
+	}
+
+	return len(todo)
 }
diff --git a/p2p/local.go b/p2p/local.go
index ab42fc0c3..861ce5e88 100644
--- a/p2p/local.go
+++ b/p2p/local.go
@@ -98,15 +98,8 @@ func (l *localListener) setupStream(local manet.Conn) {
 	l.p2p.Streams.Register(stream)
 }
 
-func (l *localListener) Close() error {
-	ok, err := l.p2p.ListenersLocal.Deregister(l.laddr.String())
-	if err != nil {
-		return err
-	}
-	if ok {
-		return l.listener.Close()
-	}
-	return nil
+func (l *localListener) close() {
+	l.listener.Close()
 }
 
 func (l *localListener) Protocol() protocol.ID {
diff --git a/p2p/remote.go b/p2p/remote.go
index 63e4848c5..924cab4c5 100644
--- a/p2p/remote.go
+++ b/p2p/remote.go
@@ -85,13 +85,7 @@ func (l *remoteListener) TargetAddress() ma.Multiaddr {
 	return l.addr
 }
 
-func (l *remoteListener) Close() error {
-	_, err := l.p2p.ListenersP2P.Deregister(string(l.proto))
-	if err != nil {
-		return err
-	}
-	return nil
-}
+func (l *remoteListener) close() {}
 
 func (l *remoteListener) key() string {
 	return string(l.proto)
diff --git a/test/sharness/t0180-p2p.sh b/test/sharness/t0180-p2p.sh
index 9b1f50393..544cd1f6f 100755
--- a/test/sharness/t0180-p2p.sh
+++ b/test/sharness/t0180-p2p.sh
@@ -285,11 +285,13 @@ test_expect_success 'S->C Setup client side (custom proto)' '
 test_server_to_client
 
 test_expect_success 'C->S Close local listener' '
-  ipfsi 0 p2p close -p /p2p-test
-  ipfsi 0 p2p ls > actual &&
+  ipfsi 1 p2p close -p /p2p-test
+  ipfsi 1 p2p ls > actual &&
   test_must_be_empty actual
 '
 
+check_test_ports
+
 test_expect_success 'stop iptb' '
   iptb stop
 '