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
|
{-# 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
instance Serialize BSS.ShortByteString where
put = put . BSS.fromShort
get = BSS.toShort <$> get
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'
|
