{-# LANGUAGE OverloadedStrings #-}
module Data.Conduit.LogFile
    ( RotatingLog
    , openRotatingLog
    , addChunk
    , close
    , defaultMaxTotal
    , dummy
    ) where

import           Control.Concurrent             (forkIO)
import           Control.Concurrent.STM         (atomically)
import           Control.Concurrent.STM.TBQueue
import           Control.Concurrent.STM.TVar
import           Control.Exception              (bracket, bracketOnError,
                                                 finally)
import           Control.Monad                  (void, when)
import qualified Data.ByteString                as S
import           Data.Time                      (UTCTime, getCurrentTime)
import           Data.Word                      (Word)
import           System.Directory               (createDirectoryIfMissing,
                                                 doesFileExist, renameFile)
import           System.FilePath                ((<.>), (</>))
import qualified System.IO                      as SIO
import           System.IO.Unsafe               (unsafePerformIO)
import           System.Mem.Weak                (addFinalizer)

data Command = AddChunk !S.ByteString
             | Close

-- | Represents a folder used for totating log files.
--
-- Since 0.2.1
data RotatingLog = RotatingLog !(TVar State)
-- Use a data instead of a newtype so that we can attach a finalizer.

-- | A @RotatingLog@ which performs no logging.
--
-- Since 0.2.1
dummy :: RotatingLog
dummy :: RotatingLog
dummy = TVar State -> RotatingLog
RotatingLog (TVar State -> RotatingLog) -> TVar State -> RotatingLog
forall a b. (a -> b) -> a -> b
$! IO (TVar State) -> TVar State
forall a. IO a -> a
unsafePerformIO (IO (TVar State) -> TVar State) -> IO (TVar State) -> TVar State
forall a b. (a -> b) -> a -> b
$! State -> IO (TVar State)
forall a. a -> IO (TVar a)
newTVarIO State
Closed

data State = Closed
           | Running !SIO.Handle !(TBQueue Command)

queue :: Command -> RotatingLog -> IO ()
queue :: Command -> RotatingLog -> IO ()
queue Command
cmd (RotatingLog TVar State
ts) = STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
    State
s <- TVar State -> STM State
forall a. TVar a -> STM a
readTVar TVar State
ts
    case State
s of
        State
Closed -> () -> STM ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
        Running Handle
_ TBQueue Command
q -> TBQueue Command -> Command -> STM ()
forall a. TBQueue a -> a -> STM ()
writeTBQueue TBQueue Command
q Command
cmd

addChunk :: RotatingLog -> S.ByteString -> IO ()
addChunk :: RotatingLog -> ByteString -> IO ()
addChunk RotatingLog
lf ByteString
bs = Command -> RotatingLog -> IO ()
queue (ByteString -> Command
AddChunk ByteString
bs) RotatingLog
lf

close :: RotatingLog -> IO ()
close :: RotatingLog -> IO ()
close = Command -> RotatingLog -> IO ()
queue Command
Close

-- | Create a new @RotatingLog@.
--
-- Since 0.2.1
openRotatingLog :: FilePath -- ^ folder to contain logs
                -> Word -- ^ maximum log file size, in bytes
                -> IO RotatingLog
openRotatingLog :: FilePath -> Word -> IO RotatingLog
openRotatingLog FilePath
dir Word
maxTotal = do
    Bool -> FilePath -> IO ()
createDirectoryIfMissing Bool
True FilePath
dir
    IO Handle
-> (Handle -> IO ())
-> (Handle -> IO RotatingLog)
-> IO RotatingLog
forall a b c. IO a -> (a -> IO b) -> (a -> IO c) -> IO c
bracketOnError (FilePath -> IO Handle
moveCurrent FilePath
dir) Handle -> IO ()
SIO.hClose ((Handle -> IO RotatingLog) -> IO RotatingLog)
-> (Handle -> IO RotatingLog) -> IO RotatingLog
forall a b. (a -> b) -> a -> b
$ \Handle
handle -> do
        TBQueue Command
queue' <- Natural -> IO (TBQueue Command)
forall a. Natural -> IO (TBQueue a)
newTBQueueIO Natural
5
        let s :: State
s = Handle -> TBQueue Command -> State
Running Handle
handle TBQueue Command
queue'
        TVar State
ts <- State -> IO (TVar State)
forall a. a -> IO (TVar a)
newTVarIO State
s
        IO ThreadId -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO ThreadId -> IO ()) -> IO ThreadId -> IO ()
forall a b. (a -> b) -> a -> b
$ IO () -> IO ThreadId
forkIO (IO () -> IO ThreadId) -> IO () -> IO ThreadId
forall a b. (a -> b) -> a -> b
$ FilePath -> TVar State -> Word -> IO ()
loop FilePath
dir TVar State
ts Word
maxTotal
        let rl :: RotatingLog
rl = TVar State -> RotatingLog
RotatingLog TVar State
ts
        RotatingLog -> IO () -> IO ()
forall key. key -> IO () -> IO ()
addFinalizer RotatingLog
rl (STM () -> IO ()
forall a. STM a -> IO a
atomically (TBQueue Command -> Command -> STM ()
forall a. TBQueue a -> a -> STM ()
writeTBQueue TBQueue Command
queue' Command
Close))
        RotatingLog -> IO RotatingLog
forall (m :: * -> *) a. Monad m => a -> m a
return RotatingLog
rl

current :: FilePath -- ^ folder containing logs
        -> FilePath
current :: FilePath -> FilePath
current = (FilePath -> FilePath -> FilePath
</> FilePath
"current.log")

moveCurrent :: FilePath -- ^ folder containing logs
            -> IO SIO.Handle -- ^ new handle
moveCurrent :: FilePath -> IO Handle
moveCurrent FilePath
dir = do
    let curr :: FilePath
curr = FilePath -> FilePath
current FilePath
dir
    Bool
x <- FilePath -> IO Bool
doesFileExist FilePath
curr
    Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
x (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
        UTCTime
now <- IO UTCTime
getCurrentTime
        FilePath -> FilePath -> IO ()
renameFile FilePath
curr (FilePath -> IO ()) -> FilePath -> IO ()
forall a b. (a -> b) -> a -> b
$ FilePath
dir FilePath -> FilePath -> FilePath
</> UTCTime -> FilePath
suffix UTCTime
now
    FilePath -> IOMode -> IO Handle
SIO.openFile FilePath
curr IOMode
SIO.WriteMode

suffix :: UTCTime -> FilePath
suffix :: UTCTime -> FilePath
suffix UTCTime
now =
    ((Char -> FilePath) -> FilePath -> FilePath
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Char -> FilePath
fix (FilePath -> FilePath) -> FilePath -> FilePath
forall a b. (a -> b) -> a -> b
$ (Char -> Bool) -> FilePath -> FilePath
forall a. (a -> Bool) -> [a] -> [a]
takeWhile (Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
/= Char
'.') (FilePath -> FilePath) -> FilePath -> FilePath
forall a b. (a -> b) -> a -> b
$ UTCTime -> FilePath
forall a. Show a => a -> FilePath
show UTCTime
now) FilePath -> FilePath -> FilePath
<.> FilePath
"log"
  where
    fix :: Char -> FilePath
fix Char
' ' = FilePath
"_"
    fix Char
c | Char
'0' Char -> Char -> Bool
forall a. Ord a => a -> a -> Bool
<= Char
c Bool -> Bool -> Bool
&& Char
c Char -> Char -> Bool
forall a. Ord a => a -> a -> Bool
<= Char
'9' = [Char
c]
    fix Char
_ = FilePath
""

loop :: FilePath -- ^ folder containing logs
     -> TVar State
     -> Word -- ^ maximum total log size
     -> IO ()
loop :: FilePath -> TVar State -> Word -> IO ()
loop FilePath
dir TVar State
ts Word
maxTotal =
    Word -> IO ()
go Word
0 IO () -> IO () -> IO ()
forall a b. IO a -> IO b -> IO a
`finally` (IO ()
closeCurrentHandle IO () -> IO () -> IO ()
forall a b. IO a -> IO b -> IO a
`finally` STM () -> IO ()
forall a. STM a -> IO a
atomically (TVar State -> State -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar TVar State
ts State
Closed))
  where
    closeCurrentHandle :: IO ()
closeCurrentHandle = IO (Maybe Handle)
-> (Maybe Handle -> IO ()) -> (Maybe Handle -> IO ()) -> IO ()
forall a b c. IO a -> (a -> IO b) -> (a -> IO c) -> IO c
bracket
        (STM (Maybe Handle) -> IO (Maybe Handle)
forall a. STM a -> IO a
atomically (STM (Maybe Handle) -> IO (Maybe Handle))
-> STM (Maybe Handle) -> IO (Maybe Handle)
forall a b. (a -> b) -> a -> b
$ do
            State
s <- TVar State -> STM State
forall a. TVar a -> STM a
readTVar TVar State
ts
            case State
s of
                State
Closed -> Maybe Handle -> STM (Maybe Handle)
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe Handle
forall a. Maybe a
Nothing
                Running Handle
h TBQueue Command
_ -> Maybe Handle -> STM (Maybe Handle)
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe Handle -> STM (Maybe Handle))
-> Maybe Handle -> STM (Maybe Handle)
forall a b. (a -> b) -> a -> b
$! Handle -> Maybe Handle
forall a. a -> Maybe a
Just Handle
h)
        (IO () -> (Handle -> IO ()) -> Maybe Handle -> IO ()
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (() -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()) Handle -> IO ()
SIO.hClose)
        (IO () -> Maybe Handle -> IO ()
forall a b. a -> b -> a
const (IO () -> Maybe Handle -> IO ()) -> IO () -> Maybe Handle -> IO ()
forall a b. (a -> b) -> a -> b
$ () -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return ())

    go :: Word -> IO ()
go Word
total = do
        Maybe (Handle, TBQueue Command, ByteString)
res <- STM (Maybe (Handle, TBQueue Command, ByteString))
-> IO (Maybe (Handle, TBQueue Command, ByteString))
forall a. STM a -> IO a
atomically (STM (Maybe (Handle, TBQueue Command, ByteString))
 -> IO (Maybe (Handle, TBQueue Command, ByteString)))
-> STM (Maybe (Handle, TBQueue Command, ByteString))
-> IO (Maybe (Handle, TBQueue Command, ByteString))
forall a b. (a -> b) -> a -> b
$ do
            State
s <- TVar State -> STM State
forall a. TVar a -> STM a
readTVar TVar State
ts
            case State
s of
                State
Closed -> Maybe (Handle, TBQueue Command, ByteString)
-> STM (Maybe (Handle, TBQueue Command, ByteString))
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe (Handle, TBQueue Command, ByteString)
forall a. Maybe a
Nothing
                Running Handle
handle TBQueue Command
queue' -> do
                    Command
cmd <- TBQueue Command -> STM Command
forall a. TBQueue a -> STM a
readTBQueue TBQueue Command
queue'
                    case Command
cmd of
                        Command
Close -> Maybe (Handle, TBQueue Command, ByteString)
-> STM (Maybe (Handle, TBQueue Command, ByteString))
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe (Handle, TBQueue Command, ByteString)
forall a. Maybe a
Nothing
                        AddChunk ByteString
bs -> Maybe (Handle, TBQueue Command, ByteString)
-> STM (Maybe (Handle, TBQueue Command, ByteString))
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe (Handle, TBQueue Command, ByteString)
 -> STM (Maybe (Handle, TBQueue Command, ByteString)))
-> Maybe (Handle, TBQueue Command, ByteString)
-> STM (Maybe (Handle, TBQueue Command, ByteString))
forall a b. (a -> b) -> a -> b
$! (Handle, TBQueue Command, ByteString)
-> Maybe (Handle, TBQueue Command, ByteString)
forall a. a -> Maybe a
Just (Handle
handle, TBQueue Command
queue', ByteString
bs)
        case Maybe (Handle, TBQueue Command, ByteString)
res of
            Maybe (Handle, TBQueue Command, ByteString)
Nothing -> () -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
            Just (Handle
handle, TBQueue Command
queue', ByteString
bs) -> do
                let total' :: Word
total' = Word
total Word -> Word -> Word
forall a. Num a => a -> a -> a
+ Int -> Word
forall a b. (Integral a, Num b) => a -> b
fromIntegral (ByteString -> Int
S.length ByteString
bs)
                Handle -> ByteString -> IO ()
S.hPut Handle
handle ByteString
bs
                Handle -> IO ()
SIO.hFlush Handle
handle
                if Word
total' Word -> Word -> Bool
forall a. Ord a => a -> a -> Bool
> Word
maxTotal
                    then do
                        IO Handle -> (Handle -> IO ()) -> (Handle -> IO ()) -> IO ()
forall a b c. IO a -> (a -> IO b) -> (a -> IO c) -> IO c
bracket
                            (Handle -> IO ()
SIO.hClose Handle
handle IO () -> IO Handle -> IO Handle
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> FilePath -> IO Handle
moveCurrent FilePath
dir)
                            (\Handle
handle' -> STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ TVar State -> State -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar TVar State
ts (State -> STM ()) -> State -> STM ()
forall a b. (a -> b) -> a -> b
$ Handle -> TBQueue Command -> State
Running Handle
handle' TBQueue Command
queue')
                            (IO () -> Handle -> IO ()
forall a b. a -> b -> a
const (IO () -> Handle -> IO ()) -> IO () -> Handle -> IO ()
forall a b. (a -> b) -> a -> b
$ () -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return ())
                        Word -> IO ()
go Word
0
                    else Word -> IO ()
go Word
total'

defaultMaxTotal :: Word
defaultMaxTotal :: Word
defaultMaxTotal = Word
5 Word -> Word -> Word
forall a. Num a => a -> a -> a
* Word
1024 Word -> Word -> Word
forall a. Num a => a -> a -> a
* Word
1024 -- 5 MB