streamly-lmdb 0.3.0 → 0.4.0
raw patch · 8 files changed
+711/−462 lines, 8 filesPVP ok
version bump matches the API change (PVP)
API changes (from Hackage documentation)
+ Streamly.External.LMDB: abortReadOnlyTxn :: ReadOnlyTxn -> IO ()
+ Streamly.External.LMDB: beginReadOnlyTxn :: Environment mode -> IO ReadOnlyTxn
+ Streamly.External.LMDB: data ReadOnlyTxn
- Streamly.External.LMDB: readLMDB :: (MonadIO m, Mode mode) => Database mode -> ReadOptions -> Unfold m Void (ByteString, ByteString)
+ Streamly.External.LMDB: readLMDB :: (MonadIO m, Mode mode) => Database mode -> Maybe ReadOnlyTxn -> ReadOptions -> Unfold m Void (ByteString, ByteString)
- Streamly.External.LMDB: unsafeReadLMDB :: (MonadIO m, Mode mode) => Database mode -> ReadOptions -> (CStringLen -> IO k) -> (CStringLen -> IO v) -> Unfold m Void (k, v)
+ Streamly.External.LMDB: unsafeReadLMDB :: (MonadIO m, Mode mode) => Database mode -> Maybe ReadOnlyTxn -> ReadOptions -> (CStringLen -> IO k) -> (CStringLen -> IO v) -> Unfold m Void (k, v)
Files
- ChangeLog.md +4/−0
- LICENSE +1/−1
- README.md +35/−26
- src/Streamly/External/LMDB.hs +383/−240
- src/Streamly/External/LMDB/Internal.hs +4/−2
- streamly-lmdb.cabal +3/−3
- test/Streamly/External/LMDB/Tests.hs +252/−174
- test/TestSuite.hs +29/−16
ChangeLog.md view
@@ -1,3 +1,7 @@+## 0.4.0++* Added support for precreating read-only transactions.+ ## 0.3.0 * Updated for Streamly 0.8.0.
LICENSE view
@@ -1,4 +1,4 @@-Copyright Shlok Datye (c) 2020+Copyright Shlok Datye (c) 2022 All rights reserved.
README.md view
@@ -20,44 +20,53 @@ module Main where import Streamly.External.LMDB- (Limits (mapSize), WriteOptions (writeTransactionSize), defaultLimits,- defaultReadOptions, defaultWriteOptions, getDatabase, openEnvironment,- readLMDB, tebibyte, writeLMDB)+ ( Limits (mapSize),+ WriteOptions (writeTransactionSize),+ defaultLimits,+ defaultReadOptions,+ defaultWriteOptions,+ getDatabase,+ openEnvironment,+ readLMDB,+ tebibyte,+ writeLMDB,+ ) import qualified Streamly.Prelude as S main :: IO () main = do- -- Open an environment. There should already exist a file or- -- directory at the given path. (Empty for a new environment.)- env <- openEnvironment "/path/to/lmdb-database" $- defaultLimits { mapSize = tebibyte }+ -- Open an environment. There should already exist a file or+ -- directory at the given path. (Empty for a new environment.)+ env <-+ openEnvironment "/path/to/lmdb-database" $+ defaultLimits {mapSize = tebibyte} - -- Get the main database.- -- Note: It is common practice with LMDB to create the database- -- once and reuse it for the remainder of the program’s execution.- db <- getDatabase env Nothing+ -- Get the main database.+ -- Note: It is common practice with LMDB to create the database+ -- once and reuse it for the remainder of the program’s execution.+ db <- getDatabase env Nothing - -- Stream key-value pairs into the database.- let fold' = writeLMDB db defaultWriteOptions { writeTransactionSize = 1 }- let writeStream = S.fromList [("baz", "a"), ("foo", "b"), ("bar", "c")]- _ <- S.fold fold' writeStream+ -- Stream key-value pairs into the database.+ let fold' = writeLMDB db defaultWriteOptions {writeTransactionSize = 1}+ let writeStream = S.fromList [("baz", "a"), ("foo", "b"), ("bar", "c")]+ _ <- S.fold fold' writeStream - -- Stream key-value pairs out of the- -- database, printing them along the way.- -- Output:- -- ("bar","c")- -- ("baz","a")- -- ("foo","b")- let unfold' = readLMDB db defaultReadOptions- let readStream = S.unfold unfold' undefined- S.mapM_ print readStream+ -- Stream key-value pairs out of the+ -- database, printing them along the way.+ -- Output:+ -- ("bar","c")+ -- ("baz","a")+ -- ("foo","b")+ let unfold' = readLMDB db Nothing defaultReadOptions+ let readStream = S.unfold unfold' undefined+ S.mapM_ print readStream ``` ## Benchmarks See `bench/README.md`. Summary (with rough figures from our machine<sup>†</sup>): -* **Reading.** For reading a fully cached LMDB database, this library (when `unsafeReadLMDB` is used instead of `readLMDB`) has a 10 ns/pair overhead compared to plain Haskell `IO` code, which has another 10 ns/pair overhead compared to C. (The first two being similar fulfills the promise of [streamly](https://hackage.haskell.org/package/streamly) and stream fusion.) We deduce that if your total workload per pair takes longer than 20 ns, your bottleneck will not be your usage of this library as opposed to C.-* **Writing**. Writing with plain Haskell `IO` code and with this library is, respectively, 10% and 20% slower than writing with C. We have not dug further into these differences because this write performance is currently good enough for our purposes.+* **Reading.** For reading a fully cached LMDB database, this library (when `unsafeReadLMDB` is used instead of `readLMDB`) has roughly a 15 ns/pair overhead compared to plain Haskell `IO` code, which has roughly another 10 ns/pair overhead compared to C. (The first two being similar fulfills the promise of [streamly](https://hackage.haskell.org/package/streamly) and stream fusion.) We deduce that if your total workload per pair takes longer than around 25 ns, your bottleneck will not be your usage of this library as opposed to C.+* **Writing**. Writing with plain Haskell `IO` code and with this library is, respectively, around 30% and 50% slower than writing with C. We have not dug further into these differences because this write performance is currently good enough for our purposes. <sup>†</sup> [Linode](https://linode.com); Debian 10, Dedicated 32GB: 16 CPU, 640GB Storage, 32GB RAM.
src/Streamly/External/LMDB.hs view
@@ -1,350 +1,493 @@-{-# LANGUAGE BangPatterns, ScopedTypeVariables #-}+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE ScopedTypeVariables #-} --- | The functionality for the limits and getting the environment and database were mostly--- obtained from the [lmdb-simple](https://hackage.haskell.org/package/lmdb-simple) library.+-- | = Acknowledgments+--+-- The functionality for the limits and getting the environment and database, in particular the idea+-- of specifying the read-only or read-write mode at the type level, was mostly obtained from the+-- [lmdb-simple](https://hackage.haskell.org/package/lmdb-simple) library. module Streamly.External.LMDB- (- -- ** Types- Database,+ ( -- * Environment++ -- | With LMDB, one first creates a so-called “environment,” which one can think of as a file or+ -- folder on disk. Environment,- Limits(..),- LMDB_Error(..),- MDB_ErrCode(..),+ openEnvironment,+ isReadOnlyEnvironment,++ -- ** Mode Mode, ReadWrite, ReadOnly,- ReadDirection(..),- ReadOptions(..),- OverwriteOptions(..),- WriteOptions(..), - -- ** Environment and database+ -- ** Limits+ Limits (..), defaultLimits,- openEnvironment,- isReadOnlyEnvironment,- getDatabase,-- -- ** Utility gibibyte, tebibyte,++ -- * Database++ -- | After creating an environment, one creates within it one or more databases.+ --+ -- Note: We currently have no functionality here for closing and disposing of databases or+ -- environments because we have had no need for it yet. In any case, it is a common practice+ -- with LMDB to create one’s environments and databases once and reuse them for the remainder of+ -- the program’s execution.+ Database,+ getDatabase, clearDatabase, - -- ** Reading- defaultReadOptions,+ -- * Reading readLMDB, unsafeReadLMDB, - -- ** Writing+ -- ** Read-only transactions+ ReadOnlyTxn,+ beginReadOnlyTxn,+ abortReadOnlyTxn,++ -- ** Read options+ ReadOptions (..),+ defaultReadOptions,+ ReadDirection (..),++ -- * Writing+ writeLMDB,+ WriteOptions (..), defaultWriteOptions,- writeLMDB) where+ OverwriteOptions (..), + -- * Error types+ LMDB_Error (..),+ MDB_ErrCode (..),+ )+where+ import Control.Concurrent (isCurrentThreadBound, myThreadId) import Control.Concurrent.Async (asyncBound, wait)-import Control.Exception (Exception, catch, tryJust, mask_, throw)+import Control.Exception (Exception, catch, mask_, throw, tryJust) import Control.Monad (guard, unless, when) import Control.Monad.IO.Class (MonadIO, liftIO) import Data.ByteString (ByteString, packCStringLen) import Data.ByteString.Unsafe (unsafePackCStringLen, unsafeUseAsCStringLen)-import Data.Maybe (fromJust)+import Data.Maybe (fromJust, isNothing) import Data.Void (Void) import Foreign (Ptr, alloca, free, malloc, nullPtr, peek) import Foreign.C (Errno (Errno), eNOTDIR) import Foreign.C.String (CStringLen) import Foreign.Marshal.Utils (with) import Foreign.Storable (poke)+import Streamly.External.LMDB.Internal (Database (..), Mode (..), ReadOnly, ReadWrite)+import Streamly.External.LMDB.Internal.Foreign+ ( LMDB_Error (..),+ MDB_ErrCode (..),+ MDB_env,+ MDB_txn,+ MDB_val (MDB_val, mv_data, mv_size),+ c_mdb_cursor_close,+ c_mdb_cursor_get,+ c_mdb_get,+ c_mdb_txn_abort,+ combineOptions,+ mdb_append,+ mdb_clear,+ mdb_create,+ mdb_cursor_open,+ mdb_dbi_open,+ mdb_env_create,+ mdb_env_open,+ mdb_env_set_mapsize,+ mdb_env_set_maxdbs,+ mdb_env_set_maxreaders,+ mdb_first,+ mdb_last,+ mdb_next,+ mdb_nooverwrite,+ mdb_nosubdir,+ mdb_notfound,+ mdb_notls,+ mdb_prev,+ mdb_put_,+ mdb_rdonly,+ mdb_set_range,+ mdb_txn_begin,+ mdb_txn_commit,+ throwLMDBErrNum,+ ) import Streamly.Internal.Data.Fold (Fold (Fold), Step (Partial)) import Streamly.Internal.Data.IOFinalizer (newIOFinalizer, runIOFinalizer) import Streamly.Internal.Data.Stream.StreamD.Type (Step (Stop, Yield)) import Streamly.Internal.Data.Unfold (supply) import Streamly.Internal.Data.Unfold.Type (Unfold (Unfold)) -import Streamly.External.LMDB.Internal-import Streamly.External.LMDB.Internal.Foreign- newtype Environment mode = Environment (Ptr MDB_env) isReadOnlyEnvironment :: Mode mode => Environment mode -> Bool isReadOnlyEnvironment = isReadOnlyMode . mode- where- mode :: Environment mode -> mode- mode = undefined+ where+ mode :: Environment mode -> mode+ mode = undefined --- | LMDB environments have various limits on the size and number of databases and concurrent readers.+-- | LMDB environments have various limits on the size and number of databases and concurrent+-- readers. data Limits = Limits- { mapSize :: !Int -- ^ Memory map size, in bytes (also the maximum size of all databases).- , maxDatabases :: !Int -- ^ Maximum number of named databases.- , maxReaders :: !Int -- ^ Maximum number of concurrent 'ReadOnly' transactions- -- (also the number of slots in the lock table).+ { -- | Memory map size, in bytes (also the maximum size of all databases).+ mapSize :: !Int,+ -- | Maximum number of named databases.+ maxDatabases :: !Int,+ -- | Maximum number of concurrent 'ReadOnly' transactions+ -- (also the number of slots in the lock table).+ maxReaders :: !Int } --- | The default limits are 1 MiB map size, 0 named databases, and 126 concurrent readers. These can be adjusted--- freely, and in particular the 'mapSize' may be set very large (limited only by available address space). However,--- LMDB is not optimized for a large number of named databases so 'maxDatabases' should be kept to a minimum.+-- | The default limits are 1 MiB map size, 0 named databases, and 126 concurrent readers. These can+-- be adjusted freely, and in particular the 'mapSize' may be set very large (limited only by+-- available address space). However, LMDB is not optimized for a large number of named databases so+-- 'maxDatabases' should be kept to a minimum. ----- The default 'mapSize' is intentionally small, and should be changed to something appropriate for your application.--- It ought to be a multiple of the OS page size, and should be chosen as large as possible to accommodate future--- growth of the database(s). Once set for an environment, this limit cannot be reduced to a value smaller than--- the space already consumed by the environment, however it can later be increased.+-- The default 'mapSize' is intentionally small, and should be changed to something appropriate for+-- your application. It ought to be a multiple of the OS page size, and should be chosen as large as+-- possible to accommodate future growth of the database(s). Once set for an environment, this limit+-- cannot be reduced to a value smaller than the space already consumed by the environment, however+-- it can later be increased. ----- If you are going to use any named databases then you will need to change 'maxDatabases'--- to the number of named databases you plan to use. However, you do not need to change--- this field if you are only going to use the single main (unnamed) database.+-- If you are going to use any named databases then you will need to change 'maxDatabases' to the+-- number of named databases you plan to use. However, you do not need to change this field if you+-- are only going to use the single main (unnamed) database. defaultLimits :: Limits-defaultLimits = Limits- { mapSize = 1024 * 1024 -- 1 MiB.- , maxDatabases = 0- , maxReaders = 126+defaultLimits =+ Limits+ { mapSize = 1024 * 1024, -- 1 MiB.+ maxDatabases = 0,+ maxReaders = 126 } --- A convenience constant for obtaining a 1 GiB map size.+-- | A convenience constant for obtaining a 1 GiB map size. gibibyte :: Int gibibyte = 1024 * 1024 * 1024 --- A convenience constant for obtaining a 1 TiB map size.+-- | A convenience constant for obtaining a 1 TiB map size. tebibyte :: Int tebibyte = 1024 * 1024 * 1024 * 1024 --- | Open an LMDB environment in either 'ReadWrite' or 'ReadOnly' mode. The 'FilePath' argument--- may be either a directory or a regular file, but it must already exist. If a regular file,--- an additional file with "-lock" appended to the name is used for the reader lock table.+-- | Open an LMDB environment in either 'ReadWrite' or 'ReadOnly' mode. The 'FilePath' argument may+-- be either a directory or a regular file, but it must already exist. If a regular file, an+-- additional file with "-lock" appended to the name is used for the reader lock table. ----- Note that an environment must have been opened in 'ReadWrite'--- mode at least once before it can be opened in 'ReadOnly' mode.+-- Note that an environment must have been opened in 'ReadWrite' mode at least once before it can be+-- opened in 'ReadOnly' mode. ----- An environment opened in 'ReadOnly' mode may still modify the reader lock table--- (except when the filesystem is read-only, in which case no locks are used).+-- An environment opened in 'ReadOnly' mode may still modify the reader lock table (except when the+-- filesystem is read-only, in which case no locks are used). openEnvironment :: Mode mode => FilePath -> Limits -> IO (Environment mode) openEnvironment path limits = do- penv <- mdb_env_create+ penv <- mdb_env_create - mdb_env_set_mapsize penv (mapSize limits)- let maxDbs = maxDatabases limits in when (maxDbs /= 0) $ mdb_env_set_maxdbs penv maxDbs- mdb_env_set_maxreaders penv (maxReaders limits)+ mdb_env_set_mapsize penv (mapSize limits)+ let maxDbs = maxDatabases limits in when (maxDbs /= 0) $ mdb_env_set_maxdbs penv maxDbs+ mdb_env_set_maxreaders penv (maxReaders limits) - -- Always use MDB_NOTLS.- let env = Environment penv :: Mode mode => Environment mode- flags = mdb_notls : [mdb_rdonly | isReadOnlyEnvironment env]+ -- Always use MDB_NOTLS; this is crucial for Haskell applications. (See+ -- https://github.com/LMDB/lmdb/blob/8d0cbbc936091eb85972501a9b31a8f86d4c51a7/libraries/liblmdb/lmdb.h#L615)+ let env = Environment penv :: Mode mode => Environment mode+ flags = mdb_notls : [mdb_rdonly | isReadOnlyEnvironment env] - let isNotDirectoryError :: LMDB_Error -> Bool- isNotDirectoryError LMDB_Error { e_code = Left code }- | Errno (fromIntegral code) == eNOTDIR = True- isNotDirectoryError _ = False+ let isNotDirectoryError :: LMDB_Error -> Bool+ isNotDirectoryError LMDB_Error {e_code = Left code}+ | Errno (fromIntegral code) == eNOTDIR = True+ isNotDirectoryError _ = False - r <- tryJust (guard . isNotDirectoryError) $ mdb_env_open penv path (combineOptions flags)- case r of- Left _ -> mdb_env_open penv path (combineOptions $ mdb_nosubdir : flags)- Right _ -> return ()+ r <- tryJust (guard . isNotDirectoryError) $ mdb_env_open penv path (combineOptions flags)+ case r of+ Left _ -> mdb_env_open penv path (combineOptions $ mdb_nosubdir : flags)+ Right _ -> return () - return env+ return env +-- | Gets a database with the given name. When creating a database (i.e., getting it for the first+-- time), one must do so in 'ReadWrite' mode.+--+-- If only one database is desired within the environment, the name can be 'Nothing' (known as the+-- “unnamed database”).+--+-- If one or more named databases (a database with a 'Just' name) are desired, the 'maxDatabases' of+-- the environment’s limits should have been adjusted accordingly. The unnamed database will in this+-- case contain the names of the named databases as keys, which one is allowed to read but not+-- write. getDatabase :: (Mode mode) => Environment mode -> Maybe String -> IO (Database mode) getDatabase env@(Environment penv) name = do- ptxn <- mdb_txn_begin penv nullPtr (combineOptions $ [mdb_rdonly | isReadOnlyEnvironment env])- dbi <- mdb_dbi_open ptxn name (combineOptions $ [mdb_create | not $ isReadOnlyEnvironment env])- mdb_txn_commit ptxn- return $ Database penv dbi+ ptxn <- mdb_txn_begin penv nullPtr (combineOptions $ [mdb_rdonly | isReadOnlyEnvironment env])+ dbi <- mdb_dbi_open ptxn name (combineOptions $ [mdb_create | not $ isReadOnlyEnvironment env])+ mdb_txn_commit ptxn+ return $ Database penv dbi -- | Clears, i.e., removes all key-value pairs from, the given database. clearDatabase :: (Mode mode) => Database mode -> IO ()-clearDatabase (Database penv dbi) = asyncBound (do- ptxn <- mdb_txn_begin penv nullPtr 0- mdb_clear ptxn dbi- mdb_txn_commit ptxn) >>= wait---- | Direction of key iteration.-data ReadDirection = Forward | Backward deriving (Show)--data ReadOptions = ReadOptions- { readDirection :: !ReadDirection- -- | If 'Nothing', a forward [backward] iteration starts at the beginning [end] of the database.- -- Otherwise, it starts at the first key that is greater [less] than or equal to the 'Just' key.- , readStart :: !(Maybe ByteString) } deriving (Show)--defaultReadOptions :: ReadOptions-defaultReadOptions = ReadOptions- { readDirection = Forward- , readStart = Nothing }+clearDatabase (Database penv dbi) =+ asyncBound+ ( do+ ptxn <- mdb_txn_begin penv nullPtr 0+ mdb_clear ptxn dbi+ mdb_txn_commit ptxn+ )+ >>= wait -- | Creates an unfold with which we can stream key-value pairs from the given database. ----- A read transaction is kept open for the duration of the unfold; one should therefore--- bear in mind LMDB's [caveats regarding long-lived transactions](https://git.io/JJZE6).+-- If an existing read-only transaction is not provided, a read-only transaction is automatically+-- created and kept open for the duration of the unfold; we suggest doing this as a first option.+-- However, if you find this to be a bottleneck (e.g., if you find upon profiling that a significant+-- time is being spent at @mdb_txn_begin@, or if you find yourself having to increase 'maxReaders'+-- in the environment’s limits because the transactions are not being garbage collected fast+-- enough), consider precreating a transaction using 'beginReadOnlyTxn'. ----- If you don’t want the overhead of intermediate 'ByteString's (on your--- way to your eventual data structures), use 'unsafeReadLMDB' instead.+-- In any case, bear in mind at all times LMDB’s [caveats regarding long-lived+-- transactions](https://github.com/LMDB/lmdb/blob/8d0cbbc936091eb85972501a9b31a8f86d4c51a7/libraries/liblmdb/lmdb.h#L107).+--+-- If you don’t want the overhead of intermediate @ByteString@s (on your way to your eventual data+-- structures), use 'unsafeReadLMDB' instead. {-# INLINE readLMDB #-}-readLMDB :: (MonadIO m, Mode mode) => Database mode -> ReadOptions -> Unfold m Void (ByteString, ByteString)-readLMDB db ropts = unsafeReadLMDB db ropts packCStringLen packCStringLen+readLMDB ::+ (MonadIO m, Mode mode) =>+ Database mode ->+ Maybe ReadOnlyTxn ->+ ReadOptions ->+ Unfold m Void (ByteString, ByteString)+readLMDB db mtxn ropts = unsafeReadLMDB db mtxn ropts packCStringLen packCStringLen --- | Creates an unfold with which we can stream key-value pairs from the given database.------ A read transaction is kept open for the duration of the unfold; one should therefore--- bear in mind LMDB's [caveats regarding long-lived transactions](https://git.io/JJZE6).+-- | Similar to 'readLMDB', except that the keys and values are not automatically converted into+-- Haskell @ByteString@s. ----- To ensure safety, make sure that the memory pointed to by the 'CStringLen' for each key/value mapping function--- call is (a) only read (and not written to); and (b) not used after the mapping function has returned. One way to--- transform the 'CStringLen's to your desired data structures is to use 'Data.ByteString.Unsafe.unsafePackCStringLen'.+-- To ensure safety, make sure that the memory pointed to by the 'CStringLen' for each key/value+-- mapping function call is (a) only read (and not written to); and (b) not used after the mapping+-- function has returned. One way to transform the 'CStringLen's to your desired data structures is+-- to use 'Data.ByteString.Unsafe.unsafePackCStringLen'. {-# INLINE unsafeReadLMDB #-}-unsafeReadLMDB :: (MonadIO m, Mode mode)- => Database mode- -> ReadOptions- -> (CStringLen -> IO k)- -> (CStringLen -> IO v)- -> Unfold m Void (k, v)-unsafeReadLMDB (Database penv dbi) ropts kmap vmap =- let (firstOp, subsequentOp) = case (readDirection ropts, readStart ropts) of- (Forward, Nothing) -> (mdb_first, mdb_next)- (Forward, Just _) -> (mdb_set_range, mdb_next)- (Backward, Nothing) -> (mdb_last, mdb_prev)- (Backward, Just _) -> (mdb_set_range, mdb_prev)- in supply firstOp $ Unfold- (\(op, pcurs, pk, pv, ref) -> do- rc <- liftIO $- if op == mdb_set_range && subsequentOp == mdb_prev then do- -- Reverse MDB_SET_RANGE.- kfst' <- peek pk- kfst <- packCStringLen (mv_data kfst', fromIntegral $ mv_size kfst')- rc <- c_mdb_cursor_get pcurs pk pv op- if rc /= 0 && rc == mdb_notfound then- c_mdb_cursor_get pcurs pk pv mdb_last- else if rc == 0 then do- k' <- peek pk- k <- unsafePackCStringLen (mv_data k', fromIntegral $ mv_size k')- if k /= kfst then- c_mdb_cursor_get pcurs pk pv mdb_prev+unsafeReadLMDB ::+ (MonadIO m, Mode mode) =>+ Database mode ->+ Maybe ReadOnlyTxn ->+ ReadOptions ->+ (CStringLen -> IO k) ->+ (CStringLen -> IO v) ->+ Unfold m Void (k, v)+unsafeReadLMDB (Database penv dbi) mtxn ropts kmap vmap =+ let (firstOp, subsequentOp) = case (readDirection ropts, readStart ropts) of+ (Forward, Nothing) -> (mdb_first, mdb_next)+ (Forward, Just _) -> (mdb_set_range, mdb_next)+ (Backward, Nothing) -> (mdb_last, mdb_prev)+ (Backward, Just _) -> (mdb_set_range, mdb_prev)+ in supply firstOp $+ Unfold+ ( \(op, pcurs, pk, pv, ref) -> do+ rc <-+ liftIO $+ if op == mdb_set_range && subsequentOp == mdb_prev+ then do+ -- Reverse MDB_SET_RANGE.+ kfst' <- peek pk+ kfst <- packCStringLen (mv_data kfst', fromIntegral $ mv_size kfst')+ rc <- c_mdb_cursor_get pcurs pk pv op+ if rc /= 0 && rc == mdb_notfound+ then c_mdb_cursor_get pcurs pk pv mdb_last else- return rc- else- return rc- else- c_mdb_cursor_get pcurs pk pv op+ if rc == 0+ then do+ k' <- peek pk+ k <- unsafePackCStringLen (mv_data k', fromIntegral $ mv_size k')+ if k /= kfst+ then c_mdb_cursor_get pcurs pk pv mdb_prev+ else return rc+ else return rc+ else c_mdb_cursor_get pcurs pk pv op - found <- liftIO $- if rc /= 0 && rc /= mdb_notfound then do- runIOFinalizer ref- throwLMDBErrNum "mdb_cursor_get" rc- else- return $ rc /= mdb_notfound+ found <-+ liftIO $+ if rc /= 0 && rc /= mdb_notfound+ then do+ runIOFinalizer ref+ throwLMDBErrNum "mdb_cursor_get" rc+ else return $ rc /= mdb_notfound - if found then do- !k <- liftIO $ (\x -> kmap (mv_data x, fromIntegral $ mv_size x)) =<< peek pk- !v <- liftIO $ (\x -> vmap (mv_data x, fromIntegral $ mv_size x)) =<< peek pv- return $ Yield (k, v) (subsequentOp, pcurs, pk, pv, ref)- else do- runIOFinalizer ref- return Stop)- (\op -> do- (pcurs, pk, pv, ref) <- liftIO $ mask_ $ do- ptxn <- liftIO $ mdb_txn_begin penv nullPtr mdb_rdonly- pcurs <- liftIO $ mdb_cursor_open ptxn dbi- pk <- liftIO malloc- pv <- liftIO malloc+ if found+ then do+ !k <- liftIO $ (\x -> kmap (mv_data x, fromIntegral $ mv_size x)) =<< peek pk+ !v <- liftIO $ (\x -> vmap (mv_data x, fromIntegral $ mv_size x)) =<< peek pv+ return $ Yield (k, v) (subsequentOp, pcurs, pk, pv, ref)+ else do+ runIOFinalizer ref+ return Stop+ )+ ( \op -> do+ (pcurs, pk, pv, ref) <- liftIO $+ mask_ $ do+ ptxn <- case mtxn of+ Nothing -> liftIO $ mdb_txn_begin penv nullPtr mdb_rdonly+ Just (ReadOnlyTxn ptxn') -> return ptxn'+ pcurs <- liftIO $ mdb_cursor_open ptxn dbi+ pk <- liftIO malloc+ pv <- liftIO malloc - _ <- case readStart ropts of+ _ <- case readStart ropts of Nothing -> return () Just k -> unsafeUseAsCStringLen k $ \(kp, kl) ->- poke pk (MDB_val (fromIntegral kl) kp)+ poke pk (MDB_val (fromIntegral kl) kp) - ref <- liftIO . newIOFinalizer $ do+ ref <- liftIO . newIOFinalizer $ do free pv >> free pk+ -- Note: If no transaction is provided to this function, it could happen that+ -- mdb_cursor_close() gets called after mdb_txn_abort() -- the former being+ -- called during garbage collection. This is allowed by LMDB for read-only+ -- transactions. c_mdb_cursor_close pcurs- -- No need to commit this read-only transaction.- c_mdb_txn_abort ptxn- return (pcurs, pk, pv, ref)- return (op, pcurs, pk, pv, ref))+ when (isNothing mtxn) $+ -- There is no need to commit this read-only transaction.+ c_mdb_txn_abort ptxn+ return (pcurs, pk, pv, ref)+ return (op, pcurs, pk, pv, ref)+ ) +newtype ReadOnlyTxn = ReadOnlyTxn (Ptr MDB_txn)++-- | Begins an LMDB read-only transaction for use with 'readLMDB' or 'unsafeReadLMDB'. It is your+-- responsibility to (a) use the transaction only on databases in the same environment, (b) make+-- sure that those databases were already obtained before the transaction was begun, and (c) dispose+-- of the transaction with 'abortReadOnlyTxn'.+beginReadOnlyTxn :: Environment mode -> IO ReadOnlyTxn+beginReadOnlyTxn (Environment penv) = ReadOnlyTxn <$> mdb_txn_begin penv nullPtr mdb_rdonly++-- | Disposes of a read-only transaction created with 'beginReadOnlyTxn'.+abortReadOnlyTxn :: ReadOnlyTxn -> IO ()+abortReadOnlyTxn (ReadOnlyTxn ptxn) = c_mdb_txn_abort ptxn++data ReadOptions = ReadOptions+ { readDirection :: !ReadDirection,+ -- | If 'Nothing', a forward [backward] iteration starts at the beginning [end] of the database.+ -- Otherwise, it starts at the first key that is greater [less] than or equal to the 'Just' key.+ readStart :: !(Maybe ByteString)+ }+ deriving (Show)++-- | By default, we start reading from the beginning of the database (i.e., from the smallest key).+defaultReadOptions :: ReadOptions+defaultReadOptions =+ ReadOptions+ { readDirection = Forward,+ readStart = Nothing+ }++-- | Direction of key iteration.+data ReadDirection = Forward | Backward deriving (Show)+ data OverwriteOptions = OverwriteAllow | OverwriteAllowSame | OverwriteDisallow deriving (Eq) data WriteOptions = WriteOptions- { writeTransactionSize :: !Int- , overwriteOptions :: !OverwriteOptions- , writeAppend :: !Bool }+ { writeTransactionSize :: !Int,+ overwriteOptions :: !OverwriteOptions,+ writeAppend :: !Bool+ } defaultWriteOptions :: WriteOptions-defaultWriteOptions = WriteOptions- { writeTransactionSize = 1- , overwriteOptions = OverwriteAllow- , writeAppend = False }-+defaultWriteOptions =+ WriteOptions+ { writeTransactionSize = 1,+ overwriteOptions = OverwriteAllow,+ writeAppend = False+ } newtype ExceptionString = ExceptionString String deriving (Show)+ instance Exception ExceptionString -- | Creates a fold with which we can stream key-value pairs into the given database. -- -- It is the responsibility of the user to execute the fold on a bound thread. ----- The fold currently cannot be used with a scan. (The plan is for this shortcoming to be--- remedied with or after a future release of streamly that addresses the underlying issue.)+-- The fold currently cannot be used with a scan. (The plan is for this shortcoming to be remedied+-- with or after a future release of streamly that addresses the underlying issue.) ----- Please specify a suitable transaction size in the write options; the default of 1 (one write transaction for each--- key-value pair) could yield suboptimal performance. One could try, e.g., 100 KB chunks and benchmark from there.+-- Please specify a suitable transaction size in the write options; the default of 1 (one write+-- transaction for each key-value pair) could yield suboptimal performance. One could try, e.g., 100+-- KB chunks and benchmark from there. {-# INLINE writeLMDB #-} writeLMDB :: (MonadIO m) => Database ReadWrite -> WriteOptions -> Fold m (ByteString, ByteString) () writeLMDB (Database penv dbi) options =- let txnSize = max 1 (writeTransactionSize options)- overwriteOpt = overwriteOptions options- flags = combineOptions $- [mdb_nooverwrite | overwriteOpt `elem` [OverwriteAllowSame, OverwriteDisallow]]- ++ [mdb_append | writeAppend options]- in Fold (\(threadId, iter, currChunkSz, mtxn) (k, v) -> do- -- In the first few iterations, ascertain that we are still on the same (bound) thread.- iter' <- liftIO $- if iter < 3 then do- threadId' <- myThreadId- when (threadId' /= threadId) $- throw (ExceptionString "Error: writeLMDB veered off the original bound thread")- return $ iter + 1- else- return iter+ let txnSize = max 1 (writeTransactionSize options)+ overwriteOpt = overwriteOptions options+ flags =+ combineOptions $+ [mdb_nooverwrite | overwriteOpt `elem` [OverwriteAllowSame, OverwriteDisallow]]+ ++ [mdb_append | writeAppend options]+ in Fold+ ( \(threadId, iter, currChunkSz, mtxn) (k, v) -> do+ -- In the first few iterations, ascertain that we are still on the same (bound) thread.+ iter' <-+ liftIO $+ if iter < 3+ then do+ threadId' <- myThreadId+ when (threadId' /= threadId) $+ throw+ (ExceptionString "Error: writeLMDB veered off the original bound thread")+ return $ iter + 1+ else return iter - currChunkSz' <- liftIO $- if currChunkSz >= txnSize then do- let (_, ref) = fromJust mtxn- runIOFinalizer ref- return 0- else- return currChunkSz+ currChunkSz' <-+ liftIO $+ if currChunkSz >= txnSize+ then do+ let (_, ref) = fromJust mtxn+ runIOFinalizer ref+ return 0+ else return currChunkSz - (ptxn, ref) <-- if currChunkSz' == 0 then- liftIO $ mask_ $ do+ (ptxn, ref) <-+ if currChunkSz' == 0+ then liftIO $+ mask_ $ do ptxn <- mdb_txn_begin penv nullPtr 0 ref <- newIOFinalizer $ mdb_txn_commit ptxn return (ptxn, ref)- else- return $ fromJust mtxn+ else return $ fromJust mtxn - liftIO $ unsafeUseAsCStringLen k $ \(kp, kl) -> unsafeUseAsCStringLen v $ \(vp, vl) ->- catch (mdb_put_ ptxn dbi kp (fromIntegral kl) vp (fromIntegral vl) flags)- (\(e :: LMDB_Error) -> do- -- Discard error if OverwriteAllowSame was specified and the error from LMDB- -- was due to the exact same key-value pair already existing in the database.- ok <- with (MDB_val (fromIntegral kl) kp) $ \pk ->- alloca $ \pv -> do- rc <- c_mdb_get ptxn dbi pk pv- if rc == 0 then do- v' <- peek pv- vbs <- unsafePackCStringLen (mv_data v', fromIntegral $ mv_size v')- return $ overwriteOpt == OverwriteAllowSame- && e_code e == Right MDB_KEYEXIST- && vbs == v- else- return False- unless ok $ runIOFinalizer ref >> throw e)- return $ Partial (threadId, iter', currChunkSz' + 1, Just (ptxn, ref)))- (do- isBound <- liftIO isCurrentThreadBound- threadId <- liftIO myThreadId- if isBound then- return $ Partial (threadId, 0 :: Int, 0, Nothing)- else- throw $ ExceptionString "Error: writeLMDB should be executed on a bound thread")- -- This final part is incompatible with scans.- (\(threadId, _, _, mtxn) -> liftIO $ do- threadId' <- myThreadId- when (threadId' /= threadId) $- throw (ExceptionString "Error: writeLMDB veered off the original bound thread at the end")- case mtxn of- Nothing -> return ()- Just (_, ref) -> runIOFinalizer ref)+ liftIO $+ unsafeUseAsCStringLen k $ \(kp, kl) -> unsafeUseAsCStringLen v $ \(vp, vl) ->+ catch+ (mdb_put_ ptxn dbi kp (fromIntegral kl) vp (fromIntegral vl) flags)+ ( \(e :: LMDB_Error) -> do+ -- Discard error if OverwriteAllowSame was specified and the error from LMDB+ -- was due to the exact same key-value pair already existing in the database.+ ok <- with (MDB_val (fromIntegral kl) kp) $ \pk ->+ alloca $ \pv -> do+ rc <- c_mdb_get ptxn dbi pk pv+ if rc == 0+ then do+ v' <- peek pv+ vbs <- unsafePackCStringLen (mv_data v', fromIntegral $ mv_size v')+ return $+ overwriteOpt == OverwriteAllowSame+ && e_code e == Right MDB_KEYEXIST+ && vbs == v+ else return False+ unless ok $ runIOFinalizer ref >> throw e+ )+ return $ Partial (threadId, iter', currChunkSz' + 1, Just (ptxn, ref))+ )+ ( do+ isBound <- liftIO isCurrentThreadBound+ threadId <- liftIO myThreadId+ if isBound+ then return $ Partial (threadId, 0 :: Int, 0, Nothing)+ else throw $ ExceptionString "Error: writeLMDB should be executed on a bound thread"+ )+ -- This final part is incompatible with scans.+ ( \(threadId, _, _, mtxn) -> liftIO $ do+ threadId' <- myThreadId+ when (threadId' /= threadId) $+ throw+ (ExceptionString "Error: writeLMDB veered off the original bound thread at the end")+ case mtxn of+ Nothing -> return ()+ Just (_, ref) -> runIOFinalizer ref+ )
src/Streamly/External/LMDB/Internal.hs view
@@ -6,12 +6,14 @@ -- This is in a separate internal module because the tests make use of the Database constructor. class Mode a where- isReadOnlyMode :: a -> Bool+ isReadOnlyMode :: a -> Bool data ReadWrite+ data ReadOnly instance Mode ReadWrite where isReadOnlyMode _ = False-instance Mode ReadOnly where isReadOnlyMode _ = True++instance Mode ReadOnly where isReadOnlyMode _ = True data Database mode = Database (Ptr MDB_env) MDB_dbi_t
streamly-lmdb.cabal view
@@ -1,11 +1,11 @@ cabal-version: 1.12 --- This file has been generated from package.yaml by hpack version 0.34.4.+-- This file has been generated from package.yaml by hpack version 0.34.6. -- -- see: https://github.com/sol/hpack name: streamly-lmdb-version: 0.3.0+version: 0.4.0 synopsis: Stream data to or from LMDB databases using the streamly library. description: Please see the README on GitHub at <https://github.com/shlok/streamly-lmdb#readme> category: Database, Streaming, Streamly@@ -13,7 +13,7 @@ bug-reports: https://github.com/shlok/streamly-lmdb/issues author: Shlok Datye maintainer: sd-haskell@quant.is-copyright: 2021 Shlok Datye+copyright: 2022 Shlok Datye license: BSD3 license-file: LICENSE build-type: Simple
test/Streamly/External/LMDB/Tests.hs view
@@ -3,206 +3,266 @@ module Streamly.External.LMDB.Tests (tests) where import Control.Concurrent.Async (asyncBound, wait)-import Control.Exception (SomeException, onException, try)+import Control.Exception (SomeException, bracket, onException, try) import Control.Monad (forM_) import Data.ByteString (ByteString, pack, unpack)+import qualified Data.ByteString as B import Data.ByteString.Unsafe (unsafeUseAsCStringLen) import Data.List (find, foldl', nubBy, sort) import Data.Word (Word8) import Foreign (castPtr, nullPtr, with)+import Streamly.External.LMDB+ ( Environment,+ Mode,+ OverwriteOptions (..),+ ReadDirection (..),+ ReadOnlyTxn,+ ReadOptions (..),+ ReadWrite,+ WriteOptions (..),+ abortReadOnlyTxn,+ beginReadOnlyTxn,+ clearDatabase,+ defaultReadOptions,+ defaultWriteOptions,+ readLMDB,+ unsafeReadLMDB,+ writeLMDB,+ )+import Streamly.External.LMDB.Internal (Database (..))+import Streamly.External.LMDB.Internal.Foreign+ ( MDB_val (..),+ combineOptions,+ mdb_nooverwrite,+ mdb_put,+ mdb_txn_begin,+ mdb_txn_commit,+ ) import Streamly.Prelude (fromList, toList, unfold)-import Test.QuickCheck (NonEmptyList (..), Gen, choose, elements, frequency)+import qualified Streamly.Prelude as S+import Test.QuickCheck (Gen, NonEmptyList (..), choose, elements, frequency) import Test.QuickCheck.Monadic (PropertyM, monadicIO, pick, run) import Test.Tasty (TestTree) import Test.Tasty.QuickCheck (arbitrary, testProperty) -import qualified Data.ByteString as B-import qualified Streamly.Prelude as S--import Streamly.External.LMDB (Mode, ReadWrite, OverwriteOptions (..), ReadDirection (..), ReadOptions (..),- WriteOptions(..), clearDatabase, defaultReadOptions, defaultWriteOptions, readLMDB, unsafeReadLMDB, writeLMDB)-import Streamly.External.LMDB.Internal (Database (..))-import Streamly.External.LMDB.Internal.Foreign (MDB_val (..), combineOptions,- mdb_nooverwrite, mdb_put, mdb_txn_begin, mdb_txn_commit)+tests :: IO (Database ReadWrite, Environment ReadWrite) -> [TestTree]+tests dbenv =+ [ testReadLMDB dbenv,+ testUnsafeReadLMDB dbenv,+ testWriteLMDB dbenv,+ testWriteLMDB_2 dbenv,+ testWriteLMDB_3 dbenv,+ testBetween+ ] -tests :: IO (Database ReadWrite) -> [TestTree]-tests db =- [ testReadLMDB db- , testUnsafeReadLMDB db- , testWriteLMDB db- , testWriteLMDB_2 db- , testWriteLMDB_3 db- , testBetween ]+withReadOnlyTxn :: (Mode mode) => Environment mode -> (ReadOnlyTxn -> IO r) -> IO r+withReadOnlyTxn env = bracket (beginReadOnlyTxn env) abortReadOnlyTxn -- | Clear the database, write key-value pairs to it in a normal manner, read -- them back using our library, and make sure the result is what we wrote.-testReadLMDB :: (Mode mode) => IO (Database mode) -> TestTree+testReadLMDB :: (Mode mode) => IO (Database mode, Environment mode) -> TestTree testReadLMDB res = testProperty "readLMDB" . monadicIO $ do- db <- run res- keyValuePairs <- arbitraryKeyValuePairs''- run $ clearDatabase db+ (db, env) <- run res+ keyValuePairs <- arbitraryKeyValuePairs''+ run $ clearDatabase db - run $ writeChunk db False keyValuePairs- let keyValuePairsInDb = sort . removeDuplicateKeys $ keyValuePairs+ run $ writeChunk db False keyValuePairs+ let keyValuePairsInDb = sort . removeDuplicateKeys $ keyValuePairs - (readOpts, expectedResults) <- pick $ readOptionsAndResults keyValuePairsInDb- results <- run . toList $ unfold (readLMDB db readOpts) undefined+ (readOpts, expectedResults) <- pick $ readOptionsAndResults keyValuePairsInDb+ let unf txn = toList $ unfold (readLMDB db txn readOpts) undefined+ results <- run $ unf Nothing+ resultsTxn <- run $ withReadOnlyTxn env (unf . Just) - return $ results == expectedResults+ return $ results == expectedResults && resultsTxn == expectedResults -- | Similar to 'testReadLMDB', except that it tests the unsafe function in a different manner.-testUnsafeReadLMDB :: (Mode mode) => IO (Database mode) -> TestTree+testUnsafeReadLMDB :: (Mode mode) => IO (Database mode, Environment mode) -> TestTree testUnsafeReadLMDB res = testProperty "unsafeReadLMDB" . monadicIO $ do- db <- run res- keyValuePairs <- arbitraryKeyValuePairs''- run $ clearDatabase db+ (db, env) <- run res+ keyValuePairs <- arbitraryKeyValuePairs''+ run $ clearDatabase db - run $ writeChunk db False keyValuePairs- let keyValuePairsInDb = sort . removeDuplicateKeys $ keyValuePairs+ run $ writeChunk db False keyValuePairs+ let keyValuePairsInDb = sort . removeDuplicateKeys $ keyValuePairs - (readOpts, expectedResults) <- pick $ readOptionsAndResults keyValuePairsInDb- let expectedLengths = map (\(k, v) -> (B.length k, B.length v)) expectedResults- lengths <- run . toList $ unfold (unsafeReadLMDB db readOpts (return . snd) (return . snd)) undefined+ (readOpts, expectedResults) <- pick $ readOptionsAndResults keyValuePairsInDb+ let expectedLengths = map (\(k, v) -> (B.length k, B.length v)) expectedResults+ let unf txn =+ toList $+ unfold (unsafeReadLMDB db txn readOpts (return . snd) (return . snd)) undefined+ lengths <- run $ unf Nothing+ lengthsTxn <- run $ withReadOnlyTxn env (unf . Just) - return $ lengths == expectedLengths+ return $ lengths == expectedLengths && lengthsTxn == expectedLengths --- | Clear the database, write key-value pairs to it using our library with key overwriting allowed, read--- them back using our library (already covered by 'testReadLMDB'), and make sure the result is what we wrote.-testWriteLMDB :: IO (Database ReadWrite) -> TestTree+-- | Clear the database, write key-value pairs to it using our library with key overwriting allowed,+-- read them back using our library (already covered by 'testReadLMDB'), and make sure the result is+-- what we wrote.+testWriteLMDB :: IO (Database ReadWrite, Environment ReadWrite) -> TestTree testWriteLMDB res = testProperty "writeLMDB" . monadicIO $ do- db <- run res- keyValuePairs <- arbitraryKeyValuePairs- run $ clearDatabase db+ (db, _) <- run res+ keyValuePairs <- arbitraryKeyValuePairs+ run $ clearDatabase db - chunkSz <- pick arbitrary+ chunkSz <- pick arbitrary - let fol' = writeLMDB db $ defaultWriteOptions { writeTransactionSize = chunkSz- , overwriteOptions = OverwriteAllow }+ let fol' =+ writeLMDB db $+ defaultWriteOptions+ { writeTransactionSize = chunkSz,+ overwriteOptions = OverwriteAllow+ } - -- TODO: Run with new "bound" functionality in streamly.- run $ asyncBound (S.fold fol' (fromList keyValuePairs)) >>= wait- let keyValuePairsInDb = sort . removeDuplicateKeys $ keyValuePairs+ -- TODO: Run with new "bound" functionality in streamly.+ run $ asyncBound (S.fold fol' (fromList keyValuePairs)) >>= wait+ let keyValuePairsInDb = sort . removeDuplicateKeys $ keyValuePairs - readPairsAll <- run . toList $ unfold (readLMDB db defaultReadOptions) undefined+ readPairsAll <- run . toList $ unfold (readLMDB db Nothing defaultReadOptions) undefined - return $ keyValuePairsInDb == readPairsAll+ return $ keyValuePairsInDb == readPairsAll -- | Clear the database, write key-value pairs to it using our library with key overwriting -- disallowed, and make sure an exception occurs iff we had a duplicate key in our pairs. -- Furthermore make sure that key-value pairs prior to a duplicate key are actually in the database.-testWriteLMDB_2 :: IO (Database ReadWrite) -> TestTree+testWriteLMDB_2 :: IO (Database ReadWrite, Environment ReadWrite) -> TestTree testWriteLMDB_2 res = testProperty "writeLMDB_2" . monadicIO $ do- db <- run res- keyValuePairs <- arbitraryKeyValuePairs'- run $ clearDatabase db+ (db, _) <- run res+ keyValuePairs <- arbitraryKeyValuePairs'+ run $ clearDatabase db - chunkSz <- pick arbitrary+ chunkSz <- pick arbitrary - -- TODO: Run with new "bound" functionality in streamly.- let fol' = writeLMDB db $ defaultWriteOptions { writeTransactionSize = chunkSz- , overwriteOptions = OverwriteDisallow }- e <- run $ try @SomeException $ (asyncBound (S.fold fol' (fromList keyValuePairs)) >>= wait)- exceptionAsExpected <-- case e of- Left _ -> return $ hasDuplicateKeys keyValuePairs- Right _ -> return . not $ hasDuplicateKeys keyValuePairs+ -- TODO: Run with new "bound" functionality in streamly.+ let fol' =+ writeLMDB db $+ defaultWriteOptions+ { writeTransactionSize = chunkSz,+ overwriteOptions = OverwriteDisallow+ }+ e <- run $ try @SomeException $ (asyncBound (S.fold fol' (fromList keyValuePairs)) >>= wait)+ exceptionAsExpected <-+ case e of+ Left _ -> return $ hasDuplicateKeys keyValuePairs+ Right _ -> return . not $ hasDuplicateKeys keyValuePairs - let keyValuePairsInDb = sort . prefixBeforeDuplicate $ keyValuePairs- readPairsAll <- run . toList $ unfold (readLMDB db defaultReadOptions) undefined- let pairsAsExpected = keyValuePairsInDb == readPairsAll+ let keyValuePairsInDb = sort . prefixBeforeDuplicate $ keyValuePairs+ readPairsAll <- run . toList $ unfold (readLMDB db Nothing defaultReadOptions) undefined+ let pairsAsExpected = keyValuePairsInDb == readPairsAll - return $ exceptionAsExpected && pairsAsExpected+ return $ exceptionAsExpected && pairsAsExpected -- | Clear the database, write key-value pairs to it using our library with key overwriting -- disallowed except when attempting to replace an existing key-value pair, and make sure an--- exception occurs iff we had a duplicate key with different values in our pairs. Furthermore--- make sure that key-value pairs prior to a such a duplicate key are actually in the database.-testWriteLMDB_3 :: IO (Database ReadWrite) -> TestTree+-- exception occurs iff we had a duplicate key with different values in our pairs. Furthermore make+-- sure that key-value pairs prior to a such a duplicate key are actually in the database.+testWriteLMDB_3 :: IO (Database ReadWrite, Environment ReadWrite) -> TestTree testWriteLMDB_3 res = testProperty "writeLMDB_3" . monadicIO $ do- db <- run res- keyValuePairs <- arbitraryKeyValuePairs'- run $ clearDatabase db+ (db, _) <- run res+ keyValuePairs <- arbitraryKeyValuePairs'+ run $ clearDatabase db - chunkSz <- pick arbitrary+ chunkSz <- pick arbitrary - -- TODO: Run with new "bound" functionality in streamly.- let fol' = writeLMDB db $ defaultWriteOptions { writeTransactionSize = chunkSz- , overwriteOptions = OverwriteAllowSame }- e <- run $ try @SomeException $ (asyncBound (S.fold fol' (fromList keyValuePairs)) >>= wait)- exceptionAsExpected <-- case e of- Left _ -> return $ hasDuplicateKeysWithDiffVals keyValuePairs- Right _ -> return . not $ hasDuplicateKeysWithDiffVals keyValuePairs+ -- TODO: Run with new "bound" functionality in streamly.+ let fol' =+ writeLMDB db $+ defaultWriteOptions+ { writeTransactionSize = chunkSz,+ overwriteOptions = OverwriteAllowSame+ }+ e <- run $ try @SomeException $ (asyncBound (S.fold fol' (fromList keyValuePairs)) >>= wait)+ exceptionAsExpected <-+ case e of+ Left _ -> return $ hasDuplicateKeysWithDiffVals keyValuePairs+ Right _ -> return . not $ hasDuplicateKeysWithDiffVals keyValuePairs - let keyValuePairsInDb = sort . removeDuplicateKeys . prefixBeforeDuplicateWithDiffVal $ keyValuePairs- readPairsAll <- run . toList $ unfold (readLMDB db defaultReadOptions) undefined- let pairsAsExpected = keyValuePairsInDb == readPairsAll+ let keyValuePairsInDb =+ sort . removeDuplicateKeys . prefixBeforeDuplicateWithDiffVal $+ keyValuePairs+ readPairsAll <- run . toList $ unfold (readLMDB db Nothing defaultReadOptions) undefined+ let pairsAsExpected = keyValuePairsInDb == readPairsAll - return $ exceptionAsExpected && pairsAsExpected+ return $ exceptionAsExpected && pairsAsExpected arbitraryKeyValuePairs :: PropertyM IO [(ByteString, ByteString)] arbitraryKeyValuePairs =- map (\(ws1, ws2) -> (pack ws1, pack ws2))+ map (\(ws1, ws2) -> (pack ws1, pack ws2)) . filter (\(ws1, _) -> not (null ws1)) -- LMDB does not allow empty keys.- <$> pick arbitrary+ <$> pick arbitrary -- A variation that makes duplicate keys more likely. arbitraryKeyValuePairs' :: PropertyM IO [(ByteString, ByteString)] arbitraryKeyValuePairs' = do- arb <- arbitraryKeyValuePairs- b <- pick arbitrary- if not (null arb) && b then do- let (k, v) = head arb- b' <- pick arbitrary- v' <- if b' then return v else pack <$> pick arbitrary- i <- pick $ choose (negate $ length arb, 2 * length arb)- let (arb1, arb2) = splitAt i arb- let arb' = arb1 ++ [(k, v')] ++ arb2- return arb'- else- return arb+ arb <- arbitraryKeyValuePairs+ b <- pick arbitrary+ if not (null arb) && b+ then do+ let (k, v) = head arb+ b' <- pick arbitrary+ v' <- if b' then return v else pack <$> pick arbitrary+ i <- pick $ choose (negate $ length arb, 2 * length arb)+ let (arb1, arb2) = splitAt i arb+ let arb' = arb1 ++ [(k, v')] ++ arb2+ return arb'+ else return arb --- A variation that makes more likely keys with same the prefix and a difference of trailing zero bytes.+-- A variation that makes more likely keys with same the prefix and a difference of trailing zero+-- bytes. arbitraryKeyValuePairs'' :: PropertyM IO [(ByteString, ByteString)] arbitraryKeyValuePairs'' = do- arb <- arbitraryKeyValuePairs- if null arb then- return arb- else pick $ frequency [(1, return arb), (3, do- let (k, v) = head arb- b' <- arbitrary- v' <- if b' then return v else pack <$> arbitrary- i <- choose (0, length arb - 1)- let (arb1, arb2) = splitAt i arb- let arb3 = map (\i' -> (k `B.append` B.replicate i' 0, v')) [1..(i+1)]- let arb' = arb1 ++ arb3 ++ arb2- return arb' )]+ arb <- arbitraryKeyValuePairs+ if null arb+ then return arb+ else+ pick $+ frequency+ [ (1, return arb),+ ( 3,+ do+ let (k, v) = head arb+ b' <- arbitrary+ v' <- if b' then return v else pack <$> arbitrary+ i <- choose (0, length arb - 1)+ let (arb1, arb2) = splitAt i arb+ let arb3 = map (\i' -> (k `B.append` B.replicate i' 0, v')) [1 .. (i + 1)]+ let arb' = arb1 ++ arb3 ++ arb2+ return arb'+ )+ ] -- | Note that this function retains the last value for each key. removeDuplicateKeys :: (Eq a) => [(a, b)] -> [(a, b)]-removeDuplicateKeys = foldl' (\acc (a, b) -> if any ((== a) . fst) acc then acc else (a, b) : acc) [] . reverse+removeDuplicateKeys =+ foldl' (\acc (a, b) -> if any ((== a) . fst) acc then acc else (a, b) : acc) [] . reverse hasDuplicateKeys :: (Eq a) => [(a, b)] -> Bool hasDuplicateKeys l =- let l2 = nubBy (\(a1, _) (a2, _) -> a1 == a2) l- in length l /= length l2+ let l2 = nubBy (\(a1, _) (a2, _) -> a1 == a2) l+ in length l /= length l2 hasDuplicateKeysWithDiffVals :: (Eq a, Eq b) => [(a, b)] -> Bool hasDuplicateKeysWithDiffVals l =- let l2 = nubBy (\(a1, b1) (a2, b2) -> a1 == a2 && b1 /= b2) l- in length l /= length l2+ let l2 = nubBy (\(a1, b1) (a2, b2) -> a1 == a2 && b1 /= b2) l+ in length l /= length l2 prefixBeforeDuplicate :: (Eq a) => [(a, b)] -> [(a, b)] prefixBeforeDuplicate xs =- let fstDup = snd <$> find (\((a, _), i) -> a `elem` map fst (take i xs)) (zip xs [0..])- in case fstDup of+ let fstDup = snd <$> find (\((a, _), i) -> a `elem` map fst (take i xs)) (zip xs [0 ..])+ in case fstDup of Nothing -> xs Just i -> take i xs prefixBeforeDuplicateWithDiffVal :: (Eq a, Eq b) => [(a, b)] -> [(a, b)] prefixBeforeDuplicateWithDiffVal xs =- let fstDup = snd <$> find (\((a, b), i) -> any (\(a', b') -> a == a' && b /= b') (take i xs)) (zip xs [0..])- in case fstDup of+ let fstDup =+ snd+ <$> find+ ( \((a, b), i) ->+ any (\(a', b') -> a == a' && b /= b') (take i xs)+ )+ (zip xs [0 ..])+ in case fstDup of Nothing -> xs Just i -> take i xs @@ -210,84 +270,102 @@ between :: [Word8] -> [Word8] -> [Word8] -> Maybe [Word8] between [] [] _ = error "first = second" between _ [] _ = error "first > second"-between [] (w:ws) commonPrefixRev- | w == 0 && null ws = Nothing- | w == 0 = between [] ws (w:commonPrefixRev)- | otherwise = Just $ reverse (0:commonPrefixRev)-between (w1:ws1) (w2:ws2) commonPrefixRev- | w1 == w2 = between ws1 ws2 (w1:commonPrefixRev)- | w1 > w2 = error "first > second"- | otherwise = Just $ reverse commonPrefixRev ++ [w1] ++ ws1 ++ [0]+between [] (w : ws) commonPrefixRev+ | w == 0 && null ws = Nothing+ | w == 0 = between [] ws (w : commonPrefixRev)+ | otherwise = Just $ reverse (0 : commonPrefixRev)+between (w1 : ws1) (w2 : ws2) commonPrefixRev+ | w1 == w2 = between ws1 ws2 (w1 : commonPrefixRev)+ | w1 > w2 = error "first > second"+ | otherwise = Just $ reverse commonPrefixRev ++ [w1] ++ ws1 ++ [0] testBetween :: TestTree testBetween = testProperty "testBetween" $ \ws1 ws2 ->- (ws1 == ws2) ||- let (smaller, bigger) = if ws1 < ws2 then (ws1, ws2) else (ws2, ws1)+ (ws1 == ws2)+ || let (smaller, bigger) = if ws1 < ws2 then (ws1, ws2) else (ws2, ws1) in case between smaller bigger [] of- Nothing -> drop (length ws1) ws2 == replicate (length ws2 - length ws1) 0- Just betw -> smaller < betw && betw < bigger+ Nothing -> drop (length ws1) ws2 == replicate (length ws2 - length ws1) 0+ Just betw -> smaller < betw && betw < bigger betweenBs :: ByteString -> ByteString -> Maybe ByteString betweenBs bs1 bs2 = between (unpack bs1) (unpack bs2) [] >>= (return . pack) type PairsInDatabase = [(ByteString, ByteString)]+ type ExpectedReadResult = [(ByteString, ByteString)] -- | Given database pairs, randomly generates read options and corresponding expected results. readOptionsAndResults :: PairsInDatabase -> Gen (ReadOptions, ExpectedReadResult) readOptionsAndResults pairsInDb = do- forw <- arbitrary- let dir = if forw then Forward else Backward- let len = length pairsInDb- readAll <- frequency [ (1, return True), (3, return False) ]- let ropts = defaultReadOptions { readDirection = dir }- if readAll then- return (ropts { readStart = Nothing }, (if forw then id else reverse) pairsInDb)- else if len == 0 then do- bs <- arbitrary >>= \(NonEmpty ws) -> return $ pack ws- return (ropts { readStart = Just bs }, [])- else do- idx <- if len < 3 then choose (0, len - 1) else- frequency [ (1, choose (1, len - 2)), (3, elements [0, len - 1]) ]- let keyAt i = fst $ pairsInDb !! i- let nextKey+ forw <- arbitrary+ let dir = if forw then Forward else Backward+ let len = length pairsInDb+ readAll <- frequency [(1, return True), (3, return False)]+ let ropts = defaultReadOptions {readDirection = dir}+ if readAll+ then return (ropts {readStart = Nothing}, (if forw then id else reverse) pairsInDb)+ else+ if len == 0+ then do+ bs <- arbitrary >>= \(NonEmpty ws) -> return $ pack ws+ return (ropts {readStart = Just bs}, [])+ else do+ idx <-+ if len < 3+ then choose (0, len - 1)+ else frequency [(1, choose (1, len - 2)), (3, elements [0, len - 1])]+ let keyAt i = fst $ pairsInDb !! i+ let nextKey | idx + 1 <= len - 1 = betweenBs (keyAt idx) (keyAt $ idx + 1) | otherwise = Just $ keyAt (len - 1) `B.append` B.singleton 0- let prevKey- | idx == 0 && keyAt idx /= B.singleton 0 = Just $ B.singleton 0 -- Keys are known to be non-empty.+ let prevKey+ -- Keys are known to be non-empty.+ | idx == 0 && keyAt idx /= B.singleton 0 = Just $ B.singleton 0 | idx == 0 = Nothing | otherwise = betweenBs (keyAt $ idx - 1) (keyAt idx)- let forwEq = (ropts { readStart = Just $ keyAt idx }, drop idx pairsInDb)- let backwEq = (ropts { readStart = Just $ keyAt idx }, reverse $ take (idx + 1) pairsInDb)- ord <- arbitrary @Ordering -- Proximity to the key at idx (if possible).- return $ case (ord, dir) of+ let forwEq = (ropts {readStart = Just $ keyAt idx}, drop idx pairsInDb)+ let backwEq = (ropts {readStart = Just $ keyAt idx}, reverse $ take (idx + 1) pairsInDb)+ ord <- arbitrary @Ordering -- Proximity to the key at idx (if possible).+ return $ case (ord, dir) of (EQ, Forward) -> forwEq (EQ, Backward) -> backwEq (GT, Forward) -> case nextKey of- Nothing -> forwEq- Just nextKey' -> (ropts { readStart = Just nextKey' }, drop (idx + 1) pairsInDb)+ Nothing -> forwEq+ Just nextKey' -> (ropts {readStart = Just nextKey'}, drop (idx + 1) pairsInDb) (GT, Backward) -> case nextKey of- Nothing -> backwEq- Just nextKey' -> (ropts { readStart = Just nextKey' }, reverse $ take (idx + 1) pairsInDb)+ Nothing -> backwEq+ Just nextKey' ->+ (ropts {readStart = Just nextKey'}, reverse $ take (idx + 1) pairsInDb) (LT, Forward) -> case prevKey of- Nothing -> forwEq- Just prevKey' -> (ropts { readStart = Just prevKey' }, drop idx pairsInDb)+ Nothing -> forwEq+ Just prevKey' -> (ropts {readStart = Just prevKey'}, drop idx pairsInDb) (LT, Backward) -> case prevKey of- Nothing -> backwEq- Just prevKey' -> (ropts { readStart = Just prevKey' }, reverse $ take idx pairsInDb)+ Nothing -> backwEq+ Just prevKey' -> (ropts {readStart = Just prevKey'}, reverse $ take idx pairsInDb) -- Writes the given key-value pairs to the given database.-writeChunk :: (Foldable t, Mode mode) => Database mode -> Bool -> t (ByteString, ByteString) -> IO ()+writeChunk ::+ (Foldable t, Mode mode) =>+ Database mode ->+ Bool ->+ t (ByteString, ByteString) ->+ IO () writeChunk (Database penv dbi) noOverwrite' keyValuePairs =- let flags = combineOptions $ [mdb_nooverwrite | noOverwrite']- in asyncBound (do- ptxn <- mdb_txn_begin penv nullPtr 0- onException (forM_ keyValuePairs $ \(k, v) ->- marshalOut k $ \k' -> marshalOut v $ \v' -> with k' $ \k'' -> with v' $ \v'' ->- mdb_put ptxn dbi k'' v'' flags)- (mdb_txn_commit ptxn) -- Make sure the key-value pairs we have so far are committed.- mdb_txn_commit ptxn) >>= wait+ let flags = combineOptions $ [mdb_nooverwrite | noOverwrite']+ in asyncBound+ ( do+ ptxn <- mdb_txn_begin penv nullPtr 0+ onException+ ( forM_ keyValuePairs $ \(k, v) ->+ marshalOut k $ \k' -> marshalOut v $ \v' -> with k' $ \k'' -> with v' $ \v'' ->+ mdb_put ptxn dbi k'' v'' flags+ )+ (mdb_txn_commit ptxn) -- Make sure the key-value pairs we have so far are committed.+ mdb_txn_commit ptxn+ )+ >>= wait {-# INLINE marshalOut #-} marshalOut :: ByteString -> (MDB_val -> IO ()) -> IO ()-marshalOut bs f = unsafeUseAsCStringLen bs $ \(ptr, len) -> f $ MDB_val (fromIntegral len) (castPtr ptr)+marshalOut bs f =+ unsafeUseAsCStringLen bs $ \(ptr, len) -> f $ MDB_val (fromIntegral len) (castPtr ptr)
test/TestSuite.hs view
@@ -1,27 +1,40 @@ module Main where +import Streamly.External.LMDB+ ( Database,+ Environment,+ ReadWrite,+ defaultLimits,+ getDatabase,+ mapSize,+ openEnvironment,+ tebibyte,+ )+import qualified Streamly.External.LMDB.Tests (tests) import System.Directory (removeDirectoryRecursive) import System.Environment (setEnv) import System.IO.Temp (createTempDirectory, getCanonicalTemporaryDirectory) import Test.Tasty (TestTree, defaultMain, testGroup, withResource) -import Streamly.External.LMDB (Database, ReadWrite, defaultLimits, getDatabase, mapSize, openEnvironment, tebibyte)-import qualified Streamly.External.LMDB.Tests (tests)- main :: IO () main = do- setEnv "TASTY_NUM_THREADS" "1" -- Multiple tests use the same LMDB database.- defaultMain $ withResource- (do tmpParent <- getCanonicalTemporaryDirectory- tmpDir <- createTempDirectory tmpParent "streamly-lmdb-tests"- env <- openEnvironment tmpDir $ defaultLimits { mapSize = tebibyte }- db <- getDatabase env Nothing- return (tmpDir, db))- (\(tmpDir, _) -> removeDirectoryRecursive tmpDir)- (\io -> tests $ snd <$> io)+ setEnv "TASTY_NUM_THREADS" "1" -- Multiple tests use the same LMDB database.+ defaultMain $+ withResource+ ( do+ tmpParent <- getCanonicalTemporaryDirectory+ tmpDir <- createTempDirectory tmpParent "streamly-lmdb-tests"+ env <- openEnvironment tmpDir $ defaultLimits {mapSize = tebibyte}+ db <- getDatabase env Nothing+ return (tmpDir, (db, env))+ )+ (\(tmpDir, _) -> removeDirectoryRecursive tmpDir)+ (\io -> tests $ snd <$> io) -tests :: IO (Database ReadWrite) -> TestTree+tests :: IO (Database ReadWrite, Environment ReadWrite) -> TestTree tests res =- testGroup "Tests"- [ testGroup "Streamly.External.LMDB.Tests" $- Streamly.External.LMDB.Tests.tests res ]+ testGroup+ "Tests"+ [ testGroup "Streamly.External.LMDB.Tests" $+ Streamly.External.LMDB.Tests.tests res+ ]