@@ -1,5 +1,4 @@
11{-# LANGUAGE BangPatterns #-}
2−{-# LANGUAGE GADTs #-}
32{-# LANGUAGE GeneralizedNewtypeDeriving #-}
43{-# LANGUAGE OverloadedStrings #-}
54{-# LANGUAGE RankNTypes #-}
@@ -39,34 +38,33 @@ module Katip.Scribes.LogzIO.HTTPS
3938
4039-------------------------------------------------------------------------------
4140import Control.Applicative
42−import qualified Control.Concurrent.Async as Async
43−import qualified Control.Concurrent.STM as STM
44−import qualified Control.Concurrent.STM.TBMQueue as STM
45−import qualified Control.Error as E
46−import qualified Control.Exception.Safe as EX
47−import Control.Monad
48−import qualified Control.Retry as Retry
49−import qualified Data.Aeson as A
50−import qualified Data.ByteString.Builder as BB
51−import qualified Data.ByteString.Lazy as LBS
52−import qualified Data.ByteString.Lazy.Char8 as LBS8
53−import qualified Data.HashMap.Strict as HM
41+import qualified Control.Concurrent.Async as Async
42+import qualified Control.Concurrent.STM as STM
43+import qualified Control.Error as E
44+import qualified Control.Exception.Safe as EX
45+import qualified Control.Retry as Retry
46+import qualified Data.Aeson as A
47+import qualified Data.ByteString.Builder as BB
48+import qualified Data.ByteString.Lazy as LBS
49+import qualified Data.ByteString.Lazy.Char8 as LBS8
50+import qualified Data.HashMap.Strict as HM
5451import Data.Int
55−import qualified Data.Scientific as Scientific
56−import Data.Semigroup as Semigroup
57−import Data.String (IsString)
58−import qualified Data.Text as T
59−import qualified Data.Text.Encoding as TE
60−import qualified Data.Text.Lazy as TL
61−import qualified Data.Text.Lazy.Builder as TB
62−import qualified Data.Time as Time
63−import qualified Katip as K
64−import Katip.Core (LocJs (..))
65−import qualified Network.HTTP.Client as HTTP
66−import qualified Network.HTTP.Client.TLS as HTTPS
67−import qualified Network.HTTP.Types as HTypes
68−import qualified System.Posix.Types as POSIX
69−import qualified URI.ByteString as URIBS
52+import qualified Data.Scientific as Scientific
53+import Data.Semigroup as Semigroup
54+import Data.String (IsString)
55+import qualified Data.Text as T
56+import qualified Data.Text.Encoding as TE
57+import qualified Data.Text.Lazy as TL
58+import qualified Data.Text.Lazy.Builder as TB
59+import qualified Data.Time as Time
60+import qualified Katip as K
61+import Katip.Core (LocJs (..))
62+import qualified Network.HTTP.Client as HTTP
63+import qualified Network.HTTP.Client.TLS as HTTPS
64+import qualified Network.HTTP.Types as HTypes
65+import qualified System.Posix.Types as POSIX
66+import qualified URI.ByteString as URIBS
67+-------------------------------------------------------------------------------
7068-------------------------------------------------------------------------------
7169
7270
@@ -209,115 +207,132 @@ mkLogzIOScribe
209207 -> K.Verbosity
210208 -> IO K.Scribe
211209mkLogzIOScribe config permitItem verbosity = do
212− -- This size is actually somewhat arbitrary. We're just making sure
213− -- it isn't infinite so we don't consume the whole backlog right
214− -- away and take up lots of memory.. Katip above us does bounded
215− -- buffering and load-shedding when there's too much of a
216− -- backlog. Our writes are blocking to apply backpressure up to
217− -- katip. This value just holds onto values temporarily while we
218− -- concatenate them into a buffer.
219− ingestionQueue <- STM.newTBMQueueIO (itemBufferSize * 10)
220−
221− let newTimer = STM.registerDelay itemBufferTimeoutMicros
222− timerRef <- STM.newTVarIO =<< newTimer
210+ -- Accumulating buffer for each item
211+ bufferState <- STM.newTVarIO (mempty :: BulkBuffer)
212+ -- Single-slot TMVar when something is ready to write. Katip buffers
213+ -- for us upstream so we can simplify shutdown and flush guarantees
214+ -- by synchronously writing when we need to.
215+ readyMVar <- STM.newEmptyTMVarIO :: IO (STM.TMVar BulkBuffer)
216+
217+ -- We create a TVar that is set to ShutDown when its time to stop
218+ killSwitch <- STM.newTVarIO Continue
219+
220+ -- Create a mutable ref to whether the timer has expired. This lets
221+ -- us replace the timer when its time to start timing again. We
222+ -- reset the timer __after__ a write goes out, not as soon as it expires
223+ timerRef <- STM.newTVarIO =<< STM.registerDelay itemBufferTimeoutMicros
223224
224225 -- Set up a connection manager for requests
225226 mgr <- case scheme of
226227 HTTPS -> HTTPS.newTlsManager
227228 HTTP -> HTTP.newManager HTTP.defaultManagerSettings
228229
229− -- An STM transaction that will return true when writes are stopped
230− -- and backlog is emptied.
231− let workExhausted :: STM.STM Bool
232− workExhausted = (&&)
233− <$> STM.isClosedTBMQueue ingestionQueue
234− <*> STM.isEmptyTBMQueue ingestionQueue
235− -- Block until there's no more work
236− let waitWorkExhausted :: STM.STM ()
237− waitWorkExhausted = STM.check =<< workExhausted
238− let pop :: STM.STM (Tick AnyLogItem)
239− pop = maybe WorkExhausted NewItem <$> STM.readTBMQueue ingestionQueue
240−
241− let timeExpired :: STM.STM (Tick a)
242− timeExpired = do
243− isExpired <- STM.readTVar =<< STM.readTVar timerRef
244− STM.check isExpired
245− pure TimeExpired
246−
247− -- Circular transaction that checks for completion of work, time
248− -- expiration, or new events.
249− let nextTick :: STM.STM (Tick AnyLogItem)
250− nextTick =
251− timeExpired <|>
252− pop <|>
253− (WorkExhausted <$ waitWorkExhausted)
254−
255− -- Blocking push. This applies backpressure upstream to katip, which does its own buffering
256− let push :: AnyLogItem -> STM.STM ()
257− push = void . STM.writeTBMQueue ingestionQueue
258−
259− let sealQueue = STM.atomically (STM.closeTBMQueue ingestionQueue)
260−
261− -- Replace the timer with a new one
262− let resetTimer = STM.atomically . STM.writeTVar timerRef =<< newTimer
263−
264− -- Send the buffer and then reset the timer
265− let flush curBuffer = do
266− res <- flushBuffer config mgr curBuffer
230+ -- In STM: resolve an action with the buffers
231+ let processLogAction :: LogAction BulkBuffer -> STM.STM ()
232+ processLogAction logAction = case logAction of
233+ -- No write needed right now
234+ Buffered !newBuffer -> STM.writeTVar bufferState newBuffer
235+ -- Time to flush and put excess in the next buffer
236+ FlushNow !toFlush !newBuffer -> do
237+ STM.writeTVar bufferState newBuffer
238+ STM.putTMVar readyMVar toFlush
239+
240+ -- Recompute the new buffer state from a new log item and take action
241+ let push :: (K.LogItem a) => K.Item a -> STM.STM ()
242+ push !logItem = do
243+ !curBuffer <- STM.readTVar bufferState
244+ processLogAction (bufferItem itemBufferSize verbosity logItem curBuffer)
245+
246+
247+ -- A transaction that retries until its time to shutdown
248+ let waitShutdown :: STM.STM ()
249+ waitShutdown = do
250+ killStatus <- STM.readTVar killSwitch
251+ STM.check (killStatus == ShutDown)
252+
253+ -- Retry transaction until timer has expired. When it does, take action on the buffer.
254+ let waitTimer :: STM.STM ()
255+ waitTimer = do
256+ -- get the timer we'll be using
257+ timer <- STM.readTVar timerRef
258+ -- do not proceed until its ready
259+ timerExpired <- STM.readTVar timer
260+ STM.check timerExpired
261+ -- do not wake again until this message is out the door
262+ STM.writeTVar timer False
263+ -- we waited long enough, grab the current buffer
264+ !curBuffer <- STM.readTVar bufferState
265+ -- and tell it to flush right now
266+ processLogAction (forceFlush curBuffer)
267+
268+ -- run the timer in a loop until its time to shut down
269+ let timedFlushLoop :: IO ()
270+ timedFlushLoop = do
271+ res <- STM.atomically $
272+ raceAlt
273+ waitShutdown
274+ waitTimer
267275 case res of
268− Left e -> onErrorSafe e
269− Right () -> pure ()
270− resetTimer
271−
272− let flushLoop :: BulkBuffer -> IO ()
273− flushLoop curBuffer = do
274− tick <- STM.atomically nextTick
275− case tick of
276− WorkExhausted -> do
277− flush curBuffer -- flush what you've got
278− pure () -- stop looping
279− TimeExpired -> do
280− flush curBuffer
281− flushLoop mempty
282− NewItem (AnyLogItem item) -> do
283− case bufferItem (logzIOScribeConfiguration_bufferItems config) verbosity item curBuffer of
284− Buffered newBuffer -> flushLoop newBuffer
285− FlushNow flushThis newBuffer -> do
286− flush flushThis
287− flushLoop newBuffer
288−
289− flushThread <- Async.async (flushLoop mempty)
290−
291− let close = do
292− sealQueue
293− _ <- Async.waitCatch flushThread
276+ Left () -> pure () -- exit loop
277+ Right () -> timedFlushLoop -- wait for another timer or shutdown
278+
279+ timedFlusher <- Async.async timedFlushLoop
280+
281+ -- flush the buffer until its time to shutdown and we're out of items to log
282+ let flushLoop :: IO ()
283+ flushLoop = do
284+ -- hold until something comes in to flush or we're told to stop. If we're told to stop but get a log first, we'll catch it on the next loop
285+ readyBufferOrDie <- STM.atomically $
286+ raceAlt
287+ (STM.takeTMVar readyMVar)
288+ waitShutdown
289+ case readyBufferOrDie of
290+ Right () -> pure () -- we caught shutdown before any new data arrived
291+ -- There's some logs to write
292+ Left readyBuffer -> do
293+ -- flush in IO
294+ res <- flushBuffer config mgr readyBuffer
295+ case res of
296+ Left ex -> onErrorSafe ex
297+ Right () -> pure ()
298+ -- reset the timeout before the next flush on each write
299+ !newTimer <- STM.registerDelay itemBufferTimeoutMicros
300+ STM.atomically $ do
301+ STM.writeTVar timerRef newTimer
302+ -- start over
303+ flushLoop
304+ flusher <- Async.async flushLoop
305+
306+ -- Blocking call which sends a poison pill to the workers and waits
307+ -- for them to exit, swallowing exceptions. This follows the
308+ -- semantics that when katip calls the finalizer, no more log
309+ -- entries will be sent.
310+ let shutdown :: IO ()
311+ shutdown = do
312+ -- tell workers this is the last loop
313+ STM.atomically (STM.writeTVar killSwitch ShutDown)
314+ -- wait for workers to finish
315+ _ <- Async.waitCatch timedFlusher
316+ _ <- Async.waitCatch flusher
317+ -- stop blocking
294318 pure ()
295319
320+
296321 pure $ K.Scribe
297− { K.liPush = STM.atomically . push . AnyLogItem
298− , K.scribeFinalizer = close
322+ { K.liPush = STM.atomically . push
323+ , K.scribeFinalizer = shutdown
299324 , K.scribePermitItem = permitItem
300325 }
301326 where
302327 itemBufferSize = logzIOScribeConfiguration_bufferItems config
303328 itemBufferTimeoutMicros = ndtToMicros (logzIOScribeConfiguration_bufferTimeout config)
304329 scheme = logzIOScribeConfiguration_scheme config
330+ -- should this be async?
305331 onErrorSafe ex = do
306332 _ <- EX.tryAny (logzIOScribeConfiguration_onError config ex)
307333 pure ()
308334
309335
310−data AnyLogItem where
311− AnyLogItem :: K.LogItem a => K.Item a -> AnyLogItem
312−
313−
314−-------------------------------------------------------------------------------
315−data Tick a =
316− TimeExpired
317− | NewItem !a
318− | WorkExhausted
319−
320−
321336-------------------------------------------------------------------------------
322337-- Match the native resolution of NominalDiffTime, which is excessive
323338-- for our needs
@@ -335,6 +350,23 @@ ndtToMicros t = round (((fromIntegral (ndtPicos t)) :: Double) / picosInMicro)
335350 picosInMicro = 10 ^ (3 :: Int)
336351
337352
353+-------------------------------------------------------------------------------
354+-- | An internal mode flag used to shut down worker
355+-- threads. Internally, we assume htere is no going back from
356+-- ShutDown.
357+data KillSwitch =
358+ Continue
359+ | ShutDown
360+ deriving (Eq)
361+
362+
363+-------------------------------------------------------------------------------
364+-- | When applied to STM, it will try the left side first and if that
365+-- retries, the right, and if that retries, start from the beginning.
366+raceAlt :: (Alternative m) => m a -> m b -> m (Either a b)
367+raceAlt a b = (Left <$> a) <|> (Right <$> b)
368+
369+
338370-------------------------------------------------------------------------------
339371data LoggingError =
340372 URIError HTTP.HttpException