obelisk

commit 5319d00f55699500d9ae9c97350bd150a559a4e2
parent 7cecf2f6906774f4ecaa234516b4f9f769c10113
Author: parazyd <parazyd@dyne.org>
Date:   Wed,  7 Apr 2021 22:17:45 +0200

Try exiting cleanly on ^C

Diffstat:
M
electrumobelisk/protocol.py
 | 
5
+++++
M
electrumobelisk/zeromq.py
 | 
1
+
M
obelisk.py
 | 
18
+++++++++++++++---

3 files changed, 21 insertions(+), 3 deletions(-)
diff --git a/electrumobelisk/protocol.py b/electrumobelisk/protocol.py
@@ -97,6 +97,11 @@ class ElectrumProtocol(asyncio.Protocol):  # pylint: disable=R0904,R0902
             "server.version": self.server_version,
         }
 
+    async def stop(self):
+        self.log.debug("ElectrumProtocol.stop()")
+        if self.bx:
+            await self.bx.stop()
+
     async def recv(self, reader, writer):
         recv_buf = bytearray()
         while True:
diff --git a/electrumobelisk/zeromq.py b/electrumobelisk/zeromq.py
@@ -211,6 +211,7 @@ class Client:
                                                      self._settings._loop)
 
     async def stop(self):
+        self.log.debug("zmq Client.stop()")
         self._query_socket.close()
         self._block_socket.close()
         return await self._request_collection.stop()
diff --git a/obelisk.py b/obelisk.py
@@ -26,6 +26,9 @@ from pkg_resources import resource_filename
 
 from electrumobelisk.protocol import ElectrumProtocol, VERSION
 
+# Used for destructor/cleanup
+PROTOCOL = None
+
 
 def logger_config(log, config):
     """Setup logging"""
@@ -100,9 +103,10 @@ async def run_electrum_server(config, chain):
     server_cfg["server_port"] = port
     server_cfg["using_tls"] = usetls
 
-    protocol = ElectrumProtocol(log, chain, endpoints, server_cfg)
+    global PROTOCOL
+    PROTOCOL = ElectrumProtocol(log, chain, endpoints, server_cfg)
 
-    server = await asyncio.start_server(protocol.recv, host, port)
+    server = await asyncio.start_server(PROTOCOL.recv, host, port)
     async with server:
         await server.serve_forever()
 
@@ -131,7 +135,15 @@ def main():
         log.error("chain is not 'mainnet' or 'testnet'")
         return 1
 
-    asyncio.run(run_electrum_server(config, chain))
+    try:
+        asyncio.run(run_electrum_server(config, chain))
+    except KeyboardInterrupt:
+        print("\r", end="")
+        log.debug("Caught KeyboardInterrupt, exiting...")
+        if PROTOCOL:
+            asyncio.run(PROTOCOL.stop())
+        return 0
+
     return 1