Persistence.hs

{-# LANGUAGE NamedFieldPuns #-}
{- |
  This module contains some basis persistence strategies useful for
  testing, or getting started.
-}
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


{- |
  A convenient memory-based persistence layer. Good for testing or for
  applications (like caches) that don't have durability requirements.
-}
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))


{- |
  A convenient way to persist partition states to disk. This persistence
  implementation is very simplistic, and may not be suitable for all uses.
-}
diskPersistence :: (Binary e, Binary s)
  => FilePath
     {- ^ The directory under which partition states will be stored. -}
  -> (Peer -> ClusterPowerState -> IO ())
     {- ^
       A user-supplied mechanism to for saving the cluster state, for
       recovery purposes.
     -}
  -> 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 -- . filter notSys
        -- notSys = not . (`elem` [".", ".."])

    {- |
      Convert a key to a path
    -}
    toPath :: PartitionKey -> FilePath
    toPath key = directory ++ "/" ++ toHex key


{- | Convert a `PartitionKey` into a hex string. -}
toHex :: PartitionKey -> String
toHex (K (Word256 (Word128 a b) (Word128 c d))) =
  concatMap toHex64 [a, b, c, d]


{- |
  Convert a `Word64` into a hex string.

  I know I'm going to hell for this, but I just can't abide the
  @hexstring@ package pulling @aeson@ into our dependency tree.
-}
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'


{- | Maybe convert a hex string into a partition key -}
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)