{-# LANGUAGE ScopedTypeVariables, PatternGuards, RecordWildCards, FlexibleInstances, MultiParamTypeClasses #-}
module Development.Shake.Storage(
withStorage
) where
import General.Binary
import Development.Shake.Types
import General.Timing
import General.FileLock
import Data.Tuple.Extra
import Control.Exception.Extra
import Control.Monad.Extra
import Control.Concurrent.Extra
import Data.Binary.Get
import Data.Binary.Put
import Data.Time
import Data.Char
import Development.Shake.Classes
import qualified Data.HashMap.Strict as Map
import Data.List
import Numeric
import System.Directory
import System.Exit
import System.FilePath
import System.IO
import qualified Data.ByteString.Lazy.Char8 as LBS
import qualified Data.ByteString.Lazy as LBS8
type Map = Map.HashMap
databaseVersion :: String -> String
databaseVersion x = "SHAKE-DATABASE-11-" ++ s ++ "\r\n"
where s = tail $ init $ show x
splitVersion :: LBS.ByteString -> (LBS.ByteString, LBS.ByteString)
splitVersion abc = (a `LBS.append` b, c)
where (a,bc) = LBS.break (== '\r') abc
(b,c) = LBS.splitAt 2 bc
withStorage
:: (Show k, Show v, Eq w, Eq k, Hashable k
,Binary w, BinaryWith w k, BinaryWith w v)
=> ShakeOptions
-> (String -> IO ())
-> w
-> (Map k v -> (k -> v -> IO ()) -> IO a)
-> IO a
withStorage ShakeOptions{..} diagnostic witness act = do
diagnostic $ "Before fileLock on " ++ shakeFiles </> ".shake.lock"
withLockFile (shakeFiles </> ".shake.lock") $ do
diagnostic "After fileLock"
let dbfile = shakeFiles </> ".shake.database"
bupfile = shakeFiles </> ".shake.backup"
createDirectoryIfMissing True shakeFiles
b <- doesFileExist bupfile
when b $ do
unexpected "Backup file exists, restoring over the previous file\n"
diagnostic "Backup file move to original"
ignore $ removeFile dbfile
renameFile bupfile dbfile
addTiming "Database read"
withBinaryFile dbfile ReadWriteMode $ \h -> do
n <- hFileSize h
diagnostic $ "Reading file of size " ++ show n
(oldVer,src) <- fmap splitVersion $ LBS.hGet h $ fromInteger n
verEqual <- evaluate $ ver == oldVer
if not verEqual && not shakeVersionIgnore then do
unless (n == 0) $ do
let limit x = let (a,b) = splitAt 200 x in a ++ (if null b then "" else "...")
let disp = map (\x -> if isPrint x && isAscii x then x else '?') . takeWhile (`notElem` "\r\n")
outputErr $ unlines
["Error when reading Shake database - invalid version stamp detected:"
," File: " ++ dbfile
," Expected: " ++ disp (LBS.unpack ver)
," Found: " ++ disp (limit $ LBS.unpack oldVer)
,"All rules will be rebuilt"]
continue h Map.empty
else
join $ handleBool (not . asyncException) (\err -> do
msg <- showException err
outputErr $ unlines $
("Error when reading Shake database " ++ dbfile) :
map (" "++) (lines msg) ++
["All files will be rebuilt"]
when shakeStorageLog $ do
hSeek h AbsoluteSeek 0
i <- hFileSize h
bs <- LBS.hGet h $ fromInteger i
let cor = shakeFiles </> ".shake.corrupt"
LBS.writeFile cor bs
unexpected $ "Backup of corrupted file stored at " ++ cor ++ ", " ++ show i ++ " bytes\n"
return $ continue h Map.empty) $
case readChunks src of
([], slop) -> do
when (LBS.length slop > 0) $ unexpected $ "Last " ++ show slop ++ " bytes do not form a whole record\n"
diagnostic $ "Read 0 chunks, plus " ++ show slop ++ " slop"
return $ continue h Map.empty
(w:xs, slopRaw) -> do
let slop = fromIntegral $ LBS.length slopRaw
when (slop > 0) $ unexpected $ "Last " ++ show slop ++ " bytes do not form a whole record\n"
diagnostic $ "Read " ++ show (length xs + 1) ++ " chunks, plus " ++ show slop ++ " slop"
let ws = decode w
f mp (k, v) = Map.insert k v mp
ents = map (runGet $ getWith ws) xs
mp = foldl' f Map.empty ents
when (shakeVerbosity == Diagnostic) $ do
let raw x = "[len " ++ show (LBS.length x) ++ "] " ++ concat
[['0' | length c == 1] ++ c | x <- LBS8.unpack x, let c = showHex x ""]
let pretty (Left x) = "FAILURE: " ++ show x
pretty (Right x) = x
diagnostic $ "Witnesses " ++ raw w
forM_ (zip3 [1..] xs ents) $ \(i,x,ent) -> do
x2 <- try_ $ evaluate $ let s = show ent in rnf s `seq` s
diagnostic $ "Chunk " ++ show i ++ " " ++ raw x ++ " " ++ pretty x2
diagnostic $ "Slop " ++ raw slopRaw
diagnostic $ "Found " ++ show (Map.size mp) ++ " real entries"
if verEqual && (Map.null mp || (ws == witness && Map.size mp * 2 > length xs - 2)) then do
when (not (Map.null mp) && slop /= 0) $ do
diagnostic $ "Dropping last " ++ show slop ++ " bytes of database (incomplete)"
now <- hFileSize h
hSetFileSize h $ now - slop
hSeek h AbsoluteSeek $ now - slop
hFlush h
diagnostic "Drop complete"
return $ continue h mp
else do
addTiming "Database compression"
unexpected "Compressing database\n"
diagnostic "Compressing database"
hClose h
return $ do
renameFile dbfile bupfile
withBinaryFile dbfile ReadWriteMode $ \h -> do
reset h mp
removeFile bupfile
diagnostic "Compression complete"
continue h mp
where
unexpected x = when shakeStorageLog $ do
t <- getCurrentTime
appendFile (shakeFiles </> ".shake.storage.log") $ "\n[" ++ show t ++ "]: " ++ x
outputErr x = do
when (shakeVerbosity >= Quiet) $ shakeOutput Quiet x
unexpected x
ver = LBS.pack $ databaseVersion shakeVersion
writeChunk h s = do
diagnostic $ "Writing chunk " ++ show (LBS.length s)
LBS.hPut h $ toChunk s
reset h mp = do
diagnostic $ "Resetting database to " ++ show (Map.size mp) ++ " elements"
hSetFileSize h 0
hSeek h AbsoluteSeek 0
LBS.hPut h ver
writeChunk h $ encode witness
mapM_ (writeChunk h . runPut . putWith witness) $ Map.toList mp
hFlush h
diagnostic "Flush"
continue h mp = do
when (Map.null mp) $
reset h mp
flushThread shakeFlush h $ \out -> do
addTiming "With database"
act mp $ \k v -> out $ toChunk $ runPut $ putWith witness (k, v)
flushThread :: Maybe Double -> Handle -> ((LBS.ByteString -> IO ()) -> IO a) -> IO a
flushThread flush h act = do
chan <- newChan
kick <- newEmptyMVar
died <- newBarrier
flusher <- case flush 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 <- forkIO $ handle_ (\e -> signalBarrier died () >> throwTo root e) $
whileM $ join $ readChan chan
(act $ \s -> do
evaluate $ LBS.length s
writeChan chan $ LBS.hPut h s >> tryPutMVar kick () >> return True)
`finally` do
maybe (return ()) killThread flusher
writeChan chan $ signalBarrier died () >> return False
waitBarrier died
readChunks :: LBS.ByteString -> ([LBS.ByteString], LBS.ByteString)
readChunks x
| Just (n, x) <- grab 4 x
, Just (y, x) <- grab (fromIntegral (decode n :: Word32)) x
= first (y :) $ readChunks x
| otherwise = ([], x)
where
grab i x | LBS.length a == i = Just (a, b)
| otherwise = Nothing
where (a,b) = LBS.splitAt i x
toChunk :: LBS.ByteString -> LBS.ByteString
toChunk x = n `LBS.append` x
where n = encode (fromIntegral $ LBS.length x :: Word32)
asyncException :: SomeException -> Bool
asyncException e
| Just (_ :: AsyncException) <- fromException e = True
| Just (_ :: ExitCode) <- fromException e = True
| otherwise = False