add Terminate command for Watcher and using it

coingaming · Sep 2, 2021 · b5a9b78 · b5a9b78
1 parent 476379d
commit b5a9b78
Show file tree

Hide file tree

Showing 4 changed files with 51 additions and 8 deletions.
diff --git a/src/LndClient/LndTest.hs b/src/LndClient/LndTest.hs
@@ -13,6 +13,8 @@ module LndClient.LndTest
     -- * TestEnv
     TestEnv,
     newTestEnv,
+    deleteTestEnv,
+    withTestEnv,
     spawnLinkChannelWatcher,
     spawnLinkInvoiceWatcher,
     spawnLinkSingleInvoiceWatcher,
@@ -147,6 +149,30 @@ newTestEnv lnd loc = do
         testSingleInvoiceWatcher = siw
       }
 
+deleteTestEnv ::
+  ( KatipContext m,
+    MonadUnliftIO m
+  ) =>
+  TestEnv ->
+  m ()
+deleteTestEnv env = do
+  Watcher.terminate $ testChannelWatcher env
+  Watcher.terminate $ testInvoiceWatcher env
+  Watcher.terminate $ testSingleInvoiceWatcher env
+
+withTestEnv ::
+  ( KatipContext m,
+    MonadUnliftIO m
+  ) =>
+  LndEnv ->
+  NodeLocation ->
+  (TestEnv -> m ()) ->
+  m ()
+withTestEnv lnd loc this = do
+  env <- newTestEnv lnd loc
+  this env
+  deleteTestEnv env
+
 class
   ( KatipContext m,
     MonadUnliftIO m,

diff --git a/src/LndClient/Watcher.hs b/src/LndClient/Watcher.hs
@@ -15,7 +15,7 @@ module LndClient.Watcher
     unWatch,
     unWatchUnit,
     dupLndTChan,
-    delete,
+    terminate,
   )
 where
 
@@ -40,6 +40,7 @@ data Watcher req res
 data Cmd req
   = Watch req
   | UnWatch req
+  | Terminate (Async ())
 
 data Event req res
   = EventCmd (Cmd req)
@@ -139,9 +140,10 @@ dupLndTChan = atomically . dupTChan . watcherLndChan
 
 --
 -- TODO : atomically cancel all linked processes
+-- in case of root process crash/cancel (graceful cleanup)
 --
-delete :: (MonadUnliftIO m) => Watcher req res -> m ()
-delete (Watcher _ _ proc) = liftIO $ cancel proc
+terminate :: (MonadUnliftIO m) => Watcher req res -> m ()
+terminate (Watcher _ _ proc) = liftIO $ cancel proc
 
 loop ::
   (Ord req, MonadUnliftIO m, KatipContext m) =>
@@ -205,6 +207,17 @@ applyCmd w = \case
       Just t -> do
         liftIO $ cancel t
         loop w {watcherStateTasks = Map.delete x ts}
+  Terminate x -> do
+    $(logTM) (newSev w InfoS) $
+      logStr
+        ( "Watcher - applying Cmd Terminate "
+            <> show (asyncThreadId x) ::
+            Text
+        )
+    liftIO $ do
+      mapM_ cancel $ Map.elems ts
+      cancel x
+    loop w {watcherStateTasks = mempty}
   where
     ts = watcherStateTasks w
 

diff --git a/test/LndClient/RPCSpec.hs b/test/LndClient/RPCSpec.hs
@@ -136,7 +136,7 @@ spec = do
         (AddInvoice.rHash inv)
         Invoice.OPEN
         chan
-    Watcher.delete w
+    Watcher.terminate w
     liftIO $ res `shouldSatisfy` isRight
   it "watchUnit" $ withEnv $ do
     bob <- getLndEnv Bob
@@ -161,7 +161,7 @@ spec = do
             }
     cp <- liftLndResult =<< openChannelSync alice openChannelRequest
     res <- receiveActiveChannel proxyOwner cp chan
-    Watcher.delete w
+    Watcher.terminate w
     liftIO $ res `shouldSatisfy` isRight
   it "unWatch" $ withEnv $ do
     lnd <- getLndEnv Bob
@@ -179,7 +179,7 @@ spec = do
     purgeChan chan
     void . liftLndResult =<< addInvoice lnd addInvoiceRequest
     res <- readTChanTimeout (MicroSecondsDelay 500000) chan
-    Watcher.delete w
+    Watcher.terminate w
     liftIO $ res `shouldBe` Nothing
   it "ensureHodlInvoice" $ withEnv $ do
     r <- newRPreimage
@@ -352,7 +352,7 @@ spec = do
     --  alice <- getLndEnv Alice
     --  void $ liftLndResult =<< sendPayment alice spr
     --  res <- readTChanTimeout (MicroSecondsDelay 500000) chan
-    --  Watcher.delete w
+    --  Watcher.terminate w
     --  liftIO $ res `shouldSatisfy` isJust
     subscribeInvoicesRequest =
       SubscribeInvoicesRequest (Just $ AddIndex 1) Nothing

diff --git a/test/LndClient/TestApp.hs b/test/LndClient/TestApp.hs
@@ -111,7 +111,11 @@ readEnv = do
 withEnv :: AppM IO () -> IO ()
 withEnv this = do
   env <- readEnv
-  runApp env $ setupZeroChannels proxyOwner >> this
+  runApp env $ do
+    setupZeroChannels proxyOwner
+    this
+    deleteTestEnv $ envBob env
+    deleteTestEnv $ envAlice env
   void . closeScribes $ envKatipLE env
 
 btcEnv :: BtcEnv