@@ -1,5 +1,4 @@
11{-# LANGUAGE BangPatterns #-}
2−{-# LANGUAGE GADTs #-}
32{-# LANGUAGE GeneralizedNewtypeDeriving #-}
43{-# LANGUAGE OverloadedStrings #-}
54{-# LANGUAGE RankNTypes #-}
@@ -39,34 +38,32 @@ 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.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+-------------------------------------------------------------------------------
7067-------------------------------------------------------------------------------
7168
7269
@@ -200,7 +197,7 @@ httpPort = URIBS.Port 8070
200197-- | A reasonable retry policy: exponential backoff with 25ms base
201198-- delay up to 5 retries, for a total cumulative delay of 775ms.
202199defaultRetryPolicy :: (Monad m) => Retry.RetryPolicyM m
203−defaultRetryPolicy = Retry.exponentialBackoff 25000 `mappend` Retry.limitRetries 5
200+defaultRetryPolicy = Retry.exponentialBackoff 25000 <> Retry.limitRetries 5
204201
205202-------------------------------------------------------------------------------
206203mkLogzIOScribe
@@ -209,115 +206,132 @@ mkLogzIOScribe
209206 -> K.Verbosity
210207 -> IO K.Scribe
211208mkLogzIOScribe 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
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
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− -- 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
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
267274 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
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
294317 pure ()
295318
319+
296320 pure $ K.Scribe
297− { K.liPush = STM.atomically . push . AnyLogItem
298− , K.scribeFinalizer = close
321+ { K.liPush = STM.atomically . push
322+ , K.scribeFinalizer = shutdown
299323 , K.scribePermitItem = permitItem
300324 }
301325 where
302326 itemBufferSize = logzIOScribeConfiguration_bufferItems config
303327 itemBufferTimeoutMicros = ndtToMicros (logzIOScribeConfiguration_bufferTimeout config)
304328 scheme = logzIOScribeConfiguration_scheme config
329+ -- should this be async?
305330 onErrorSafe ex = do
306331 _ <- EX.tryAny (logzIOScribeConfiguration_onError config ex)
307332 pure ()
308333
309334
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−
321335-------------------------------------------------------------------------------
322336-- Match the native resolution of NominalDiffTime, which is excessive
323337-- for our needs
@@ -335,6 +349,23 @@ ndtToMicros t = round (((fromIntegral (ndtPicos t)) :: Double) / picosInMicro)
335349 picosInMicro = 10 ^ (3 :: Int)
336350
337351
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+
338369-------------------------------------------------------------------------------
339370data LoggingError =
340371 URIError HTTP.HttpException
@@ -411,7 +442,7 @@ flushBuffer config mgr bulkBuffer
411442-------------------------------------------------------------------------------
412443newtype Bytes = Bytes
413444 { bytes :: Int64
414− } deriving (Show, Eq, Num, Ord, Bounded)
445+ } deriving (Show, Eq, Num, Ord)
415446
416447
417448-- | How big of a body can we send? The limit is defined
@@ -513,7 +544,7 @@ data BulkBuffer = BulkBuffer
513544 }
514545
515546
516−instance Semigroup.Semigroup BulkBuffer where
547+instance Semigroup BulkBuffer where
517548 (BulkBuffer bytesUsedA bufferA itemCountA) <>
518549 (BulkBuffer bytesUsedB bufferB itemCountB) = BulkBuffer
519550 (bytesUsedA + bytesUsedB)
@@ -567,7 +598,7 @@ bufferItem' customMaxPayload customMaxItem maxItems verb item bulkBuffer =
567598 let (encodedLine, spaceNeeded) = renderLineTruncated' customMaxItem verb item
568599 newBytesUsed = bulkBuffer_bytesUsed bulkBuffer + spaceNeeded
569600 newItemCount = bulkBuffer_itemCount bulkBuffer + 1
570− in if newItemCount >= maxItems || newBytesUsed >= customMaxPayload
601+ in if newItemCount > maxItems || newBytesUsed > customMaxPayload
571602 then FlushNow
572603 bulkBuffer
573604 BulkBuffer