module Metro.Utils
( getEpochTime
, setupLog
, recvEnough
) where
import Control.Monad (when)
import Data.ByteString (ByteString)
import qualified Data.ByteString as B (concat, drop, empty, length,
null, take)
import Data.Int (Int64)
import Data.UnixTime (getUnixTime, toEpochTime)
import Foreign.C.Types (CTime (..))
import Metro.Class (Transport (..), TransportError (..))
import System.IO (stderr)
import System.Log.Formatter (simpleLogFormatter)
import System.Log.Handler (setFormatter)
import System.Log.Handler.Simple (streamHandler)
import System.Log.Logger
import UnliftIO (MonadIO (..), TVar, atomically,
readTVar, throwIO, writeTVar)
getEpochTime :: MonadIO m => m Int64
getEpochTime :: m Int64
getEpochTime = IO Int64 -> m Int64
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO Int64 -> m Int64) -> IO Int64 -> m Int64
forall a b. (a -> b) -> a -> b
$ CTime -> Int64
un (CTime -> Int64) -> (UnixTime -> CTime) -> UnixTime -> Int64
forall b c a. (b -> c) -> (a -> b) -> a -> c
. UnixTime -> CTime
toEpochTime (UnixTime -> Int64) -> IO UnixTime -> IO Int64
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IO UnixTime
getUnixTime
where un :: CTime -> Int64
un :: CTime -> Int64
un (CTime Int64
t) = Int64
t
setupLog :: Priority -> IO ()
setupLog :: Priority -> IO ()
setupLog Priority
logLevel = do
IO ()
removeAllHandlers
GenericHandler Handle
handle <- Handle -> Priority -> IO (GenericHandler Handle)
streamHandler Handle
stderr Priority
logLevel IO (GenericHandler Handle)
-> (GenericHandler Handle -> IO (GenericHandler Handle))
-> IO (GenericHandler Handle)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \GenericHandler Handle
lh -> GenericHandler Handle -> IO (GenericHandler Handle)
forall (m :: * -> *) a. Monad m => a -> m a
return (GenericHandler Handle -> IO (GenericHandler Handle))
-> GenericHandler Handle -> IO (GenericHandler Handle)
forall a b. (a -> b) -> a -> b
$
GenericHandler Handle
-> LogFormatter (GenericHandler Handle) -> GenericHandler Handle
forall a. LogHandler a => a -> LogFormatter a -> a
setFormatter GenericHandler Handle
lh (String -> LogFormatter (GenericHandler Handle)
forall a. String -> LogFormatter a
simpleLogFormatter String
"[$time : $loggername : $prio] $msg")
String -> (Logger -> Logger) -> IO ()
updateGlobalLogger String
rootLoggerName (GenericHandler Handle -> Logger -> Logger
forall a. LogHandler a => a -> Logger -> Logger
addHandler GenericHandler Handle
handle (Logger -> Logger) -> (Logger -> Logger) -> Logger -> Logger
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Priority -> Logger -> Logger
setLevel Priority
logLevel)
recvEnough :: (MonadIO m, Transport tp) => TVar ByteString -> tp -> Int -> m ByteString
recvEnough :: TVar ByteString -> tp -> Int -> m ByteString
recvEnough TVar ByteString
buffer tp
tp Int
nbytes = do
ByteString
buf <- STM ByteString -> m ByteString
forall (m :: * -> *) a. MonadIO m => STM a -> m a
atomically (STM ByteString -> m ByteString) -> STM ByteString -> m ByteString
forall a b. (a -> b) -> a -> b
$ do
ByteString
bf <- TVar ByteString -> STM ByteString
forall a. TVar a -> STM a
readTVar TVar ByteString
buffer
TVar ByteString -> ByteString -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar TVar ByteString
buffer (ByteString -> STM ()) -> ByteString -> STM ()
forall a b. (a -> b) -> a -> b
$! Int -> ByteString -> ByteString
B.drop Int
nbytes ByteString
bf
ByteString -> STM ByteString
forall (m :: * -> *) a. Monad m => a -> m a
return (ByteString -> STM ByteString) -> ByteString -> STM ByteString
forall a b. (a -> b) -> a -> b
$! Int -> ByteString -> ByteString
B.take Int
nbytes ByteString
bf
if ByteString -> Int
B.length ByteString
buf Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
nbytes then ByteString -> m ByteString
forall (m :: * -> *) a. Monad m => a -> m a
return ByteString
buf
else do
ByteString
otherBuf <- IO ByteString -> m ByteString
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO ByteString -> m ByteString) -> IO ByteString -> m ByteString
forall a b. (a -> b) -> a -> b
$ Int -> IO ByteString
readBuf (Int
nbytes Int -> Int -> Int
forall a. Num a => a -> a -> a
- ByteString -> Int
B.length ByteString
buf)
let out :: ByteString
out = [ByteString] -> ByteString
B.concat [ ByteString
buf, ByteString
otherBuf ]
STM () -> m ()
forall (m :: * -> *) a. MonadIO m => STM a -> m a
atomically (STM () -> m ()) -> (ByteString -> STM ()) -> ByteString -> m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TVar ByteString -> ByteString -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar TVar ByteString
buffer (ByteString -> m ()) -> ByteString -> m ()
forall a b. (a -> b) -> a -> b
$! Int -> ByteString -> ByteString
B.drop Int
nbytes ByteString
out
ByteString -> m ByteString
forall (m :: * -> *) a. Monad m => a -> m a
return (ByteString -> m ByteString) -> ByteString -> m ByteString
forall a b. (a -> b) -> a -> b
$! Int -> ByteString -> ByteString
B.take Int
nbytes ByteString
out
where readBuf :: Int -> IO ByteString
readBuf :: Int -> IO ByteString
readBuf Int
0 = ByteString -> IO ByteString
forall (m :: * -> *) a. Monad m => a -> m a
return ByteString
B.empty
readBuf Int
nb = do
ByteString
buf <- tp -> Int -> IO ByteString
forall transport.
Transport transport =>
transport -> Int -> IO ByteString
recvData tp
tp (Int -> IO ByteString) -> Int -> IO ByteString
forall a b. (a -> b) -> a -> b
$ Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
4096 Int
nb
Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (ByteString -> Bool
B.null ByteString
buf) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ TransportError -> IO ()
forall (m :: * -> *) e a. (MonadIO m, Exception e) => e -> m a
throwIO TransportError
TransportClosed
if ByteString -> Int
B.length ByteString
buf Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
nb then ByteString -> IO ByteString
forall (m :: * -> *) a. Monad m => a -> m a
return ByteString
buf
else do
ByteString
otherBuf <- Int -> IO ByteString
readBuf (Int
nb Int -> Int -> Int
forall a. Num a => a -> a -> a
- ByteString -> Int
B.length ByteString
buf)
ByteString -> IO ByteString
forall (m :: * -> *) a. Monad m => a -> m a
return (ByteString -> IO ByteString) -> ByteString -> IO ByteString
forall a b. (a -> b) -> a -> b
$! [ByteString] -> ByteString
B.concat [ ByteString
buf, ByteString
otherBuf ]