module Network.Legion.Persistence (
newMemoryPersistence,
diskPersistence,
) where
import Prelude hiding (lookup, readFile, writeFile)
import Control.Concurrent.STM (atomically, newTVar, modifyTVar', readTVar,
TVar)
import Control.Monad.Trans.Class (lift)
import Data.Attoparsec.ByteString (parseOnly, atEnd)
import Data.Attoparsec.ByteString.Char8 (hexadecimal)
import Data.Binary (Binary, encode, decode)
import Data.Bits (testBit)
import Data.Bool (bool)
import Data.ByteString (readFile, writeFile)
import Data.ByteString.Char8 (pack)
import Data.ByteString.Lazy (toStrict, fromStrict)
import Data.Conduit (Source, (=$=), awaitForever, yield)
import Data.Conduit.List (sourceList)
import Data.DoubleWord (Word256(Word256), Word128(Word128))
import Data.Either (rights)
import Data.Map (Map, insert, lookup)
import Data.Word (Word64)
import Network.Legion (Persistence(Persistence, getState, saveState,
list, saveCluster), PartitionPowerState, Peer, ClusterPowerState,
PartitionKey(K))
import System.Directory (removeFile, doesFileExist, getDirectoryContents)
import qualified Data.Map as Map
newMemoryPersistence :: IO (Persistence e o s)
newMemoryPersistence = do
cacheT <- atomically (newTVar Map.empty)
return Persistence {
saveCluster = const . const $ return (),
getState = fetchState cacheT,
saveState = saveState_ cacheT,
list = list_ cacheT
}
where
saveState_
:: TVar (Map PartitionKey (PartitionPowerState e o s))
-> PartitionKey
-> Maybe (PartitionPowerState e o s)
-> IO ()
saveState_ cacheT key (Just state) =
(atomically . modifyTVar' cacheT . insert key) state
saveState_ cacheT key Nothing =
(atomically . modifyTVar' cacheT . Map.delete) key
fetchState
:: TVar (Map PartitionKey (PartitionPowerState e o s))
-> PartitionKey
-> IO (Maybe (PartitionPowerState e o s))
fetchState cacheT key = atomically $
lookup key <$> readTVar cacheT
list_
:: TVar (Map PartitionKey (PartitionPowerState e o s))
-> Source IO (PartitionKey, PartitionPowerState e o s)
list_ cacheT =
sourceList . Map.toList =<< lift (atomically (readTVar cacheT))
diskPersistence :: (Binary e, Binary s)
=> FilePath
-> (Peer -> ClusterPowerState -> IO ())
-> Persistence e o s
diskPersistence directory saveCluster = Persistence {
saveCluster,
getState,
saveState,
list
}
where
getState :: (Binary e, Binary s)
=> PartitionKey
-> IO (Maybe (PartitionPowerState e o s))
getState key =
let path = toPath key in
doesFileExist path >>= bool
(return Nothing)
((Just . decode . fromStrict) <$> readFile path)
saveState :: (Binary e, Binary s)
=> PartitionKey
-> Maybe (PartitionPowerState e o s)
-> IO ()
saveState key (Just state) =
writeFile (toPath key) (toStrict (encode state))
saveState key Nothing =
let path = toPath key in
doesFileExist path >>= bool
(return ())
(removeFile path)
list :: (Binary e, Binary s)
=> Source IO (PartitionKey, PartitionPowerState e o s)
list = do
keys <- lift $ readHexList <$> getDirectoryContents directory
sourceList keys =$= fillData
where
fillData = awaitForever (\key -> do
let path = toPath key
state <- lift ((decode . fromStrict) <$> readFile path)
yield (key, state)
)
readHexList = rights . fmap fromHex
toPath :: PartitionKey -> FilePath
toPath key = directory ++ "/" ++ toHex key
toHex :: PartitionKey -> String
toHex (K (Word256 (Word128 a b) (Word128 c d))) =
concatMap toHex64 [a, b, c, d]
toHex64 :: Word64 -> String
toHex64 w = fmap (digit . quad) [15, 14..0]
where
quad :: Int -> (Int, Int, Int, Int)
quad n = let base = n * 4 in (base + 3, base + 2, base + 1, base)
digit :: (Int, Int, Int, Int) -> Char
digit (a, b, c, d) =
case (testBit w a, testBit w b, testBit w c, testBit w d) of
(False, False, False, False) -> '0'
(False, False, False, True) -> '1'
(False, False, True, False) -> '2'
(False, False, True, True) -> '3'
(False, True, False, False) -> '4'
(False, True, False, True) -> '5'
(False, True, True, False) -> '6'
(False, True, True, True) -> '7'
(True, False, False, False) -> '8'
(True, False, False, True) -> '9'
(True, False, True, False) -> 'a'
(True, False, True, True) -> 'b'
(True, True, False, False) -> 'c'
(True, True, False, True) -> 'd'
(True, True, True, False) -> 'e'
(True, True, True, True) -> 'f'
fromHex :: String -> Either String PartitionKey
fromHex str
| length str > 64 =
Left "trailing characters while parsing hex PartitionKey"
| otherwise =
K <$> parseOnly parser (pack str)
where
parser = do
w <- hexadecimal
atEnd >>= bool
(fail "not a valid hex string")
(return w)