1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
|
{-# LANGUAGE CPP #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilyDependencies #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE TypeSynonymInstances #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE DeriveAnyClass #-}
{-# LANGUAGE StrictData #-}
{-# OPTIONS_GHC -Wno-orphans #-}
-- | Read-only on-disk key-value store
module Store where
import Control.DeepSeq (NFData)
import Control.Monad.IO.Class (liftIO)
import qualified Control.Monad.State.Strict as S
import qualified Data.ByteString as BS
import qualified Data.ByteString.Short as BSS
import Data.Either (Either)
import qualified Data.Map.Strict as M
import Data.Serialize (Serialize, decode, encode, get, put)
import GHC.Generics (Generic)
import Prelude hiding (lookup)
import System.Directory (doesFileExist)
import System.FilePath ((</>))
import System.IO (Handle, IOMode(..), hTell, withFile)
import System.IO.MMap (mmapFileByteString)
data Store = Store
{ index :: M.Map BSS.ShortByteString Location
, values :: BS.ByteString
}
data Location = Location
{ offset :: Int
, length :: Int
} deriving (Show, Eq, Ord, Generic, NFData)
instance Serialize Location
#if MIN_VERSION_cereal(0,5,8)
#else
instance Serialize BSS.ShortByteString where
put = put . BSS.fromShort
get = BSS.toShort <$> get
#endif
class StoreItem item where
toByteString :: item -> BS.ByteString
fromByteString :: BS.ByteString -> Either String item
type KeyArgs item = arg | arg -> item
itemKey :: KeyArgs item -> BSS.ShortByteString
indexFileName :: FilePath
indexFileName = "index"
valuesFileName :: FilePath
valuesFileName = "values"
data ReadMode
= ReadEntireFile
| MemoryMapFile
deriving (Show, Eq)
load :: FilePath -> ReadMode -> IO (Either String Store)
load directoryPath readMode = do
let valuesFilePath = directoryPath </> valuesFileName
indexFilePath = directoryPath </> indexFileName
(valuesFileExists, indexFileExists) <-
(,) <$> doesFileExist indexFilePath <*> doesFileExist valuesFilePath
case (valuesFileExists, indexFileExists) of
(True, True) -> do
indexFile <- BS.readFile indexFilePath
valuesFile <-
case readMode of
ReadEntireFile -> BS.readFile valuesFilePath
MemoryMapFile -> mmapFileByteString valuesFilePath Nothing
let eitherIndex = decode @(M.Map BSS.ShortByteString Location) indexFile
case eitherIndex of
Right locMap -> return . Right $ Store {index = locMap, values = valuesFile}
Left err -> return . Left $ "Error while decoding index : " ++ err
(False, False) ->
return . Left $ "Cannot find index and values in " ++ directoryPath
(True, False) -> return . Left $ "Cannot find index in " ++ directoryPath
(False, True) -> return . Left $ "Cannot find values in " ++ directoryPath
lookup :: (StoreItem item) => KeyArgs item -> Store -> Either String item
lookup keyArgs = lookupByteString (itemKey keyArgs)
lookupByteString ::
(StoreItem item) => BSS.ShortByteString -> Store -> Either String item
lookupByteString key store =
case M.lookup key (index store) of
Just (Location off len) ->
fromByteString . BS.take len . BS.drop off $ values store
Nothing -> Left $ "Cannot find item with key " ++ show key
data State =
State (M.Map BSS.ShortByteString Location)
Handle
add :: (StoreItem item) => item -> KeyArgs item -> S.StateT State IO ()
add item keyArgs = do
let bs = toByteString item
len = BS.length bs
State locMap handle <- S.get
off <- liftIO $ fromIntegral <$> hTell handle
_ <- liftIO $ BS.hPut handle bs
S.put $ State (M.insert (itemKey keyArgs) (Location off len) locMap) handle
createStore ::
FilePath -> (Handle -> IO (M.Map BSS.ShortByteString Location)) -> IO ()
createStore directoryPath action =
withFile (directoryPath </> valuesFileName) WriteMode $ \valuesHandle -> do
locMap <- action valuesHandle
BS.writeFile (directoryPath </> indexFileName) (encode locMap)
writeValues ::
Handle
-> M.Map BSS.ShortByteString Location
-> S.StateT State IO ()
-> IO (M.Map BSS.ShortByteString Location)
writeValues handle locMap action = do
State locMap' _ <- S.execStateT action (State locMap handle)
return locMap'
|
