@@ -3,8 +3,9 @@ module DMS (runDMS) where
33
44import Chronos (Time, getTime, now)
55import Control.Applicative
6−import Control.Concurrent (forkIO, threadDelay)
6+import Control.Concurrent (forkFinally, threadDelay)
77import Control.Concurrent.Async (race_)
8+import Control.Concurrent.MVar
89import Control.Concurrent.STM
910import Control.Exception (throwIO, try, Exception)
1011import Control.Monad
@@ -16,6 +17,7 @@ import Data.Aeson (Value (..), encode, toJSON)
1617import Data.Coerce
1718import Data.Conduit (ConduitT, yield)
1819import Data.Int (Int64)
20+import Data.IORef
1921import Data.Maybe
2022import Data.Pool
2123import Data.Text (lines, pack, unpack, unwords)
@@ -31,6 +33,7 @@ import Servant
3133import System.Environment (lookupEnv)
3234import System.Exit
3335import System.Log.FastLogger
36+import System.Posix.Signals (sigTERM)
3437import System.Process.Typed
3538
3639
@@ -38,6 +41,8 @@ import API
3841import Common.API
3942import DMS.Args
4043import DMS.AWS
44+import DMS.Logger
45+import DMS.Unix
4146import Orphans ()
4247import TLS (createTLSOpts)
4348import Types
@@ -47,18 +52,21 @@ type PgPool = Pool Connection
4752type AppM = ReaderT AppState Handler
4853
4954data AppState = AppState
50− { pool :: PgPool
51− , logger :: TimedFastLogger
52− , eventSink :: TChan WSEvent
53− , projectName :: ProjectName
54− , baseDomain :: Domain
55− , namespace :: Namespace
56− , archiveRetention :: ArchiveRetention
57− , creationCommand :: Command
58− , updateCommand :: Command
59− , deletionCommand :: Command
60− , checkingCommand :: Command
61− , cleanupCommand :: Command }
55+ { pool :: PgPool
56+ , logger :: TimedFastLogger
57+ , eventSink :: TChan WSEvent
58+ , bgWorkersCounter :: IORef Int
59+ , gracefulShudownActivated :: IORef Bool
60+ , shutdownSem :: MVar ()
61+ , projectName :: ProjectName
62+ , baseDomain :: Domain
63+ , namespace :: Namespace
64+ , archiveRetention :: ArchiveRetention
65+ , creationCommand :: Command
66+ , updateCommand :: Command
67+ , deletionCommand :: Command
68+ , checkingCommand :: Command
69+ , cleanupCommand :: Command }
6270
6371data DeploymentException
6472 = DeploymentFailed Int
@@ -74,9 +82,15 @@ data DeploymentListType
7482
7583runDMS :: IO ()
7684runDMS = do
77− timeCache <- newTimeCache "%Y-%m-%d %T%z"
78− (logger', _) <- newTimedFastLogger timeCache (LogStdout defaultBufSize)
85+ logger' <- newLogger
7986 logInfo logger' "started"
87+ bgWorkersC <- newIORef 0
88+ gracefulShudownAct <- newIORef False
89+ shutdownS <- newEmptyMVar
90+ void $ do
91+ let
92+ termHandler = terminationHandler bgWorkersC gracefulShudownAct shutdownS
93+ installShutdownHandler logger' [sigTERM] termHandler
8094 opts <- parseArgs
8195 let
8296 a ?! e = a >>= maybe (die e) pure
@@ -99,6 +113,9 @@ runDMS = do
99113 pgPool
100114 logger'
101115 channel
116+ bgWorkersC
117+ gracefulShudownAct
118+ shutdownS
102119 projName
103120 domain
104121 ns
@@ -122,6 +139,7 @@ runDMS = do
122139 `race_` (runTLS tlsOpts warpOpts app')
123140 `race_` (run wsServerPort wsApp')
124141 `race_` (runStatusUpdater appSt)
142+ `race_` (runShutdownHandler appSt)
125143
126144initConnectionPool :: ByteString -> Int -> IO PgPool
127145initConnectionPool dbConnStr =
@@ -182,6 +200,7 @@ listH = do
182200
183201createH :: Deployment -> AppM CommandResponse
184202createH dep = do
203+ failIfGracefulShudownActivated
185204 t1 <- liftIO $ now
186205 st <- ask
187206 let
@@ -202,7 +221,7 @@ createH dep = do
202221 { errBody = validationError ["deployment already exists"] [] }
203222 Left (SqlError _ _ _ _ _) ->
204223 throwError err409 { errBody = appError "some database error" }
205− void . liftIO . forkIO $ do
224+ liftIO . runBgWorker st $ do
206225 (ec, out, err) <- createDeployment dep st
207226 t2 <- now
208227 let
@@ -287,6 +306,7 @@ selectDeployment p dName lType = do
287306
288307editH :: DeploymentName -> EnvPairs -> AppM CommandResponse
289308editH dName dEnvs = do
309+ failIfGracefulShudownActivated
290310 t1 <- liftIO $ now
291311 st <- ask
292312 let
@@ -304,7 +324,7 @@ editH dName dEnvs = do
304324 , "--tag", coerce dTag
305325 ] ++ concat [["--env", concatPair p] | p <- dEnvs]
306326 cmd = coerce $ updateCommand st
307− void . liftIO . forkIO $ do
327+ liftIO . runBgWorker st $ do
308328 void $ updateEditDeployment pgPool
309329 log $ "call " <> unwords (cmd : args)
310330 (ec, out, err) <- runCommand (unpack cmd) (unpack <$> args)
@@ -337,6 +357,7 @@ editH dName dEnvs = do
337357
338358deleteH :: DeploymentName -> AppM CommandResponse
339359deleteH dName = do
360+ failIfGracefulShudownActivated
340361 t1 <- liftIO $ now
341362 st <- ask
342363 let
@@ -350,7 +371,7 @@ deleteH dName = do
350371 cmd = coerce $ deletionCommand st
351372 arch = ArchivedFlag True
352373 dep <- selectDeployment pgPool dName AllDeployments
353− void . liftIO . forkIO $ do
374+ liftIO . runBgWorker st $ do
354375 log $ "call " <> unwords (cmd : args)
355376 (ec, out, err) <- runCommand (unpack cmd) (unpack <$> args)
356377 void $ archiveDeployment pgPool dName
@@ -374,6 +395,7 @@ archiveDeployment p dName = withResource p $ \conn -> do
374395
375396updateH :: DeploymentName -> DeploymentUpdate -> AppM CommandResponse
376397updateH dName DeploymentUpdate { newTag = dTag, newEnvs = nEnvs } = do
398+ failIfGracefulShudownActivated
377399 t1 <- liftIO $ now
378400 st <- ask
379401 let pgPool = pool st
@@ -394,7 +416,7 @@ updateH dName DeploymentUpdate { newTag = dTag, newEnvs = nEnvs } = do
394416 ] ++ concat [["--env", concatPair e] | e <- envPairs]
395417 cmd = coerce $ updateCommand st
396418 failIfImageNotFound dTag
397− void . liftIO . forkIO $ do
419+ liftIO . runBgWorker st $ do
398420 void $ updateDeployment pgPool dName dTag envPairs
399421 log $ "call " <> unwords (cmd : args)
400422 (ec, out, err) <- runCommand (unpack cmd) (unpack <$> args)
@@ -475,11 +497,12 @@ statusH dName = do
475497
476498cleanupH :: DeploymentName -> AppM CommandResponse
477499cleanupH dName = do
500+ failIfGracefulShudownActivated
478501 st <- ask
479− cleanupDeployment dName st
502+ liftIO . runBgWorker st $ cleanupDeployment dName st
480503 pure Success
481504
482−cleanupDeployment :: DeploymentName -> AppState -> AppM ()
505+cleanupDeployment :: DeploymentName -> AppState -> IO ()
483506cleanupDeployment dName st = do
484507 let
485508 log = logInfo (logger st)
@@ -490,15 +513,14 @@ cleanupDeployment dName st = do
490513 , "--name", coerce dName
491514 ]
492515 cmd = coerce $ cleanupCommand st
493− void . liftIO . forkIO $ do
494− log $ "call " <> unwords (cmd : args)
495− (ec, out, err) <- runCommand (unpack cmd) (unpack <$> args)
496− print out >> print err
497− void $ deleteDeploymentLogs pgPool dName
498− void $ deleteDeployment pgPool dName
499− log $ "deployment destroyed, name: " <> coerce dName
500− sendReloadEvent st
501− handleExitCode ec
516+ log $ "call " <> unwords (cmd : args)
517+ (ec, out, err) <- runCommand (unpack cmd) (unpack <$> args)
518+ print out >> print err
519+ void $ deleteDeploymentLogs pgPool dName
520+ void $ deleteDeployment pgPool dName
521+ log $ "deployment destroyed, name: " <> coerce dName
522+ sendReloadEvent st
523+ handleExitCode ec
502524
503525deleteDeploymentLogs :: PgPool -> DeploymentName -> IO Int64
504526deleteDeploymentLogs p n = withResource p $ \conn -> execute
@@ -513,6 +535,7 @@ deleteDeployment p n = withResource p $ \conn ->
513535
514536cleanArchiveH :: AppM CommandResponse
515537cleanArchiveH = do
538+ failIfGracefulShudownActivated
516539 st <- ask
517540 let
518541 pgPool = pool st
@@ -522,17 +545,19 @@ cleanArchiveH = do
522545 \WHERE archived = 't' AND archived_at + interval '?' second < now()"
523546 retrieved :: [Only DeploymentName] <- liftIO $
524547 withResource pgPool $ \conn -> query conn q (Only archRetention)
525− void $ for retrieved $ \(Only dName) -> cleanupDeployment dName st
548+ void . liftIO $ for retrieved $ \(Only dName) ->
549+ runBgWorker st $ cleanupDeployment dName st
526550 pure Success
527551
528552restoreH :: DeploymentName -> AppM CommandResponse
529553restoreH dName = do
554+ failIfGracefulShudownActivated
530555 t1 <- liftIO $ now
531556 st <- ask
532557 let pgPool = pool st
533558 dep <- selectDeployment pgPool dName ArchivedOnlyDeployments
534559 failIfImageNotFound $ tag dep
535− void . liftIO . forkIO $ do
560+ liftIO . runBgWorker st $ do
536561 (ec, out, err) <- createDeployment dep st
537562 let
538563 q =
@@ -693,15 +718,33 @@ newStatus (ExitFailure _) UpdatePending ts | ts > updatePeriod = Failure
693718newStatus (ExitFailure _) DeletePending _ = Failure
694719newStatus (ExitFailure _) oldStatus _ = oldStatus
695720
696−logInfo :: TimedFastLogger -> Text -> IO ()
697−logInfo l = logWithSeverity l "INFO"
698−
699−logWarning :: TimedFastLogger -> Text -> IO ()
700−logWarning l = logWithSeverity l "WARN"
701−
702−logWithSeverity :: ToLogStr msg => TimedFastLogger -> ByteString -> msg -> IO ()
703−logWithSeverity l severity msg = l $ \ft -> metadata ft <> message
721+failIfGracefulShudownActivated :: AppM ()
722+failIfGracefulShudownActivated = do
723+ gracefulShudownAct <- gracefulShudownActivated <$> ask
724+ gracefulShudown <- liftIO . readIORef $ gracefulShudownAct
725+ if gracefulShudown
726+ then throwError err405 { errBody = appError "graceful shutdown activated" }
727+ else pure ()
728+
729+terminationHandler :: IORef Int -> IORef Bool -> MVar () -> IO ()
730+terminationHandler bgWorkersC gracefulShudownAct shutdownS = do
731+ atomicWriteIORef gracefulShudownAct True
732+ c <- readIORef bgWorkersC
733+ if c == 0
734+ then putMVar shutdownS ()
735+ else pure ()
736+
737+runShutdownHandler :: AppState -> IO ()
738+runShutdownHandler = takeMVar . shutdownSem
739+
740+runBgWorker :: AppState -> IO () -> IO ()
741+runBgWorker state act = void $ forkFinally act' cleanup
704742 where
705− metadata :: ByteString -> LogStr
706− metadata ft = foldMap toLogStr ["[", ft, " ", severity, "] "]
707− message = toLogStr msg <> toLogStr ("\n" :: ByteString)
743+ act' =
744+ atomicModifyIORef' (bgWorkersCounter state) (\c -> (c + 1, c)) >> act
745+ cleanup _ = do
746+ c <- atomicModifyIORef' (bgWorkersCounter state) (\c -> (c - 1, c - 1))
747+ gracefulShudown <- readIORef (gracefulShudownActivated state)
748+ if gracefulShudown && c == 0
749+ then putMVar (shutdownSem state) ()
750+ else pure ()