git.iko.soy/typeable/katip
on master help sign in
Add finalizer to gracefully shut down
4185f47Michael Xavier
katip-elasticsearch/katip-elasticsearch.cabal+1
compare with master
unifiedsplit
@@ -24,6 +24,7 @@ library
24 , uuid >= 1.3 && < 1.4
24 , uuid >= 1.3 && < 1.4
25 , aeson >=0.6 && <0.9
25 , aeson >=0.6 && <0.9
26 , stm >= 2.4.3 && < 2.5
26 , stm >= 2.4.3 && < 2.5
27+ , stm-chans >= 3.0.0.2 && < 3.1
27 , async >= 2.0.1.0 && < 2.1
28 , async >= 2.0.1.0 && < 2.1
28 , enclosed-exceptions >= 1.0.0 && < 1.1
29 , enclosed-exceptions >= 1.0.0 && < 1.1
29 , http-client
30 , http-client
katip-elasticsearch/src/Katip/Scribes/ElasticSearch.hs+25−19
@@ -20,14 +20,15 @@ module Katip.Scribes.ElasticSearch
20
20
21
21
22-------------------------------------------------------------------------------
22-------------------------------------------------------------------------------
23−import Control.Applicative
24import Control.Concurrent
23import Control.Concurrent
25−import Control.Concurrent.STM
24+import Control.Concurrent.Async
25+import Control.Monad.STM
26+import Control.Concurrent.STM.TBMQueue
26import Control.Exception.Enclosed
27import Control.Exception.Enclosed
27import Control.Monad
28import Control.Monad
28import Control.Retry
29import Control.Retry
29import Data.Aeson
30import Data.Aeson
30−import qualified Data.Text.Encoding as T
31+import qualified Data.Text.Encoding as T
31import Data.Typeable
32import Data.Typeable
32import Data.UUID
33import Data.UUID
33import Database.Bloodhound
34import Database.Bloodhound
@@ -76,19 +77,32 @@ mkEsScribe
76 -> MappingName
77 -> MappingName
77 -> Severity
78 -> Severity
78 -> Verbosity
79 -> Verbosity
79− -> IO Scribe
80+ -> IO (Scribe, IO ())
81+ -- ^ Returns a finalizer that will gracefully flush all remaining logs before shutting down workers
80mkEsScribe cfg@EsScribeCfg {..} server ix mapping sev verb = do
82mkEsScribe cfg@EsScribeCfg {..} server ix mapping sev verb = do
81− q <- newTBQueueIO $ unEsQueueSize essQueueSize
83+ q <- newTBMQueueIO $ unEsQueueSize essQueueSize
82 withManager essManagerSettings $ \mgr -> do
84 withManager essManagerSettings $ \mgr -> do
83 let env = BHEnv { bhServer = server
85 let env = BHEnv { bhServer = server
84 , bhManager = mgr
86 , bhManager = mgr
85 }
87 }
88+ endSig <- newEmptyMVar
89+
86− replicateM_ (unEsPoolSize essPoolSize) $ forkIO $
90+ workers <- replicateM (unEsPoolSize essPoolSize) $ async $
87 startWorker cfg env ix mapping q
91 startWorker cfg env ix mapping q
88
92
93+ _ <- async $ do
94+ takeMVar endSig
95+ atomically $ closeTBMQueue q
96+ -- wait for queue to empty
97+ atomically $ check =<< isEmptyTBMQueue q
98+ mapM_ cancel workers
99+ putMVar endSig ()
100+
89− return $ Scribe $ \ i ->
101+ let scribe = Scribe $ \ i ->
90− when (_itemSeverity i >= sev) $
102+ when (_itemSeverity i >= sev) $
91− void $ atomically $ tryWriteTBQueue q (itemJson verb i)
103+ void $ atomically $ tryWriteTBMQueue q (itemJson verb i)
104+ let finalizer = putMVar endSig () >> takeMVar endSig
105+ return (scribe, finalizer)
92
106
93
107
94-------------------------------------------------------------------------------
108-------------------------------------------------------------------------------
@@ -139,21 +153,13 @@ startWorker
139 -> BHEnv
153 -> BHEnv
140 -> IndexName
154 -> IndexName
141 -> MappingName
155 -> MappingName
142− -> TBQueue Value
156+ -> TBMQueue Value
143 -> IO ()
157 -> IO ()
144startWorker EsScribeCfg {..} env ix mapping q = forever $ do
158startWorker EsScribeCfg {..} env ix mapping q = forever $ do
145− v <- atomically $ readTBQueue q
159+ v <- atomically $ readTBMQueue q
146 sendLog v `catchAny` eat
160 sendLog v `catchAny` eat
147 where
161 where
148 sendLog v = void $ recoverAll essRetryPolicy $ do
162 sendLog v = void $ recoverAll essRetryPolicy $ do
149 did <- mkDocId
163 did <- mkDocId
150 runBH env $ indexDocument ix mapping v did
164 runBH env $ indexDocument ix mapping v did
151 eat _ = return ()
165 eat _ = return ()
152−
153−
154−-------------------------------------------------------------------------------
155−tryWriteTBQueue :: TBQueue a -> a -> STM Bool
156−tryWriteTBQueue q v = do
157− ok <- not <$> isFullTBQueue q
158− when ok $ writeTBQueue q v
159− return ok