@@ -1,4 +1,5 @@
11{-# LANGUAGE BangPatterns #-}
2+{-# LANGUAGE GADTs #-}
23{-# LANGUAGE GeneralizedNewtypeDeriving #-}
34{-# LANGUAGE OverloadedStrings #-}
45{-# LANGUAGE RankNTypes #-}
@@ -38,33 +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.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−-------------------------------------------------------------------------------
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
6870-------------------------------------------------------------------------------
6971
7072
@@ -207,132 +209,115 @@ mkLogzIOScribe
207209 -> K.Verbosity
208210 -> IO K.Scribe
209211mkLogzIOScribe config permitItem verbosity = do
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
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
224223
225224 -- Set up a connection manager for requests
226225 mgr <- case scheme of
227226 HTTPS -> HTTPS.newTlsManager
228227 HTTP -> HTTP.newManager HTTP.defaultManagerSettings
229228
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
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
275267 case res of
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
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
318294 pure ()
319295
320−
321296 pure $ K.Scribe
322− { K.liPush = STM.atomically . push
323− , K.scribeFinalizer = shutdown
297+ { K.liPush = STM.atomically . push . AnyLogItem
298+ , K.scribeFinalizer = close
324299 , K.scribePermitItem = permitItem
325300 }
326301 where
327302 itemBufferSize = logzIOScribeConfiguration_bufferItems config
328303 itemBufferTimeoutMicros = ndtToMicros (logzIOScribeConfiguration_bufferTimeout config)
329304 scheme = logzIOScribeConfiguration_scheme config
330− -- should this be async?
331305 onErrorSafe ex = do
332306 _ <- EX.tryAny (logzIOScribeConfiguration_onError config ex)
333307 pure ()
334308
335309
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+
336321-------------------------------------------------------------------------------
337322-- Match the native resolution of NominalDiffTime, which is excessive
338323-- for our needs
@@ -350,23 +335,6 @@ ndtToMicros t = round (((fromIntegral (ndtPicos t)) :: Double) / picosInMicro)
350335 picosInMicro = 10 ^ (3 :: Int)
351336
352337
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−
370338-------------------------------------------------------------------------------
371339data LoggingError =
372340 URIError HTTP.HttpException