@@ -1,4 +1,5 @@
11{-# LANGUAGE BangPatterns #-}
2+{-# LANGUAGE GADTs #-}
23{-# LANGUAGE GeneralizedNewtypeDeriving #-}
34{-# LANGUAGE OverloadedStrings #-}
45{-# LANGUAGE RankNTypes #-}
@@ -38,32 +39,34 @@ module Katip.Scribes.LogzIO.HTTPS
3839
3940-------------------------------------------------------------------------------
4041import Control.Applicative
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
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
5154import Data.Int
52−import qualified Data.Scientific as Scientific
53−import Data.String (IsString)
54−import qualified Data.Text as T
55−import qualified Data.Text.Encoding as TE
56−import qualified Data.Text.Lazy as TL
57−import qualified Data.Text.Lazy.Builder as TB
58−import qualified Data.Time as Time
59−import qualified Katip as K
60−import Katip.Core (LocJs (..))
61−import qualified Network.HTTP.Client as HTTP
62−import qualified Network.HTTP.Client.TLS as HTTPS
63−import qualified Network.HTTP.Types as HTypes
64−import qualified System.Posix.Types as POSIX
65−import qualified URI.ByteString as URIBS
66−-------------------------------------------------------------------------------
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
6770-------------------------------------------------------------------------------
6871
6972
@@ -197,7 +200,7 @@ httpPort = URIBS.Port 8070
197200-- | A reasonable retry policy: exponential backoff with 25ms base
198201-- delay up to 5 retries, for a total cumulative delay of 775ms.
199202defaultRetryPolicy :: (Monad m) => Retry.RetryPolicyM m
200−defaultRetryPolicy = Retry.exponentialBackoff 25000 <> Retry.limitRetries 5
203+defaultRetryPolicy = Retry.exponentialBackoff 25000 `mappend` Retry.limitRetries 5
201204
202205-------------------------------------------------------------------------------
203206mkLogzIOScribe
@@ -206,132 +209,115 @@ mkLogzIOScribe
206209 -> K.Verbosity
207210 -> IO K.Scribe
208211mkLogzIOScribe config permitItem verbosity = do
209− -- Accumulating buffer for each item
210− bufferState <- STM.newTVarIO (mempty :: BulkBuffer)
211− -- Single-slot TMVar when something is ready to write. Katip buffers
212− -- for us upstream so we can simplify shutdown and flush guarantees
213− -- by synchronously writing when we need to.
214− readyMVar <- STM.newEmptyTMVarIO :: IO (STM.TMVar BulkBuffer)
215−
216− -- We create a TVar that is set to ShutDown when its time to stop
217− killSwitch <- STM.newTVarIO Continue
218−
219− -- Create a mutable ref to whether the timer has expired. This lets
220− -- us replace the timer when its time to start timing again. We
221− -- reset the timer __after__ a write goes out, not as soon as it expires
222− timerRef <- STM.newTVarIO =<< STM.registerDelay itemBufferTimeoutMicros
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
223223
224224 -- Set up a connection manager for requests
225225 mgr <- case scheme of
226226 HTTPS -> HTTPS.newTlsManager
227227 HTTP -> HTTP.newManager HTTP.defaultManagerSettings
228228
229− -- In STM: resolve an action with the buffers
230− let processLogAction :: LogAction BulkBuffer -> STM.STM ()
231− processLogAction logAction = case logAction of
232− -- No write needed right now
233− Buffered !newBuffer -> STM.writeTVar bufferState newBuffer
234− -- Time to flush and put excess in the next buffer
235− FlushNow !toFlush !newBuffer -> do
236− STM.writeTVar bufferState newBuffer
237− STM.putTMVar readyMVar toFlush
238−
239− -- Recompute the new buffer state from a new log item and take action
240− let push :: (K.LogItem a) => K.Item a -> STM.STM ()
241− push !logItem = do
242− !curBuffer <- STM.readTVar bufferState
243− processLogAction (bufferItem itemBufferSize verbosity logItem curBuffer)
244−
245−
246− -- A transaction that retries until its time to shutdown
247− let waitShutdown :: STM.STM ()
248− waitShutdown = do
249− killStatus <- STM.readTVar killSwitch
250− STM.check (killStatus == ShutDown)
251−
252− -- Retry transaction until timer has expired. When it does, take action on the buffer.
253− let waitTimer :: STM.STM ()
254− waitTimer = do
255− -- get the timer we'll be using
256− timer <- STM.readTVar timerRef
257− -- do not proceed until its ready
258− timerExpired <- STM.readTVar timer
259− STM.check timerExpired
260− -- do not wake again until this message is out the door
261− STM.writeTVar timer False
262− -- we waited long enough, grab the current buffer
263− !curBuffer <- STM.readTVar bufferState
264− -- and tell it to flush right now
265− processLogAction (forceFlush curBuffer)
266−
267− -- run the timer in a loop until its time to shut down
268− let timedFlushLoop :: IO ()
269− timedFlushLoop = do
270− res <- STM.atomically $
271− raceAlt
272− waitShutdown
273− waitTimer
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
274267 case res of
275− Left () -> pure () -- exit loop
276− Right () -> timedFlushLoop -- wait for another timer or shutdown
277−
278− timedFlusher <- Async.async timedFlushLoop
279−
280− -- flush the buffer until its time to shutdown and we're out of items to log
281− let flushLoop :: IO ()
282− flushLoop = do
283− -- 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
284− readyBufferOrDie <- STM.atomically $
285− raceAlt
286− (STM.takeTMVar readyMVar)
287− waitShutdown
288− case readyBufferOrDie of
289− Right () -> pure () -- we caught shutdown before any new data arrived
290− -- There's some logs to write
291− Left readyBuffer -> do
292− -- flush in IO
293− res <- flushBuffer config mgr readyBuffer
294− case res of
295− Left ex -> onErrorSafe ex
296− Right () -> pure ()
297− -- reset the timeout before the next flush on each write
298− !newTimer <- STM.registerDelay itemBufferTimeoutMicros
299− STM.atomically $ do
300− STM.writeTVar timerRef newTimer
301− -- start over
302− flushLoop
303− flusher <- Async.async flushLoop
304−
305− -- Blocking call which sends a poison pill to the workers and waits
306− -- for them to exit, swallowing exceptions. This follows the
307− -- semantics that when katip calls the finalizer, no more log
308− -- entries will be sent.
309− let shutdown :: IO ()
310− shutdown = do
311− -- tell workers this is the last loop
312− STM.atomically (STM.writeTVar killSwitch ShutDown)
313− -- wait for workers to finish
314− _ <- Async.waitCatch timedFlusher
315− _ <- Async.waitCatch flusher
316− -- stop blocking
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
317294 pure ()
318295
319−
320296 pure $ K.Scribe
321− { K.liPush = STM.atomically . push
322− , K.scribeFinalizer = shutdown
297+ { K.liPush = STM.atomically . push . AnyLogItem
298+ , K.scribeFinalizer = close
323299 , K.scribePermitItem = permitItem
324300 }
325301 where
326302 itemBufferSize = logzIOScribeConfiguration_bufferItems config
327303 itemBufferTimeoutMicros = ndtToMicros (logzIOScribeConfiguration_bufferTimeout config)
328304 scheme = logzIOScribeConfiguration_scheme config
329− -- should this be async?
330305 onErrorSafe ex = do
331306 _ <- EX.tryAny (logzIOScribeConfiguration_onError config ex)
332307 pure ()
333308
334309
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+
335321-------------------------------------------------------------------------------
336322-- Match the native resolution of NominalDiffTime, which is excessive
337323-- for our needs
@@ -349,23 +335,6 @@ ndtToMicros t = round (((fromIntegral (ndtPicos t)) :: Double) / picosInMicro)
349335 picosInMicro = 10 ^ (3 :: Int)
350336
351337
352−-------------------------------------------------------------------------------
353−-- | An internal mode flag used to shut down worker
354−-- threads. Internally, we assume htere is no going back from
355−-- ShutDown.
356−data KillSwitch =
357− Continue
358− | ShutDown
359− deriving (Eq)
360−
361−
362−-------------------------------------------------------------------------------
363−-- | When applied to STM, it will try the left side first and if that
364−-- retries, the right, and if that retries, start from the beginning.
365−raceAlt :: (Alternative m) => m a -> m b -> m (Either a b)
366−raceAlt a b = (Left <$> a) <|> (Right <$> b)
367−
368−
369338-------------------------------------------------------------------------------
370339data LoggingError =
371340 URIError HTTP.HttpException
@@ -442,7 +411,7 @@ flushBuffer config mgr bulkBuffer
442411-------------------------------------------------------------------------------
443412newtype Bytes = Bytes
444413 { bytes :: Int64
445− } deriving (Show, Eq, Num, Ord)
414+ } deriving (Show, Eq, Num, Ord, Bounded)
446415
447416
448417-- | How big of a body can we send? The limit is defined
@@ -544,7 +513,7 @@ data BulkBuffer = BulkBuffer
544513 }
545514
546515
547−instance Semigroup BulkBuffer where
516+instance Semigroup.Semigroup BulkBuffer where
548517 (BulkBuffer bytesUsedA bufferA itemCountA) <>
549518 (BulkBuffer bytesUsedB bufferB itemCountB) = BulkBuffer
550519 (bytesUsedA + bytesUsedB)
@@ -598,7 +567,7 @@ bufferItem' customMaxPayload customMaxItem maxItems verb item bulkBuffer =
598567 let (encodedLine, spaceNeeded) = renderLineTruncated' customMaxItem verb item
599568 newBytesUsed = bulkBuffer_bytesUsed bulkBuffer + spaceNeeded
600569 newItemCount = bulkBuffer_itemCount bulkBuffer + 1
601− in if newItemCount > maxItems || newBytesUsed > customMaxPayload
570+ in if newItemCount >= maxItems || newBytesUsed >= customMaxPayload
602571 then FlushNow
603572 bulkBuffer
604573 BulkBuffer