@@ -20,14 +20,15 @@ module Katip.Scribes.ElasticSearch
2020
2121
2222-------------------------------------------------------------------------------
23−import Control.Applicative
2423import Control.Concurrent
25−import Control.Concurrent.STM
24+import Control.Concurrent.Async
25+import Control.Monad.STM
26+import Control.Concurrent.STM.TBMQueue
2627import Control.Exception.Enclosed
2728import Control.Monad
2829import Control.Retry
2930import Data.Aeson
30−import qualified Data.Text.Encoding as T
31+import qualified Data.Text.Encoding as T
3132import Data.Typeable
3233import Data.UUID
3334import Database.Bloodhound
@@ -76,19 +77,32 @@ mkEsScribe
7677 -> MappingName
7778 -> Severity
7879 -> Verbosity
79− -> IO Scribe
80+ -> IO (Scribe, IO ())
81+ -- ^ Returns a finalizer that will gracefully flush all remaining logs before shutting down workers
8082mkEsScribe cfg@EsScribeCfg {..} server ix mapping sev verb = do
81− q <- newTBQueueIO $ unEsQueueSize essQueueSize
83+ q <- newTBMQueueIO $ unEsQueueSize essQueueSize
8284 withManager essManagerSettings $ \mgr -> do
8385 let env = BHEnv { bhServer = server
8486 , bhManager = mgr
8587 }
86− replicateM_ (unEsPoolSize essPoolSize) $ forkIO $
88+ endSig <- newEmptyMVar
89+
90+ workers <- replicateM (unEsPoolSize essPoolSize) $ async $
8791 startWorker cfg env ix mapping q
8892
89− return $ Scribe $ \ i ->
90− when (_itemSeverity i >= sev) $
91− void $ atomically $ tryWriteTBQueue q (itemJson verb i)
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+
101+ let scribe = Scribe $ \ i ->
102+ when (_itemSeverity i >= sev) $
103+ void $ atomically $ tryWriteTBMQueue q (itemJson verb i)
104+ let finalizer = putMVar endSig () >> takeMVar endSig
105+ return (scribe, finalizer)
92106
93107
94108-------------------------------------------------------------------------------
@@ -139,21 +153,13 @@ startWorker
139153 -> BHEnv
140154 -> IndexName
141155 -> MappingName
142− -> TBQueue Value
156+ -> TBMQueue Value
143157 -> IO ()
144158startWorker EsScribeCfg {..} env ix mapping q = forever $ do
145− v <- atomically $ readTBQueue q
159+ v <- atomically $ readTBMQueue q
146160 sendLog v `catchAny` eat
147161 where
148162 sendLog v = void $ recoverAll essRetryPolicy $ do
149163 did <- mkDocId
150164 runBH env $ indexDocument ix mapping v did
151165 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