From 7907d67797e9745600caa02ab8032705b9f8a42c Mon Sep 17 00:00:00 2001
From: Dmitriy Ryajov <dryajov@gmail.com>
Date: Mon, 5 Dec 2022 09:00:02 -0600
Subject: [PATCH] don't crash on shutdown (#323)

* don't crash on shutdown

* cleanup
---
 codex.nim | 40 ++++++++++++++++++++++++++++++++++++----
 1 file changed, 36 insertions(+), 4 deletions(-)

diff --git a/codex.nim b/codex.nim
index 528e48d2..dedc6344 100644
--- a/codex.nim
+++ b/codex.nim
@@ -19,6 +19,7 @@ import ./codex/utils/keyutils
 export codex, conf, libp2p, chronos, chronicles
 
 when isMainModule:
+  import std/sequtils
   import std/os
   import pkg/confutils/defs
   import ./codex/utils/fileutils
@@ -29,6 +30,12 @@ when isMainModule:
   when defined(posix):
     import system/ansi_c
 
+  type
+    CodexStatus {.pure.} = enum
+      Stopped,
+      Stopping,
+      Running
+
   let config = CodexConf.load(
     version = codexFullVersion
   )
@@ -59,6 +66,10 @@ when isMainModule:
 
     trace "Repo dir initialized", dir = config.dataDir / "repo"
 
+    var
+      state: CodexStatus
+      pendingFuts: seq[Future[void]]
+
     let
       keyPath =
         if isAbsolute(string config.netPrivKeyFile):
@@ -77,7 +88,11 @@ when isMainModule:
           setupForeignThreadGc()
         except Exception as exc: raiseAssert exc.msg # shouldn't happen
       notice "Shutting down after having received SIGINT"
-      waitFor server.stop()
+
+      pendingFuts.add(server.stop())
+      state = CodexStatus.Stopping
+
+      notice "Stopping Codex"
 
     try:
       setControlCHook(controlCHandler)
@@ -88,12 +103,29 @@ when isMainModule:
     when defined(posix):
       proc SIGTERMHandler(signal: cint) {.noconv.} =
         notice "Shutting down after having received SIGTERM"
-        waitFor server.stop()
-        notice "Stopped Codex"
+
+        pendingFuts.add(server.stop())
+        state = CodexStatus.Stopping
+
+        notice "Stopping Codex"
 
       c_signal(ansi_c.SIGTERM, SIGTERMHandler)
 
-    waitFor server.start()
+    pendingFuts.add(server.start())
+
+    state = CodexStatus.Running
+    while state == CodexStatus.Running:
+      # poll chronos
+      chronos.poll()
+
+    # wait fot futures to finish
+    let res = waitFor allFinished(pendingFuts)
+    state = CodexStatus.Stopped
+
+    if res.anyIt( it.failed ):
+      error "Codex didn't shutdown correctly"
+      quit QuitFailure
+
     notice "Exited codex"
 
   of StartUpCommand.initNode: