rpc_cli: Close peer connections when exiting prompt

Signed-off-by: IWASE Yusuke <iwase.yusuke0@gmail.com>
Signed-off-by: FUJITA Tomonori <fujita.tomonori@lab.ntt.co.jp>

1 files changed, 26 insertions, 7 deletions

diff --git a/ryu/cmd/rpc_cli.py b/ryu/cmd/rpc_cli.py
index 2519c8bb..2c1eb6b3 100755
--- a/ryu/cmd/rpc_cli.py
+++ b/ryu/cmd/rpc_cli.py
@@ -56,16 +56,18 @@ class Peer(object):
     def __init__(self, name, addr):
         self._name = name
         self._addr = addr
+        self.socket = None
         self.client = None
         try:
             self.connect()
-        except:
-            pass
+        except ConnectionError as e:
+            print('Exception when connecting to peer "%s": %s' % (name, e))
+            raise e
 
     def connect(self):
-        self.client = None
-        s = socket.create_connection(self._addr)
-        self.client = rpc.Client(s, notification_callback=self.notification)
+        self.socket = socket.create_connection(self._addr)
+        self.client = rpc.Client(self.socket,
+                                 notification_callback=self.notification)
 
     def try_to_connect(self, verbose=False):
         if self.client:
@@ -104,12 +106,25 @@ class Peer(object):
             print("connected.  retrying the request...")
             return g()
 
+    def close(self):
+        self.socket.close()
+
 
 peers = {}
 
 
 def add_peer(name, host, port):
-    peers[name] = Peer(name, (host, port))
+    try:
+        peer = Peer(name, (host, port))
+    except ConnectionError:
+        return
+
+    peers[name] = peer
+
+
+def close_peers():
+    for peer in peers.values():
+        peer.socket.close()
 
 
 class Cmd(cmd.Cmd):
@@ -174,7 +189,8 @@ class Cmd(cmd.Cmd):
     def complete_notify(self, text, line, begidx, endidx):
         return self._complete_peer(text, line, begidx, endidx)
 
-    def do_EOF(self, _line):
+    def do_EOF(self, _line=None):
+        close_peers()
         sys.exit(0)
 
     def emptyline(self):
@@ -205,6 +221,9 @@ class Cmd(cmd.Cmd):
         signal.signal(signal.SIGALRM, self._timeout)
         signal.alarm(1)
 
+    def postloop(self):
+        close_peers()
+
     def onecmd(self, string):
         self._in_onecmd = True
         try: