From 900d53f9df79c9abd81e6b021c8057ab343adb5a Mon Sep 17 00:00:00 2001
From: Hanno Cornelius <68783915+jm-clius@users.noreply.github.com>
Date: Thu, 15 Apr 2021 10:18:14 +0200
Subject: [PATCH] Implement graceful shutdown for `wakunode2` and `chat2`
 (#490)

* Implement graceful shutdown

* Consistent use of QuitSuccess
---
 docs/tutorial/chat2.md     |  1 +
 examples/v2/chat2.nim      | 15 +++++++--------
 waku/v2/node/wakunode2.nim | 29 ++++++++++++++++++++++++++---
 3 files changed, 34 insertions(+), 11 deletions(-)

diff --git a/docs/tutorial/chat2.md b/docs/tutorial/chat2.md
index 37be9c636..5d43d5ab6 100644
--- a/docs/tutorial/chat2.md
+++ b/docs/tutorial/chat2.md
@@ -69,6 +69,7 @@ This will bypass the random peer selection process and connect to the specified
 | `/help` | displays available in-chat commands |
 | `/connect` | interactively connect to a new peer |
 | `/nick` | change nickname for current chat session |
+| `/exit` | exits the current chat session |
 
 ## `chat2` message protobuf format
 
diff --git a/examples/v2/chat2.nim b/examples/v2/chat2.nim
index 1306c75bd..7d36e6857 100644
--- a/examples/v2/chat2.nim
+++ b/examples/v2/chat2.nim
@@ -30,6 +30,7 @@ const Help = """
   help: Prints this help
   connect: dials a remote peer
   nick: change nickname for current chat session
+  exit: exits chat session
 """
 
 const
@@ -203,14 +204,11 @@ proc writeAndPrint(c: Chat) {.async.} =
       c.nick = await readNick(c.transp)
       echo "You are now known as " & c.nick
 
-#    elif line.startsWith("/exit"):
-#      if p.connected and p.conn.closed.not:
-#        await p.conn.close()
-#        p.connected = false
-#
-#      await p.switch.stop()
-#      echo "quitting..."
-#      quit(0)
+    elif line.startsWith("/exit"):
+     await c.node.stop()
+
+     echo "quitting..."
+     quit(QuitSuccess)
     else:
       # XXX connected state problematic
       if c.started:
@@ -368,6 +366,7 @@ proc processInput(rfd: AsyncFD, rng: ref BrHmacDrbgContext) {.async.} =
   node.subscribe(topic, handler)
 
   await chat.readWriteLoop()
+
   runForever()
   #await allFuturesThrowing(libp2pFuts)
 
diff --git a/waku/v2/node/wakunode2.nim b/waku/v2/node/wakunode2.nim
index 7da1cb847..f9e9af1f2 100644
--- a/waku/v2/node/wakunode2.nim
+++ b/waku/v2/node/wakunode2.nim
@@ -417,9 +417,6 @@ proc mountRelay*(node: WakuNode, topics: seq[string] = newSeq[string](), rlnRela
   node.wakuRelay = wakuRelay
   node.switch.mount(wakuRelay)
 
-  # Reonnect to previous relay peers
-  waitFor node.peerManager.reconnectPeers(WakuRelayCodec)
-
   info "mounting relay"
 
   node.subscribe(defaultTopic, none(TopicHandler))
@@ -441,6 +438,9 @@ proc mountRelay*(node: WakuNode, topics: seq[string] = newSeq[string](), rlnRela
     waitFor node.wakuRelay.start()
 
     info "relay mounted and started successfully"
+  
+  # Reconnect to previous relay peers
+  waitFor node.peerManager.reconnectPeers(WakuRelayCodec)
 
 ## Helpers
 proc dialPeer*(n: WakuNode, address: string) {.async.} =
@@ -496,6 +496,7 @@ proc connectToNodes*(n: WakuNode, nodes: seq[PeerInfo]) {.async.} =
 
 when isMainModule:
   import
+    system/ansi_c,
     confutils, json_rpc/rpcserver, metrics,
     ./config, 
     ./jsonrpc/[admin_api,
@@ -646,5 +647,27 @@ when isMainModule:
     if conf.metricsServer:
       startMetricsServer(conf.metricsServerAddress,
         Port(conf.metricsServerPort + conf.portsShift))
+  
+  # Setup graceful shutdown
+  
+  # Handle Ctrl-C SIGINT
+  proc handleCtrlC() {.noconv.} =
+    when defined(windows):
+      # workaround for https://github.com/nim-lang/Nim/issues/4057
+      setupForeignThreadGc()
+    info "Shutting down after receiving SIGINT"
+    waitFor node.stop()
+    quit(QuitSuccess)
+  
+  setControlCHook(handleCtrlC)
+
+  # Handle SIGTERM
+  when defined(posix):
+    proc handleSigterm(signal: cint) {.noconv.} =
+      info "Shutting down after receiving SIGTERM"
+      waitFor node.stop()
+      quit(QuitSuccess)
+    
+    c_signal(SIGTERM, handleSigterm)
 
   runForever()