Slurp.hs

   1 -- This file is part of Diohsc
   2 -- Copyright (C) 2020 Martin Bays <mbays@sdf.org>
   3 --
   4 -- This program is free software: you can redistribute it and/or modify
   5 -- it under the terms of version 3 of the GNU General Public License as
   6 -- published by the Free Software Foundation, or any later version.
   7 --
   8 -- You should have received a copy of the GNU General Public License
   9 -- along with this program.  If not, see http://www.gnu.org/licenses/.
  10
  11 module Slurp where
  12
  13 import           Control.Monad        (when)
  14 import           Data.Hourglass       (timeDiffP)
  15 import           System.IO
  16 import           System.IO.Unsafe     (unsafeInterleaveIO)
  17 import           Time.System          (timeCurrentP)
  18 import           Time.Types           (ElapsedP, NanoSeconds (..), Seconds (..))
  19
  20 import qualified Data.ByteString      as BS
  21 import qualified Data.ByteString.Lazy as BL
  22
  23 -- incorporate lazy IO into a bytestring such that it will print progress as
  24 -- it is forced, if this takes more than a second.
  25 interleaveProgress :: ElapsedP -> BL.ByteString -> IO BL.ByteString
  26 interleaveProgress t0 bs = do
  27     t1 <- timeCurrentP
  28     let slurp _ _ n' [] = do
  29             t' <- timeCurrentP
  30             when (t' `timeDiffP` t1 > (Seconds 1, 0)) $ do
  31                 hPutStrLn stderr $ "\r\ESC[KReceived: " ++ humanBytes n'
  32                     ++ "  " ++ humanRate n' (t' `timeDiffP` t0)
  33                 hFlush stderr
  34             return []
  35         slurp n t n' (c:cs) = unsafeInterleaveIO $ do
  36             let n'' = n' + fromIntegral (BS.length c)
  37             t' <- timeCurrentP
  38             if t' `timeDiffP` t > (Seconds 1, 0)
  39             then do
  40                 when (t' `timeDiffP` t1 > (Seconds 1, 0)) $ do
  41                     hPutStr stderr $   "\r\ESC[KProgress: " ++ humanBytes n'
  42                         ++ "  " ++ humanRate (n' - n) (t' `timeDiffP` t)
  43                     hFlush stderr
  44                 (c:) <$> slurp n'' t' n'' cs
  45             else (c:) <$> slurp n t n'' cs
  46         humanBytes n | n < 1024 = show n ++ "B"
  47         humanBytes n | n < 1024*1024 =
  48             let (n',p) = (`divMod` 10) $ (n*10) `div` 1024
  49             in show n' ++ "." ++ show p ++ "KB"
  50         humanBytes n =
  51             let (n',p) = (`divMod` 10) $ (n*10) `div` (1024*1024)
  52             in show n' ++ "." ++ show p ++ "MB"
  53         humanRate _ (0,0) = ""
  54         humanRate n (Seconds s, NanoSeconds ns) = humanBytes r ++ "/s" where
  55             r = (billion*n)`div`((billion*s)+ns)
  56             billion = 1000000000
  57     BL.fromChunks <$> slurp 0 t0 0 (BL.toChunks bs)
  58
  59 -- |force bs, printing progress
  60 slurpNoisily :: ElapsedP -> BL.ByteString -> IO ()
  61 slurpNoisily t0 bs = do
  62     bs' <- BL.toStrict <$> interleaveProgress t0 bs
  63     seq bs' $ return ()