server: fix unexpected behaviour of waitForShutdown

With the current implementation, only the first call to `waitForShutdown` on a given
`ShutdownLatch` will return, while others will block (typically indefinitely). That's not
how one would expect a shutdown latch to work.

This isn't currently a concrete issue because we only wait once on each `ShutdownLatch`.
But in the context of #4154 we'll probably end up wanting to wait for shutdown from
multiple threads.

This adds a number of tests to verify the current behaviour, and adds a test for multiple
`waitForShutdown` calls that fails prior to the functional change.

PR-URL: https://github.com/hasura/graphql-engine-mono/pull/4162
GitOrigin-RevId: 9a108858d11390b847404f30bc7b93c06fc3f966

server/graphql-engine.cabal

@@ -834,6 +834,7 @@ test-suite graphql-engine-tests
   type: exitcode-stdio-1.0
   build-depends:
       aeson
+    , async
     , base
     , bytestring
     , case-insensitive
@@ -898,6 +899,7 @@ test-suite graphql-engine-tests
     Data.TimeSpec
     Data.TrieSpec
     Database.MSSQL.TransactionSpec
+    Hasura.AppSpec
     Hasura.Backends.DataWrapper.API.V0Spec
     Hasura.Backends.DataWrapper.API.V0.ColumnSpec
     Hasura.Backends.DataWrapper.API.V0.ExpressionSpec

server/src-lib/Hasura/App.hs

@@ -30,6 +30,7 @@ module Hasura.App
     runHGEServer,
     setCatalogStateTx,
     shutdownGracefully,
+    waitForShutdown,
 
     -- * Exported for testing
     mkHGEServer,
@@ -504,7 +505,7 @@ newShutdownLatch = fmap ShutdownLatch C.newEmptyMVar
 
 -- | Block the current thread, waiting on the latch.
 waitForShutdown :: ShutdownLatch -> IO ()
-waitForShutdown = C.takeMVar . unShutdownLatch
+waitForShutdown = C.readMVar . unShutdownLatch
 
 -- | Initiate a graceful shutdown of the server associated with the provided
 -- latch.

server/src-test/Hasura/AppSpec.hs

@@ -0,0 +1,65 @@
+{-# LANGUAGE NumericUnderscores #-}
+
+module Hasura.AppSpec (spec) where
+
+import Control.Concurrent (threadDelay)
+import Control.Concurrent.Async qualified as Async
+import Control.Exception (throwIO)
+import Hasura.App (newShutdownLatch, shutdownGracefully, waitForShutdown)
+import Hasura.Prelude
+import System.Timeout (timeout)
+import Test.Hspec
+
+spec :: Spec
+spec = do
+  describe "ShutdownLatch" shutdownLatchSpec
+
+shutdownLatchSpec :: Spec
+shutdownLatchSpec = do
+  it "waitForShutdown blocks before shutdown, not after" $ do
+    latch <- newShutdownLatch
+    timeout 10_000 (waitForShutdown latch)
+      `shouldReturn` Nothing
+    timeout 10_000 (shutdownGracefully latch)
+      `shouldReturn` Just ()
+    timeout 10_000 (waitForShutdown latch)
+      `shouldReturn` Just ()
+
+  it "allows multiple calls to shutdownGracefully" $ do
+    latch <- newShutdownLatch
+    timeout 10_000 (waitForShutdown latch)
+      `shouldReturn` Nothing
+    timeout 10_000 (shutdownGracefully latch)
+      `shouldReturn` Just ()
+    timeout 10_000 (shutdownGracefully latch)
+      `shouldReturn` Just ()
+    timeout 10_000 (waitForShutdown latch)
+      `shouldReturn` Just ()
+
+  it "allows shutting down a thread" $ do
+    latch <- newShutdownLatch
+    Async.withAsync (waitForShutdown latch >> return ("shut down" :: String)) $ \async -> do
+      threadDelay 10_000
+      pollThrow async
+        `shouldReturn` Nothing
+      shutdownGracefully latch
+      (timeout 1_000_000 $ Async.wait async)
+        `shouldReturn` Just "shut down"
+
+  it "allows multiple threads to wait for shutdown" $ do
+    latch <- newShutdownLatch
+    timeout 10_000 (waitForShutdown latch)
+      `shouldReturn` Nothing
+    shutdownGracefully latch
+    timeout 10_000 (waitForShutdown latch)
+      `shouldReturn` Just ()
+    timeout 10_000 (waitForShutdown latch)
+      `shouldReturn` Just ()
+
+pollThrow :: Async.Async a -> IO (Maybe a)
+pollThrow async = do
+  res <- Async.poll async
+  case res of
+    Just (Left e) -> throwIO e
+    Just (Right x) -> pure $ Just x
+    Nothing -> pure Nothing

server/src-test/Main.hs

@@ -25,6 +25,7 @@ import Hasura.App
     mkMSSQLSourceResolver,
     mkPgSourceResolver,
   )
+import Hasura.AppSpec qualified as AppSpec
 import Hasura.Backends.DataWrapper.API.V0Spec qualified as DataWrapper.API.V0Spec
 import Hasura.Backends.MSSQL.ErrorSpec qualified as MSSQLErrorSpec
 import Hasura.Backends.MySQL.DataLoader.ExecuteTests qualified as MySQLDataLoader
@@ -100,6 +101,7 @@ unitSpecs = do
   describe "Data.Parser.URLTemplate" URLTemplate.spec
   describe "Data.Time" TimeSpec.spec
   describe "Data.Trie" TrieSpec.spec
+  describe "Hasura.App" AppSpec.spec
   describe "Hasura.Backends.DataWrapper.API.V0" DataWrapper.API.V0Spec.spec
   describe "Hasura.Backends.MSSQL.ErrorSpec" MSSQLErrorSpec.spec
   describe "Hasura.Backends.MySQL.DataLoader.ExecuteTests" MySQLDataLoader.spec