never executed always true always false
    1 {-|
    2 Module: Flaw.Data.Lmdb
    3 Description: Simple Lmdb Haskell interface.
    4 License: MIT
    5 -}
    6 
    7 {-# LANGUAGE PatternSynonyms #-}
    8 {-# OPTIONS_GHC -fno-warn-missing-pattern-synonym-signatures #-}
    9 
   10 module Flaw.Data.Lmdb
   11   ( Lmdb()
   12   , LmdbTransaction()
   13   , lmdbOpen
   14   , lmdbRead
   15   , lmdbWrite
   16   , lmdbCommit
   17   , lmdbGet
   18   , lmdbPut
   19   , lmdbDelete
   20   , lmdbFoldPrefixRange
   21   ) where
   22 
   23 import Control.Exception
   24 import Control.Monad
   25 import Data.Bits
   26 import qualified Data.ByteString as B
   27 import qualified Data.ByteString.Unsafe as B
   28 import Data.IORef
   29 import qualified Data.Text as T
   30 import qualified Data.Text.Encoding as T
   31 import Data.Word
   32 import Foreign.C.Types
   33 import Foreign.Marshal.Alloc
   34 import Foreign.Marshal.Array
   35 import Foreign.Ptr
   36 import Foreign.Storable
   37 
   38 import Flaw.Book
   39 import Flaw.Flow
   40 
   41 data Lmdb = Lmdb
   42   { lmdbEnvPtr :: {-# UNPACK #-} !(Ptr MDB_env)
   43   , lmdbDbi :: {-# UNPACK #-} !MDB_dbi
   44   , lmdbFlow :: !Flow
   45   }
   46 
   47 data LmdbTransaction = LmdbTransaction
   48   { lmdbTransactionTxnPtr :: {-# UNPACK #-} !(Ptr MDB_txn)
   49   , lmdbTransactionDbi :: {-# UNPACK #-} !MDB_dbi
   50   , lmdbTransactionFinishedRef :: {-# UNPACK #-} !(IORef Bool)
   51   }
   52 
   53 -- | Open LMDB environment.
   54 lmdbOpen :: T.Text -> Word64 -> IO (Lmdb, IO ())
   55 lmdbOpen fileName mapSize = withSpecialBook $ \bk -> do
   56   -- create env
   57   envPtr <- alloca $ \envPtrPtr -> do
   58     lmdbCheckError $ mdb_env_create envPtrPtr
   59     peek envPtrPtr
   60   book bk $ return ((), mdb_env_close envPtr)
   61 
   62   -- set memory map size
   63   lmdbCheckError $ mdb_env_set_mapsize envPtr (fromIntegral mapSize)
   64 
   65   -- open env
   66   lmdbCheckError $ B.useAsCString (T.encodeUtf8 fileName) $ \fileNamePtr ->
   67     mdb_env_open envPtr fileNamePtr (MDB_NOSUBDIR .|. MDB_WRITEMAP .|. MDB_NOSYNC .|. MDB_NOTLS) 0o644
   68 
   69   -- open database
   70   dbi <- let
   71     acquire = alloca $ \txnPtrPtr -> do
   72       lmdbCheckError $ mdb_txn_begin envPtr nullPtr MDB_RDONLY txnPtrPtr
   73       peek txnPtrPtr
   74     in bracketOnError acquire mdb_txn_abort $ \txnPtr -> do
   75       dbi <- alloca $ \dbiPtr -> do
   76         lmdbCheckError $ mdb_dbi_open txnPtr nullPtr MDB_CREATE dbiPtr
   77         peek dbiPtr
   78       lmdbCheckError $ mdb_txn_commit txnPtr
   79       return dbi
   80   book bk $ return ((), mdb_dbi_close envPtr dbi)
   81 
   82   -- create flow
   83   flow <- book bk newFlowOS
   84 
   85   return Lmdb
   86     { lmdbEnvPtr = envPtr
   87     , lmdbDbi = dbi
   88     , lmdbFlow = flow
   89     }
   90 
   91 -- | Run read transaction.
   92 lmdbRead :: Lmdb -> (LmdbTransaction -> IO a) -> IO a
   93 lmdbRead Lmdb
   94   { lmdbEnvPtr = envPtr
   95   , lmdbDbi = dbi
   96   } io = do
   97   finishedRef <- newIORef False
   98   let
   99     acquire = alloca $ \txnPtrPtr -> do
  100       lmdbCheckError $ mdb_txn_begin envPtr nullPtr MDB_RDONLY txnPtrPtr
  101       peek txnPtrPtr
  102     release txnPtr = do
  103       finished <- readIORef finishedRef
  104       unless finished $ mdb_txn_abort txnPtr
  105   bracket acquire release $ \txnPtr -> io LmdbTransaction
  106     { lmdbTransactionTxnPtr = txnPtr
  107     , lmdbTransactionDbi = dbi
  108     , lmdbTransactionFinishedRef = finishedRef
  109     }
  110 
  111 -- | Run write transaction.
  112 lmdbWrite :: Lmdb -> (LmdbTransaction -> IO a) -> IO a
  113 lmdbWrite Lmdb
  114   { lmdbEnvPtr = envPtr
  115   , lmdbDbi = dbi
  116   , lmdbFlow = flow
  117   } io = runInFlow flow $ do
  118   finishedRef <- newIORef False
  119   let
  120     acquire = alloca $ \txnPtrPtr -> do
  121       lmdbCheckError $ mdb_txn_begin envPtr nullPtr 0 txnPtrPtr
  122       peek txnPtrPtr
  123     release txnPtr = do
  124       finished <- readIORef finishedRef
  125       unless finished $ mdb_txn_abort txnPtr
  126   bracket acquire release $ \txnPtr -> io LmdbTransaction
  127     { lmdbTransactionTxnPtr = txnPtr
  128     , lmdbTransactionDbi = dbi
  129     , lmdbTransactionFinishedRef = finishedRef
  130     }
  131 
  132 lmdbCommit :: LmdbTransaction -> IO ()
  133 lmdbCommit LmdbTransaction
  134   { lmdbTransactionTxnPtr = txnPtr
  135   , lmdbTransactionFinishedRef = finishedRef
  136   } = do
  137   lmdbCheckError $ mdb_txn_commit txnPtr
  138   writeIORef finishedRef True
  139 
  140 lmdbGet :: LmdbTransaction -> B.ByteString -> IO (Maybe B.ByteString)
  141 lmdbGet LmdbTransaction
  142   { lmdbTransactionTxnPtr = txnPtr
  143   , lmdbTransactionDbi = dbi
  144   } key = B.unsafeUseAsCStringLen key $ \(keyPtr, keyLength) -> allocaArray 2 $ \keyBufPtr -> do
  145   poke keyBufPtr $ intPtrToPtr $ fromIntegral keyLength
  146   pokeElemOff keyBufPtr 1 keyPtr
  147   allocaArray 2 $ \valueBufPtr -> do
  148     r <- mdb_get txnPtr dbi keyBufPtr valueBufPtr
  149     if r == MDB_SUCCESS then do
  150       valueLength <- fromIntegral . ptrToIntPtr <$> peek valueBufPtr
  151       valuePtr <- peekElemOff valueBufPtr 1
  152       Just <$> B.packCStringLen (valuePtr, valueLength)
  153     else if r == MDB_NOTFOUND then return Nothing
  154     else lmdbThrowError r
  155 
  156 lmdbPut :: LmdbTransaction -> B.ByteString -> B.ByteString -> IO ()
  157 lmdbPut LmdbTransaction
  158   { lmdbTransactionTxnPtr = txnPtr
  159   , lmdbTransactionDbi = dbi
  160   } key value = lmdbCheckError $
  161   B.unsafeUseAsCStringLen key $ \(keyPtr, keyLength) -> allocaArray 2 $ \keyBufPtr -> do
  162     poke keyBufPtr $ intPtrToPtr $ fromIntegral keyLength
  163     pokeElemOff keyBufPtr 1 keyPtr
  164     B.unsafeUseAsCStringLen value $ \(valuePtr, valueLength) -> allocaArray 2 $ \valueBufPtr -> do
  165       poke valueBufPtr $ intPtrToPtr $ fromIntegral valueLength
  166       pokeElemOff valueBufPtr 1 valuePtr
  167       mdb_put txnPtr dbi keyBufPtr valueBufPtr 0
  168 
  169 lmdbDelete :: LmdbTransaction -> B.ByteString -> IO ()
  170 lmdbDelete LmdbTransaction
  171   { lmdbTransactionTxnPtr = txnPtr
  172   , lmdbTransactionDbi = dbi
  173   } key = lmdbCheckError $
  174   B.unsafeUseAsCStringLen key $ \(keyPtr, keyLength) -> allocaArray 2 $ \keyBufPtr -> do
  175     poke keyBufPtr $ intPtrToPtr $ fromIntegral keyLength
  176     pokeElemOff keyBufPtr 1 keyPtr
  177     mdb_del txnPtr dbi keyBufPtr nullPtr
  178 
  179 -- | Fold key-values pairs prefixed with specified prefix.
  180 lmdbFoldPrefixRange :: LmdbTransaction -> B.ByteString -> a -> (B.ByteString -> B.ByteString -> a -> IO (Bool, a)) -> IO a
  181 lmdbFoldPrefixRange LmdbTransaction
  182   { lmdbTransactionTxnPtr = txnPtr
  183   , lmdbTransactionDbi = dbi
  184   } keyPrefix z0 step = bracket acquireCursor releaseCursor $ \cursorPtr ->
  185   allocaArray 2 $ \keyBufPtr -> allocaArray 2 $ \valueBufPtr ->
  186   B.unsafeUseAsCStringLen keyPrefix $ \(keyPrefixPtr, keyPrefixLength) -> do
  187     -- iteration function
  188     let
  189       iteration z r = do
  190         -- if we got key-value pair
  191         if r == MDB_SUCCESS then do
  192           -- get key
  193           keyLength <- fromIntegral . ptrToIntPtr <$> peek keyBufPtr
  194           keyPtr <- peekElemOff keyBufPtr 1
  195           key <- B.packCStringLen (keyPtr, keyLength)
  196           -- check that key prefix is an actual prefix
  197           if B.isPrefixOf keyPrefix key then do
  198             -- get value
  199             valueLength <- fromIntegral . ptrToIntPtr <$> peek valueBufPtr
  200             valuePtr <- peekElemOff valueBufPtr 1
  201             value <- B.packCStringLen (valuePtr, valueLength)
  202             -- call step function
  203             (continue, nz) <- step key value z
  204             if continue
  205               -- go to next key-value pair and repeat
  206               then iteration nz =<< mdb_cursor_get cursorPtr keyBufPtr valueBufPtr MDB_NEXT
  207               else return nz
  208           else return z
  209         -- else we got to an end
  210         else if r == MDB_NOTFOUND then return z
  211         -- else it's error
  212         else lmdbThrowError r
  213 
  214     -- place cursor on first item, and start iterations
  215     poke keyBufPtr $ intPtrToPtr $ fromIntegral keyPrefixLength
  216     pokeElemOff keyBufPtr 1 keyPrefixPtr
  217     iteration z0 =<< mdb_cursor_get cursorPtr keyBufPtr valueBufPtr MDB_SET_RANGE
  218   where
  219     acquireCursor = alloca $ \cursorPtrPtr -> do
  220       lmdbCheckError $ mdb_cursor_open txnPtr dbi cursorPtrPtr
  221       peek cursorPtrPtr
  222     releaseCursor = mdb_cursor_close
  223 
  224 lmdbCheckError :: IO CInt -> IO ()
  225 lmdbCheckError io = do
  226   r <- io
  227   unless (r == MDB_SUCCESS) $ lmdbThrowError r
  228 
  229 lmdbThrowError :: CInt -> IO a
  230 lmdbThrowError r = throwIO . LmdbError r . T.decodeUtf8 =<< B.packCString =<< mdb_strerror r
  231 
  232 data LmdbError
  233   = LmdbError {-# UNPACK #-} !CInt !T.Text
  234   deriving Show
  235 
  236 instance Exception LmdbError
  237 
  238 -- FFI: types
  239 
  240 data MDB_env
  241 data MDB_txn
  242 data MDB_cursor
  243 type MDB_dbi = CUInt
  244 -- MDB_val is actually struct { size_t, void* } but we use pair of pointers.
  245 type MDB_val = Ptr CChar
  246 
  247 -- FFI: functions
  248 
  249 foreign import ccall safe mdb_strerror :: CInt -> IO (Ptr CChar)
  250 
  251 foreign import ccall safe mdb_env_create :: Ptr (Ptr MDB_env) -> IO CInt
  252 foreign import ccall safe mdb_env_close :: Ptr MDB_env -> IO ()
  253 foreign import ccall safe mdb_env_open :: Ptr MDB_env -> Ptr CChar -> CUInt -> Word32 -> IO CInt
  254 foreign import ccall safe mdb_env_set_mapsize :: Ptr MDB_env -> CSize -> IO CInt
  255 
  256 foreign import ccall safe mdb_txn_begin :: Ptr MDB_env -> Ptr MDB_txn -> CUInt -> Ptr (Ptr MDB_txn) -> IO CInt
  257 foreign import ccall safe mdb_txn_commit :: Ptr MDB_txn -> IO CInt
  258 foreign import ccall safe mdb_txn_abort :: Ptr MDB_txn -> IO ()
  259 
  260 foreign import ccall safe mdb_dbi_open :: Ptr MDB_txn -> Ptr CChar -> CUInt -> Ptr MDB_dbi -> IO CInt
  261 foreign import ccall safe mdb_dbi_close :: Ptr MDB_env -> MDB_dbi -> IO ()
  262 
  263 foreign import ccall safe mdb_get :: Ptr MDB_txn -> MDB_dbi -> Ptr MDB_val -> Ptr MDB_val -> IO CInt
  264 foreign import ccall safe mdb_put :: Ptr MDB_txn -> MDB_dbi -> Ptr MDB_val -> Ptr MDB_val -> CUInt -> IO CInt
  265 foreign import ccall safe mdb_del :: Ptr MDB_txn -> MDB_dbi -> Ptr MDB_val -> Ptr MDB_val -> IO CInt
  266 
  267 foreign import ccall safe mdb_cursor_open :: Ptr MDB_txn -> MDB_dbi -> Ptr (Ptr MDB_cursor) -> IO CInt
  268 foreign import ccall safe mdb_cursor_close :: Ptr MDB_cursor -> IO ()
  269 foreign import ccall safe mdb_cursor_get :: Ptr MDB_cursor -> Ptr MDB_val -> Ptr MDB_val -> CInt -> IO CInt
  270 
  271 -- FFI: values
  272 
  273 pattern MDB_SUCCESS = 0
  274 pattern MDB_NOTFOUND = (-30798)
  275 
  276 pattern MDB_NOSUBDIR = 0x4000
  277 pattern MDB_NOSYNC = 0x10000
  278 pattern MDB_WRITEMAP = 0x80000
  279 pattern MDB_NOTLS = 0x200000
  280 
  281 pattern MDB_RDONLY = 0x20000
  282 
  283 pattern MDB_CREATE = 0x40000
  284 
  285 pattern MDB_NEXT = 8
  286 pattern MDB_SET_RANGE = 17