module General.Chunks(
Chunks,
readChunk, readChunkMax, writeChunks, writeChunk,
restoreChunksBackup, withChunks, resetChunksCompact, resetChunksCorrupt
) where
import System.Time.Extra
import System.FilePath
import Control.Concurrent.Extra
import Control.Monad.Extra
import Control.Exception.Extra
import System.IO
import System.Directory
import qualified Data.ByteString as BS
import Data.Word
import Data.Monoid
import General.Binary
data Chunks = Chunks
{chunksFileName :: FilePath
,chunksFlush :: Maybe Seconds
,chunksHandle :: MVar Handle
}
readChunk :: Chunks -> IO (Either BS.ByteString BS.ByteString)
readChunk c = readChunkMax c maxBound
readChunkMax :: Chunks -> Word32 -> IO (Either BS.ByteString BS.ByteString)
readChunkMax Chunks{..} mx = withMVar chunksHandle $ \h -> do
let slop x = do
unless (BS.null x) $ hSetFileSize h . subtract (toInteger $ BS.length x) =<< hFileSize h
return $ Left x
n <- BS.hGet h 4
if BS.length n < 4 then slop n else do
let count = fromIntegral $ min mx $ fst $ unsafeBinarySplit n
v <- BS.hGet h count
if BS.length v < count then slop (n `BS.append` v) else return $ Right v
writeChunkDirect :: Handle -> Builder -> IO ()
writeChunkDirect h x = bs `seq` BS.hPut h bs
where bs = runBuilder $ putEx (fromIntegral $ sizeBuilder x :: Word32) <> x
writeChunks :: Chunks -> ((Builder -> IO ()) -> IO a) -> IO a
writeChunks Chunks{..} act = withMVar chunksHandle $ \h -> do
chan <- newChan
kick <- newEmptyMVar
died <- newBarrier
flusher <- case chunksFlush of
Nothing -> return Nothing
Just flush -> fmap Just $ forkIO $ forever $ do
takeMVar kick
threadDelay $ ceiling $ flush * 1000000
tryTakeMVar kick
writeChan chan $ hFlush h >> return True
root <- myThreadId
writer <- flip forkFinally (\e -> do signalBarrier died (); either (throwTo root) (const $ return ()) e) $
whileM $ join $ readChan chan
(act $ \s -> do
out <- evaluate $ writeChunkDirect h s
writeChan chan $ out >> tryPutMVar kick () >> return True)
`finally` do
maybe (return ()) killThread flusher
writeChan chan $ return False
waitBarrier died
writeChunk :: Chunks -> Builder -> IO ()
writeChunk Chunks{..} x = withMVar chunksHandle $ \h -> writeChunkDirect h x
backup x = x <.> "backup"
restoreChunksBackup :: FilePath -> IO Bool
restoreChunksBackup file = do
b <- doesFileExist $ backup file
if not b then return False else do
handle (\(_ :: IOError) -> return ()) $ removeFile file
renameFile (backup file) file
return True
withChunks :: FilePath -> Maybe Seconds -> (Chunks -> IO a) -> IO a
withChunks file flush act = do
h <- newEmptyMVar
bracket_
(putMVar h =<< openFile file ReadWriteMode)
(hClose =<< takeMVar h) $
act $ Chunks file flush h
resetChunksCompact :: Chunks -> ((Builder -> IO ()) -> IO a) -> IO a
resetChunksCompact Chunks{..} act = mask $ \restore -> do
h <- takeMVar 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
return res
resetChunksCorrupt :: Maybe FilePath -> Chunks -> IO ()
resetChunksCorrupt copy Chunks{..} = mask $ \restore -> do
h <- takeMVar chunksHandle
case copy of
Nothing -> return h
Just copy -> do
flip onException (putMVar chunksHandle h) $ restore $ do
hClose h
copyFile chunksFileName copy
openFile chunksFileName ReadWriteMode
flip finally (putMVar chunksHandle h) $ do
hSetFileSize h 0
hSeek h AbsoluteSeek 0