@@ -1,5 +1,4 @@
11{-# LANGUAGE BangPatterns #-}
2−{-# LANGUAGE GADTs #-}
32{-# LANGUAGE GeneralizedNewtypeDeriving #-}
43{-# LANGUAGE OverloadedStrings #-}
54{-# LANGUAGE RankNTypes #-}
@@ -39,34 +38,31 @@ 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 Data.String (IsString)
53+import qualified Data.Text as T
54+import qualified Data.Text.Encoding as TE
55+import qualified Data.Text.Lazy as TL
56+import qualified Data.Text.Lazy.Builder as TB
57+import qualified Data.Time as Time
58+import qualified Katip as K
59+import Katip.Core (LocJs (..))
60+import qualified Network.HTTP.Client as HTTP
61+import qualified Network.HTTP.Client.TLS as HTTPS
62+import qualified Network.HTTP.Types as HTypes
63+import qualified System.Posix.Types as POSIX
64+import qualified URI.ByteString as URIBS
65+-------------------------------------------------------------------------------
7066-------------------------------------------------------------------------------
7167
7268
@@ -200,7 +196,7 @@ httpPort = URIBS.Port 8070
200196-- | A reasonable retry policy: exponential backoff with 25ms base
201197-- delay up to 5 retries, for a total cumulative delay of 775ms.
202198defaultRetryPolicy :: (Monad m) => Retry.RetryPolicyM m
203−defaultRetryPolicy = Retry.exponentialBackoff 25000 `mappend` Retry.limitRetries 5
199+defaultRetryPolicy = Retry.exponentialBackoff 25000 <> Retry.limitRetries 5
204200
205201-------------------------------------------------------------------------------
206202mkLogzIOScribe
@@ -209,115 +205,132 @@ mkLogzIOScribe
209205 -> K.Verbosity
210206 -> IO K.Scribe
211207mkLogzIOScribe 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
208+ -- Accumulating buffer for each item
209+ bufferState <- STM.newTVarIO (mempty :: BulkBuffer)
210+ -- Single-slot TMVar when something is ready to write. Katip buffers
211+ -- for us upstream so we can simplify shutdown and flush guarantees
212+ -- by synchronously writing when we need to.
213+ readyMVar <- STM.newEmptyTMVarIO :: IO (STM.TMVar BulkBuffer)
214+
215+ -- We create a TVar that is set to ShutDown when its time to stop
216+ killSwitch <- STM.newTVarIO Continue
217+
218+ -- Create a mutable ref to whether the timer has expired. This lets
219+ -- us replace the timer when its time to start timing again. We
220+ -- reset the timer __after__ a write goes out, not as soon as it expires
221+ timerRef <- STM.newTVarIO =<< STM.registerDelay itemBufferTimeoutMicros
223222
224223 -- Set up a connection manager for requests
225224 mgr <- case scheme of
226225 HTTPS -> HTTPS.newTlsManager
227226 HTTP -> HTTP.newManager HTTP.defaultManagerSettings
228227
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
228+ -- In STM: resolve an action with the buffers
229+ let processLogAction :: LogAction BulkBuffer -> STM.STM ()
230+ processLogAction logAction = case logAction of
231+ -- No write needed right now
232+ Buffered !newBuffer -> STM.writeTVar bufferState newBuffer
233+ -- Time to flush and put excess in the next buffer
234+ FlushNow !toFlush !newBuffer -> do
235+ STM.writeTVar bufferState newBuffer
236+ STM.putTMVar readyMVar toFlush
237+
238+ -- Recompute the new buffer state from a new log item and take action
239+ let push :: (K.LogItem a) => K.Item a -> STM.STM ()
240+ push !logItem = do
241+ !curBuffer <- STM.readTVar bufferState
242+ processLogAction (bufferItem itemBufferSize verbosity logItem curBuffer)
243+
244+
245+ -- A transaction that retries until its time to shutdown
246+ let waitShutdown :: STM.STM ()
247+ waitShutdown = do
248+ killStatus <- STM.readTVar killSwitch
249+ STM.check (killStatus == ShutDown)
250+
251+ -- Retry transaction until timer has expired. When it does, take action on the buffer.
252+ let waitTimer :: STM.STM ()
253+ waitTimer = do
254+ -- get the timer we'll be using
255+ timer <- STM.readTVar timerRef
256+ -- do not proceed until its ready
257+ timerExpired <- STM.readTVar timer
258+ STM.check timerExpired
259+ -- do not wake again until this message is out the door
260+ STM.writeTVar timer False
261+ -- we waited long enough, grab the current buffer
262+ !curBuffer <- STM.readTVar bufferState
263+ -- and tell it to flush right now
264+ processLogAction (forceFlush curBuffer)
265+
266+ -- run the timer in a loop until its time to shut down
267+ let timedFlushLoop :: IO ()
268+ timedFlushLoop = do
269+ res <- STM.atomically $
270+ raceAlt
271+ waitShutdown
272+ waitTimer
267273 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
274+ Left () -> pure () -- exit loop
275+ Right () -> timedFlushLoop -- wait for another timer or shutdown
276+
277+ timedFlusher <- Async.async timedFlushLoop
278+
279+ -- flush the buffer until its time to shutdown and we're out of items to log
280+ let flushLoop :: IO ()
281+ flushLoop = do
282+ -- 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
283+ readyBufferOrDie <- STM.atomically $
284+ raceAlt
285+ (STM.takeTMVar readyMVar)
286+ waitShutdown
287+ case readyBufferOrDie of
288+ Right () -> pure () -- we caught shutdown before any new data arrived
289+ -- There's some logs to write
290+ Left readyBuffer -> do
291+ -- flush in IO
292+ res <- flushBuffer config mgr readyBuffer
293+ case res of
294+ Left ex -> onErrorSafe ex
295+ Right () -> pure ()
296+ -- reset the timeout before the next flush on each write
297+ !newTimer <- STM.registerDelay itemBufferTimeoutMicros
298+ STM.atomically $ do
299+ STM.writeTVar timerRef newTimer
300+ -- start over
301+ flushLoop
302+ flusher <- Async.async flushLoop
303+
304+ -- Blocking call which sends a poison pill to the workers and waits
305+ -- for them to exit, swallowing exceptions. This follows the
306+ -- semantics that when katip calls the finalizer, no more log
307+ -- entries will be sent.
308+ let shutdown :: IO ()
309+ shutdown = do
310+ -- tell workers this is the last loop
311+ STM.atomically (STM.writeTVar killSwitch ShutDown)
312+ -- wait for workers to finish
313+ _ <- Async.waitCatch timedFlusher
314+ _ <- Async.waitCatch flusher
315+ -- stop blocking
294316 pure ()
295317
318+
296319 pure $ K.Scribe
297− { K.liPush = STM.atomically . push . AnyLogItem
298− , K.scribeFinalizer = close
320+ { K.liPush = STM.atomically . push
321+ , K.scribeFinalizer = shutdown
299322 , K.scribePermitItem = permitItem
300323 }
301324 where
302325 itemBufferSize = logzIOScribeConfiguration_bufferItems config
303326 itemBufferTimeoutMicros = ndtToMicros (logzIOScribeConfiguration_bufferTimeout config)
304327 scheme = logzIOScribeConfiguration_scheme config
328+ -- should this be async?
305329 onErrorSafe ex = do
306330 _ <- EX.tryAny (logzIOScribeConfiguration_onError config ex)
307331 pure ()
308332
309333
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−
321334-------------------------------------------------------------------------------
322335-- Match the native resolution of NominalDiffTime, which is excessive
323336-- for our needs
@@ -335,6 +348,23 @@ ndtToMicros t = round (((fromIntegral (ndtPicos t)) :: Double) / picosInMicro)
335348 picosInMicro = 10 ^ (3 :: Int)
336349
337350
351+-------------------------------------------------------------------------------
352+-- | An internal mode flag used to shut down worker
353+-- threads. Internally, we assume htere is no going back from
354+-- ShutDown.
355+data KillSwitch =
356+ Continue
357+ | ShutDown
358+ deriving (Eq)
359+
360+
361+-------------------------------------------------------------------------------
362+-- | When applied to STM, it will try the left side first and if that
363+-- retries, the right, and if that retries, start from the beginning.
364+raceAlt :: (Alternative m) => m a -> m b -> m (Either a b)
365+raceAlt a b = (Left <$> a) <|> (Right <$> b)
366+
367+
338368-------------------------------------------------------------------------------
339369data LoggingError =
340370 URIError HTTP.HttpException
@@ -411,7 +441,7 @@ flushBuffer config mgr bulkBuffer
411441-------------------------------------------------------------------------------
412442newtype Bytes = Bytes
413443 { bytes :: Int64
414− } deriving (Show, Eq, Num, Ord, Bounded)
444+ } deriving (Show, Eq, Num, Ord)
415445
416446
417447-- | How big of a body can we send? The limit is defined
@@ -442,7 +472,7 @@ fullItemObject verbosity item = HM.fromList
442472 , "thread" A..= K.getThreadIdText (K._itemThread item)
443473 , "host" A..= K._itemHost item
444474 , "pid" A..= pidInt
445− , "data" A..= annotateKeys (K.payloadObject verbosity (K._itemPayload item))
475+ , "data" A..= K.payloadObject verbosity (K._itemPayload item)
446476 -- Slight deviation, logz.io uses "message" instead of "msg"
447477 , "message" A..= TB.toLazyText (K.unLogStr (K._itemMessage item))
448478 -- Another slight deviation, logz.io uses "@timestamp" instead of
@@ -513,7 +543,7 @@ data BulkBuffer = BulkBuffer
513543 }
514544
515545
516−instance Semigroup.Semigroup BulkBuffer where
546+instance Semigroup BulkBuffer where
517547 (BulkBuffer bytesUsedA bufferA itemCountA) <>
518548 (BulkBuffer bytesUsedB bufferB itemCountB) = BulkBuffer
519549 (bytesUsedA + bytesUsedB)
@@ -567,7 +597,7 @@ bufferItem' customMaxPayload customMaxItem maxItems verb item bulkBuffer =
567597 let (encodedLine, spaceNeeded) = renderLineTruncated' customMaxItem verb item
568598 newBytesUsed = bulkBuffer_bytesUsed bulkBuffer + spaceNeeded
569599 newItemCount = bulkBuffer_itemCount bulkBuffer + 1
570− in if newItemCount >= maxItems || newBytesUsed >= customMaxPayload
600+ in if newItemCount > maxItems || newBytesUsed > customMaxPayload
571601 then FlushNow
572602 bulkBuffer
573603 BulkBuffer
@@ -587,53 +617,3 @@ bufferItem' customMaxPayload customMaxItem maxItems verb item bulkBuffer =
587617-- will always return 'FlushNow' with an empty new buffer.
588618forceFlush :: (Monoid buf) => buf -> LogAction buf
589619forceFlush buf = FlushNow buf mempty
590−
591−
592−-------------------------------------------------------------------------------
593−-- Annotation borrowed from katip-elasticsearch. There are fixed
594−-- fields in katip logs which should stay the same type forever and
595−-- thus won't need annotation. However, any code in userland may
596−-- choose to add akey to their log's metadata with a type that's
597−-- incompatible with the mapping. If the first value that's picked up
598−-- restrictive (e.g. a long), subsequent values will be rejected. By
599−-- differentiating types by their name, this is guaranteed not to
600−-- happen.
601−
602−annotateValue :: A.Value -> A.Value
603−annotateValue (A.Object o) = A.Object (annotateKeys o)
604−annotateValue (A.Array a) = A.Array (annotateValue <$> a)
605−annotateValue x = x
606−
607−
608−annotateKeys :: A.Object -> A.Object
609−annotateKeys = HM.fromList . map go . HM.toList
610− where
611− go (k, A.Object o) = (k, A.Object (annotateKeys o))
612− go (k, A.Array a) = (k, A.Array (annotateValue <$> a))
613− go (k, s@(A.String _)) = (k <> stringAnn, s)
614− go (k, n@(A.Number sci)) = if Scientific.isFloating sci
615− then (k <> doubleAnn, n)
616− else (k <> longAnn, n)
617− go (k, b@(A.Bool _)) = (k <> booleanAnn, b)
618− go (k, A.Null) = (k <> nullAnn, A.Null)
619−
620−
621−-------------------------------------------------------------------------------
622−-- Annotation Constants
623−-------------------------------------------------------------------------------
624−
625−
626−stringAnn :: T.Text
627−stringAnn = "::s"
628−
629−doubleAnn :: T.Text
630−doubleAnn = "::d"
631−
632−longAnn :: T.Text
633−longAnn = "::l"
634−
635−booleanAnn :: T.Text
636−booleanAnn = "::b"
637−
638−nullAnn :: T.Text
639−nullAnn = "::n"