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>
This commit is contained in:
IWASE Yusuke 2016-12-16 15:17:37 +09:00 committed by FUJITA Tomonori
parent daea96645a
commit 98a3b03792

View File

@ -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: