{-# LANGUAGE RecordWildCards, ScopedTypeVariables #-}

module General.Chunks(
    Chunks,
    readChunk, readChunkMax, usingWriteChunks, writeChunk,
    restoreChunksBackup, usingChunks, resetChunksCompact, resetChunksCorrupt
    ) where

import System.Time.Extra
import System.FilePath
import Control.Concurrent.Extra
import Control.Monad.Extra
import Control.Exception
import System.IO
import System.Directory
import qualified Data.ByteString as BS
import Data.Word
import Data.Monoid
import General.Binary
import General.Extra
import General.Cleanup
import General.Thread
import Prelude


data Chunks = Chunks
    {Chunks -> FilePath
chunksFileName :: FilePath
    ,Chunks -> Maybe Seconds
chunksFlush :: Maybe Seconds
    ,Chunks -> MVar Handle
chunksHandle :: MVar Handle
    }


---------------------------------------------------------------------
-- READ/WRITE OPERATIONS

readChunk :: Chunks -> IO (Either BS.ByteString BS.ByteString)
readChunk :: Chunks -> IO (Either ByteString ByteString)
readChunk Chunks
c = Chunks -> Word32 -> IO (Either ByteString ByteString)
readChunkMax Chunks
c Word32
forall a. Bounded a => a
maxBound

-- | Return either a valid chunk (Right), or a trailing suffix with no information (Left)
readChunkMax :: Chunks -> Word32 -> IO (Either BS.ByteString BS.ByteString)
readChunkMax :: Chunks -> Word32 -> IO (Either ByteString ByteString)
readChunkMax Chunks{FilePath
Maybe Seconds
MVar Handle
chunksFileName :: Chunks -> FilePath
chunksFlush :: Chunks -> Maybe Seconds
chunksHandle :: Chunks -> MVar Handle
chunksFileName :: FilePath
chunksFlush :: Maybe Seconds
chunksHandle :: MVar Handle
..} Word32
mx = MVar Handle
-> (Handle -> IO (Either ByteString ByteString))
-> IO (Either ByteString ByteString)
forall a b. MVar a -> (a -> IO b) -> IO b
withMVar MVar Handle
chunksHandle ((Handle -> IO (Either ByteString ByteString))
 -> IO (Either ByteString ByteString))
-> (Handle -> IO (Either ByteString ByteString))
-> IO (Either ByteString ByteString)
forall a b. (a -> b) -> a -> b
$ \Handle
h -> Handle -> Word32 -> IO (Either ByteString ByteString)
readChunkDirect Handle
h Word32
mx

readChunkDirect :: Handle -> Word32 -> IO (Either BS.ByteString BS.ByteString)
readChunkDirect :: Handle -> Word32 -> IO (Either ByteString ByteString)
readChunkDirect Handle
h Word32
mx = do
    let slop :: ByteString -> IO (Either ByteString b)
slop ByteString
x = do
            Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (ByteString -> Bool
BS.null ByteString
x) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ Handle -> Integer -> IO ()
hSetFileSize Handle
h (Integer -> IO ()) -> (Integer -> Integer) -> Integer -> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
subtract (Int -> Integer
forall a. Integral a => a -> Integer
toInteger (Int -> Integer) -> Int -> Integer
forall a b. (a -> b) -> a -> b
$ ByteString -> Int
BS.length ByteString
x) (Integer -> IO ()) -> IO Integer -> IO ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Handle -> IO Integer
hFileSize Handle
h
            Either ByteString b -> IO (Either ByteString b)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Either ByteString b -> IO (Either ByteString b))
-> Either ByteString b -> IO (Either ByteString b)
forall a b. (a -> b) -> a -> b
$ ByteString -> Either ByteString b
forall a b. a -> Either a b
Left ByteString
x
    n <- Handle -> Int -> IO ByteString
BS.hGet Handle
h Int
4
    if BS.length n < 4 then slop n else do
        let count = Word32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word32 -> Int) -> Word32 -> Int
forall a b. (a -> b) -> a -> b
$ Word32 -> Word32 -> Word32
forall a. Ord a => a -> a -> a
min Word32
mx (Word32 -> Word32) -> Word32 -> Word32
forall a b. (a -> b) -> a -> b
$ (Word32, ByteString) -> Word32
forall a b. (a, b) -> a
fst ((Word32, ByteString) -> Word32) -> (Word32, ByteString) -> Word32
forall a b. (a -> b) -> a -> b
$ ByteString -> (Word32, ByteString)
forall a. Storable a => ByteString -> (a, ByteString)
unsafeBinarySplit ByteString
n
        v <- BS.hGet h count
        if BS.length v < count then slop (n `BS.append` v) else pure $ Right v

writeChunkDirect :: Handle -> Builder -> IO ()
writeChunkDirect :: Handle -> Builder -> IO ()
writeChunkDirect Handle
h Builder
x = ByteString
bs ByteString -> IO () -> IO ()
forall a b. a -> b -> b
`seq` Handle -> ByteString -> IO ()
BS.hPut Handle
h ByteString
bs
    where bs :: ByteString
bs = Builder -> ByteString
runBuilder (Builder -> ByteString) -> Builder -> ByteString
forall a b. (a -> b) -> a -> b
$ Word32 -> Builder
forall a. BinaryEx a => a -> Builder
putEx (Int -> Word32
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Word32) -> Int -> Word32
forall a b. (a -> b) -> a -> b
$ Builder -> Int
sizeBuilder Builder
x :: Word32) Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> Builder
x


-- | If 'writeChunks' and any of the reopen operations are interleaved it will cause issues.
usingWriteChunks :: Cleanup -> Chunks -> IO (Builder -> IO ())
-- We avoid calling flush too often on SSD drives, as that can be slow
-- Make sure all exceptions happen on the caller, so we don't have to move exceptions back
-- Make sure we only write on one thread, otherwise async exceptions can cause partial writes
usingWriteChunks :: Cleanup -> Chunks -> IO (Builder -> IO ())
usingWriteChunks Cleanup
cleanup Chunks{FilePath
Maybe Seconds
MVar Handle
chunksFileName :: Chunks -> FilePath
chunksFlush :: Chunks -> Maybe Seconds
chunksHandle :: Chunks -> MVar Handle
chunksFileName :: FilePath
chunksFlush :: Maybe Seconds
chunksHandle :: MVar Handle
..} = do
    h <- Cleanup -> IO Handle -> (Handle -> IO ()) -> IO Handle
forall a. Cleanup -> IO a -> (a -> IO ()) -> IO a
allocate Cleanup
cleanup (MVar Handle -> IO Handle
forall a. MVar a -> IO a
takeMVar MVar Handle
chunksHandle) (MVar Handle -> Handle -> IO ()
forall a. MVar a -> a -> IO ()
putMVar MVar Handle
chunksHandle)
    chan <- newChan -- operations to perform on the file
    kick <- newEmptyMVar -- kicked whenever something is written
    died <- newBarrier -- has the writing thread finished

    whenJust chunksFlush $ \Seconds
flush ->
        Cleanup -> IO () -> IO ()
allocateThread Cleanup
cleanup (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ IO () -> IO ()
forall (f :: * -> *) a b. Applicative f => f a -> f b
forever (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
            MVar () -> IO ()
forall a. MVar a -> IO a
takeMVar MVar ()
kick
            Seconds -> IO ()
sleep Seconds
flush
            MVar () -> IO (Maybe ())
forall a. MVar a -> IO (Maybe a)
tryTakeMVar MVar ()
kick
            Chan (IO Bool) -> IO Bool -> IO ()
forall a. Chan a -> a -> IO ()
writeChan Chan (IO Bool)
chan (IO Bool -> IO ()) -> IO Bool -> IO ()
forall a b. (a -> b) -> a -> b
$ Handle -> IO ()
hFlush Handle
h IO () -> IO Bool -> IO Bool
forall a b. IO a -> IO b -> IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Bool -> IO Bool
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
True

    -- pump the thread while we are running
    -- once we abort, let everything finish flushing first
    -- the mask_ is very important - we don't want to abort until everything finishes
    allocateThread cleanup $ mask_ $ whileM $ join $ readChan chan
    -- this cleanup will run before we attempt to kill the thread
    register cleanup $ writeChan chan $ pure False

    pure $ \Builder
s -> do
        out <- IO () -> IO (IO ())
forall a. a -> IO a
evaluate (IO () -> IO (IO ())) -> IO () -> IO (IO ())
forall a b. (a -> b) -> a -> b
$ Handle -> Builder -> IO ()
writeChunkDirect Handle
h Builder
s -- ensure exceptions occur on this thread
        writeChan chan $ out >> tryPutMVar kick () >> pure True


writeChunk :: Chunks -> Builder -> IO ()
writeChunk :: Chunks -> Builder -> IO ()
writeChunk Chunks{FilePath
Maybe Seconds
MVar Handle
chunksFileName :: Chunks -> FilePath
chunksFlush :: Chunks -> Maybe Seconds
chunksHandle :: Chunks -> MVar Handle
chunksFileName :: FilePath
chunksFlush :: Maybe Seconds
chunksHandle :: MVar Handle
..} Builder
x = MVar Handle -> (Handle -> IO ()) -> IO ()
forall a b. MVar a -> (a -> IO b) -> IO b
withMVar MVar Handle
chunksHandle ((Handle -> IO ()) -> IO ()) -> (Handle -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Handle
h -> Handle -> Builder -> IO ()
writeChunkDirect Handle
h Builder
x


---------------------------------------------------------------------
-- FILENAME OPERATIONS

backup :: FilePath -> FilePath
backup FilePath
x = FilePath
x FilePath -> FilePath -> FilePath
<.> FilePath
"backup"

restoreChunksBackup :: FilePath -> IO Bool
restoreChunksBackup :: FilePath -> IO Bool
restoreChunksBackup FilePath
file = do
    -- complete a partially failed compress
    b <- FilePath -> IO Bool
doesFileExist (FilePath -> IO Bool) -> FilePath -> IO Bool
forall a b. (a -> b) -> a -> b
$ FilePath -> FilePath
backup FilePath
file
    if not b then pure False else do
        removeFile_ file
        renameFile (backup file) file
        pure True


usingChunks :: Cleanup -> FilePath -> Maybe Seconds -> IO Chunks
usingChunks :: Cleanup -> FilePath -> Maybe Seconds -> IO Chunks
usingChunks Cleanup
cleanup FilePath
file Maybe Seconds
flush = do
    h <- IO (MVar Handle)
forall a. IO (MVar a)
newEmptyMVar
    allocate cleanup
        (putMVar h =<< openFile file ReadWriteMode)
        (const $ hClose =<< takeMVar h)
    pure $ Chunks file flush h


-- | The file is being compacted, if the process fails, use a backup.
resetChunksCompact :: Chunks -> ((Builder -> IO ()) -> IO a) -> IO a
resetChunksCompact :: forall a. Chunks -> ((Builder -> IO ()) -> IO a) -> IO a
resetChunksCompact Chunks{FilePath
Maybe Seconds
MVar Handle
chunksFileName :: Chunks -> FilePath
chunksFlush :: Chunks -> Maybe Seconds
chunksHandle :: Chunks -> MVar Handle
chunksFileName :: FilePath
chunksFlush :: Maybe Seconds
chunksHandle :: MVar Handle
..} (Builder -> IO ()) -> IO a
act = ((forall a. IO a -> IO a) -> IO a) -> IO a
forall b. ((forall a. IO a -> IO a) -> IO b) -> IO b
mask (((forall a. IO a -> IO a) -> IO a) -> IO a)
-> ((forall a. IO a -> IO a) -> IO a) -> IO a
forall a b. (a -> b) -> a -> b
$ \forall a. IO a -> IO a
restore -> do
    h <- MVar Handle -> IO Handle
forall a. MVar a -> IO a
takeMVar MVar Handle
chunksHandle
    flip onException (putMVar chunksHandle h) $ restore $ do
        hClose h
        copyFile chunksFileName $ backup chunksFileName
    h <- openFile chunksFileName ReadWriteMode
    flip finally (putMVar chunksHandle h) $ restore $ do
        hSetFileSize h 0
        hSeek h AbsoluteSeek 0
        res <- act $ writeChunkDirect h
        hFlush h
        removeFile $ backup chunksFileName
        pure res


-- | The file got corrupted, return a new version.
resetChunksCorrupt :: Maybe FilePath -> Chunks -> IO ()
resetChunksCorrupt :: Maybe FilePath -> Chunks -> IO ()
resetChunksCorrupt Maybe FilePath
copy Chunks{FilePath
Maybe Seconds
MVar Handle
chunksFileName :: Chunks -> FilePath
chunksFlush :: Chunks -> Maybe Seconds
chunksHandle :: Chunks -> MVar Handle
chunksFileName :: FilePath
chunksFlush :: Maybe Seconds
chunksHandle :: MVar Handle
..} = ((forall a. IO a -> IO a) -> IO ()) -> IO ()
forall b. ((forall a. IO a -> IO a) -> IO b) -> IO b
mask (((forall a. IO a -> IO a) -> IO ()) -> IO ())
-> ((forall a. IO a -> IO a) -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \forall a. IO a -> IO a
restore -> do
    h <- MVar Handle -> IO Handle
forall a. MVar a -> IO a
takeMVar MVar Handle
chunksHandle
    h <- case copy of
        Maybe FilePath
Nothing -> Handle -> IO Handle
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Handle
h
        Just FilePath
copy -> do
            (IO () -> IO () -> IO ()) -> IO () -> IO () -> IO ()
forall a b c. (a -> b -> c) -> b -> a -> c
flip IO () -> IO () -> IO ()
forall a b. IO a -> IO b -> IO a
onException (MVar Handle -> Handle -> IO ()
forall a. MVar a -> a -> IO ()
putMVar MVar Handle
chunksHandle Handle
h) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ IO () -> IO ()
forall a. IO a -> IO a
restore (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
                Handle -> IO ()
hClose Handle
h
                FilePath -> FilePath -> IO ()
copyFile FilePath
chunksFileName FilePath
copy
            FilePath -> IOMode -> IO Handle
openFile FilePath
chunksFileName IOMode
ReadWriteMode
    flip finally (putMVar chunksHandle h) $ do
        hSetFileSize h 0
        hSeek h AbsoluteSeek 0