From 3d336cf6c3f4a4a91a21fec8066218fea65f059e Mon Sep 17 00:00:00 2001
From: web3-developer <51288821+web3-developer@users.noreply.github.com>
Date: Fri, 20 Sep 2024 09:33:16 +0800
Subject: [PATCH] Implement graceful shutdown.

---
 fluffy/fluffy.nim                             | 36 ++++++++++++++++---
 .../beacon/beacon_light_client_manager.nim    |  6 ++--
 fluffy/portal_node.nim                        | 10 +++++-
 3 files changed, 43 insertions(+), 9 deletions(-)

diff --git a/fluffy/fluffy.nim b/fluffy/fluffy.nim
index cdb4933d4..25a0bd70e 100644
--- a/fluffy/fluffy.nim
+++ b/fluffy/fluffy.nim
@@ -41,7 +41,7 @@ func optionToOpt[T](o: Option[T]): Opt[T] =
   else:
     Opt.none(T)
 
-proc run(config: PortalConf) {.raises: [CatchableError].} =
+proc run(config: PortalConf): PortalNode {.raises: [CatchableError].} =
   setupLogging(config.logLevel, config.logStdout, none(OutFile))
 
   notice "Launching Fluffy", version = fullVersionStr, cmdParams = commandLineParams()
@@ -252,7 +252,7 @@ proc run(config: PortalConf) {.raises: [CatchableError].} =
 
     setupRpcServer(rpcWsServer)
 
-  runForever()
+  return node
 
 when isMainModule:
   {.pop.}
@@ -262,6 +262,32 @@ when isMainModule:
   )
   {.push raises: [].}
 
-  case config.cmd
-  of PortalCmd.noCommand:
-    run(config)
+  let node =
+    case config.cmd
+    of PortalCmd.noCommand:
+      run(config)
+
+  # Ctrl+C handling
+  proc controlCHandler() {.noconv.} =
+    when defined(windows):
+      # workaround for https://github.com/nim-lang/Nim/issues/4057
+      try:
+        setupForeignThreadGc()
+      except Exception as exc:
+        raiseAssert exc.msg # shouldn't happen
+
+    notice "Shutting down after having received SIGINT"
+    node.state = PortalNodeState.Stopping
+
+  try:
+    setControlCHook(controlCHandler)
+  except Exception as exc: # TODO Exception
+    warn "Cannot set ctrl-c handler", msg = exc.msg
+
+  while node.state == PortalNodeState.Running:
+    try:
+      poll()
+    except CatchableError as e:
+      warn "Exception in poll()", exc = e.name, err = e.msg
+
+  waitFor node.stop()
diff --git a/fluffy/network/beacon/beacon_light_client_manager.nim b/fluffy/network/beacon/beacon_light_client_manager.nim
index b05063f99..7a0a8c99f 100644
--- a/fluffy/network/beacon/beacon_light_client_manager.nim
+++ b/fluffy/network/beacon/beacon_light_client_manager.nim
@@ -50,7 +50,7 @@ type
   GetBoolCallback* = proc(): bool {.gcsafe, raises: [].}
   GetSlotCallback* = proc(): Slot {.gcsafe, raises: [].}
 
-  LightClientManager* = object
+  LightClientManager* = ref object
     network: BeaconNetwork
     rng: ref HmacDrbgContext
     getTrustedBlockRoot: GetTrustedBlockRootCallback
@@ -315,12 +315,12 @@ proc loop(self: LightClientManager) {.async: (raises: [CancelledError]).} =
         didLatestSyncTaskProgress = didProgress,
       )
 
-proc start*(self: var LightClientManager) =
+proc start*(self: LightClientManager) =
   ## Start light client manager's loop.
   doAssert self.loopFuture == nil
   self.loopFuture = self.loop()
 
-proc stop*(self: var LightClientManager) {.async: (raises: []).} =
+proc stop*(self: LightClientManager) {.async: (raises: []).} =
   ## Stop light client manager's loop.
   if not self.loopFuture.isNil():
     await noCancel(self.loopFuture.cancelAndWait())
diff --git a/fluffy/portal_node.nim b/fluffy/portal_node.nim
index edbb8ab49..64cebb3c9 100644
--- a/fluffy/portal_node.nim
+++ b/fluffy/portal_node.nim
@@ -24,6 +24,11 @@ export
   beacon_light_client, history_network, state_network, portal_protocol_config, forks
 
 type
+  PortalNodeState* = enum
+    Starting
+    Running
+    Stopping
+
   PortalNodeConfig* = object
     accumulatorFile*: Opt[string]
     disableStateRootValidation*: bool
@@ -33,6 +38,7 @@ type
     storageCapacity*: uint64
 
   PortalNode* = ref object
+    state*: PortalNodeState
     discovery: protocol.Protocol
     contentDB: ContentDB
     streamManager: StreamManager
@@ -214,6 +220,8 @@ proc start*(n: PortalNode) =
 
   n.statusLogLoop = statusLogLoop(n)
 
+  n.state = PortalNodeState.Running
+
 proc stop*(n: PortalNode) {.async: (raises: []).} =
   debug "Stopping Portal node"
 
@@ -233,8 +241,8 @@ proc stop*(n: PortalNode) {.async: (raises: []).} =
     futures.add(n.statusLogLoop.cancelAndWait())
 
   futures.add(n.discovery.closeWait())
+  n.contentDB.close()
 
   await noCancel(allFutures(futures))
 
-  n.contentDB.close()
   n.statusLogLoop = nil