diff --git a/ChangeLog.md b/ChangeLog.md --- a/ChangeLog.md +++ b/ChangeLog.md @@ -1,3 +1,20 @@ +## 0.8.0 + +* Updated for Streamly 0.10.0. +* Added proper concurrency support, including safety upon asynchronous exceptions. (The previous versions were unreliable in concurrent settings—segfaults, deadlocks, etc.) +* A byproduct of the previous point: `writeLMDB` no longer maintains its own read-write transactions under the hood. +* Added the ability to chunk the read-only transactions (by bytes or number of pairs) that `readLMDB` can create under the hood. +* Added the ability to use read-write transactions for `readLMDB`. +* Eliminated the `Void` seed for `readLMDB`. +* Added more choices for `readStart`. +* Added the ability to stream unexpected key-value pairs in `writeLMDB` into a separate fold. +* Added `getLMDB` for getting a single key-value pair normally (without streaming). +* Added `writeLMDBChunk` for writing key-value pairs normally (without streaming). +* Added `chunkPairs`/`chunkPairsFold` convenience functions that allow chunking a stream of key-value pairs (by bytes or number of pairs). +* Added `deleteLMDB` for deleting key-value pairs. +* Made the unsafe FFI functionality internal. +* Improved documentation, esp. regarding caveats for long-lived transactions and various lower-level LMDB requirements. + ## 0.7.0 * Added `readUnsafeFFI` and `writeUnsafeFFI` options. diff --git a/LICENSE b/LICENSE --- a/LICENSE +++ b/LICENSE @@ -1,4 +1,4 @@ -Copyright Shlok Datye (c) 2023 +Copyright Shlok Datye (c) 2024 All rights reserved. diff --git a/README.md b/README.md --- a/README.md +++ b/README.md @@ -19,28 +19,18 @@ module Main where -import Data.Function ((&)) +import Data.Function import qualified Streamly.Data.Fold as F import qualified Streamly.Data.Stream.Prelude as S import Streamly.External.LMDB - ( Limits (mapSize), - WriteOptions (writeTransactionSize), - defaultLimits, - defaultReadOptions, - defaultWriteOptions, - getDatabase, - openEnvironment, - readLMDB, - tebibyte, - writeLMDB, - ) 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" $ + openEnvironment + "/path/to/lmdb-database" defaultLimits {mapSize = tebibyte} -- Get the main database. @@ -49,9 +39,10 @@ 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 + withReadWriteTransaction env $ \txn -> + [("baz", "a"), ("foo", "b"), ("bar", "c")] + & S.fromList + & S.fold (writeLMDB defaultWriteOptions db txn) -- Stream key-value pairs out of the -- database, printing them along the way. @@ -59,9 +50,8 @@ -- ("bar","c") -- ("baz","a") -- ("foo","b") - let unfold' = readLMDB db Nothing defaultReadOptions - let readStream = S.unfold unfold' undefined - S.mapM print readStream + S.unfold readLMDB (defaultReadOptions, db, LeftTxn Nothing) + & S.mapM print & S.fold F.drain ``` @@ -69,15 +59,12 @@ See `bench/README.md`. Summary (with rough figures from our machine†): -* Reading: - - For iterating through a fully cached LMDB database, this library has roughly a 110 ns/pair overhead compared to C. (Plain Haskell `IO` code has roughly a 70 ns/pair overhead compared to C. The two preceding figures being similar fulfills the promise of [streamly](https://hackage.haskell.org/package/streamly) and stream fusion.) - - By using `unsafeReadLMDB` instead of `readLMDB`, we can get the overhead down to roughly 100 ns/pair. - - By additionally using the `readUnsafeFFI` option (to use `unsafe` FFI calls under the hood), we can get the overhead down to roughly 40 ns/pair. - +* Reading (iterating through a fully cached LMDB database): + - When using the ordinary `readLMDB` (which creates intermediate key/value `ByteString`s managed by the RTS), the overhead compared to C depends on the key/value sizes; for 480-byte keys and 2400-byte values, the overhead is roughly 815 ns/pair. + - By using `unsafeReadLMDB` instead of `readLMDB` (to avoid the intermediate `ByteString`s), we can get the overhead compared to C down to roughly 90 ns/pair. (Plain Haskell `IO` code has roughly a 50 ns/pair overhead compared to C. The two preceding figures being similar fulfills the promise of [streamly](https://hackage.haskell.org/package/streamly) and stream fusion.) * Writing: - - For writing to an LMDB database, this library has roughly a 210 ns/pair overhead compared to C. (Plain Haskell `IO` code has roughly a 100 ns/pair overhead compared to C. The two preceding figures being similar fulfills the promise of [streamly](https://hackage.haskell.org/package/streamly) and stream fusion.) - - By using the `writeUnsafeFFI` option (to use `unsafe` FFI calls under the hood), we can get the overhead down to roughly 140 ns/pair. - -* For most Haskell programs, these differences will not cause problems. (For instance, note that merely opening and reading 1 byte from a file with C already takes us tens of *microseconds*.) + - The overhead of this library compared to C depends on the size of the key/value pairs (`ByteString`s managed by the RTS). For 480-byte keys and 2400-byte values, the overhead is around 4.3 μs/pair. + - For now, we don’t provide “unsafe” write functionality (to avoid the key/value `ByteString`s) because this write performance is currently good enough for our purposes. +* For reference, we note that opening and reading 1 byte [16 KiB] from a file on disk with C takes us around 2.8 μs [20 μs]. -† May 2023; [Linode](https://linode.com); Debian 11, Dedicated 32GB: 16 CPU, 640GB SSD storage, 32GB RAM. +† September 2024; NixOS 24.11; Intel i7-12700K (3.6 GHz, 12 cores); Corsair VENGEANCE LPX DDR4 RAM 64GB (2 x 32GB) 3200MHz; Samsung 970 EVO Plus SSD 2TB (M.2 NVMe). diff --git a/src/Streamly/External/LMDB.hs b/src/Streamly/External/LMDB.hs --- a/src/Streamly/External/LMDB.hs +++ b/src/Streamly/External/LMDB.hs @@ -1,552 +1,138 @@ -{-# LANGUAGE BangPatterns #-} -{-# LANGUAGE ScopedTypeVariables #-} - --- | = 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 - ( -- * Environment + ( -- * 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. + + -- * Warning + + -- | Unless you know what you’re doing, please do not use other mechanisms (in addition to the + -- public functionality of this library) in the same Haskell program to interact with your LMDB + -- databases. (If you really want to do this, you should heed all the low-level requirements we + -- have linked to in the source code of this library, and in general understand how the LMDB C + -- API works with @MDB_NOTLS@ enabled.) + + -- * Environments + -- | With LMDB, one first creates a so-called “environment,” which one can think of as a file or - -- folder on disk. + -- directory on disk. Environment, openEnvironment, isReadOnlyEnvironment, closeEnvironment, - -- ** Mode - Mode, - ReadWrite, - ReadOnly, - -- ** Limits Limits (..), defaultLimits, - gibibyte, - tebibyte, - -- * Database + -- * Databases #databases# -- | After creating an environment, one creates within it one or more databases. Database, getDatabase, - clearDatabase, closeDatabase, - -- * Reading - readLMDB, - unsafeReadLMDB, + -- * Transactions #transactions# - -- ** Read-only transactions and cursors - ReadOnlyTxn, - beginReadOnlyTxn, - abortReadOnlyTxn, + -- | In LMDB, there are two types of transactions: read-only transactions and read-write + -- transactions. On a given environment, read-only transactions do not block other transactions + -- and read-write transactions do not block read-only transactions, but read-write transactions + -- are serialized and block other read-write transactions. + -- + -- Read-only transactions attain a snapshot view of the environment; this view is not affected + -- by newer read-write transactions. + -- + -- /Warning/: Long-lived transactions are discouraged by LMDB, and it is your responsibility as + -- a user of this library to avoid them as necessary. The reasons are twofold: (a) The first one + -- we already mentioned: Read-write transactions block other read-write transactions. (b) The + -- second is more insidious: Even though read-only transactions do not block read-write + -- transactions, read-only transactions (since they attain a snapshot view of the environment) + -- prevent the reuse of pages freed by newer read-write transactions, so the database can grow + -- quickly. + Transaction (), + + -- ** Read-only transactions + beginReadOnlyTransaction, + abortReadOnlyTransaction, + withReadOnlyTransaction, + waitReaders, + + -- ** Read-write transactions + beginReadWriteTransaction, + abortReadWriteTransaction, + commitReadWriteTransaction, + withReadWriteTransaction, + + -- * Cursors Cursor, openCursor, closeCursor, + withCursor, - -- ** Read options + -- * Reading + + -- ** Stream-based reading + readLMDB, + unsafeReadLMDB, ReadOptions (..), defaultReadOptions, ReadDirection (..), + ReadStart (..), + -- ** Direct reading + getLMDB, + -- * Writing + + -- ** Stream-based writing writeLMDB, WriteOptions (..), defaultWriteOptions, OverwriteOptions (..), - -- * Error types - LMDB_Error (..), - MDB_ErrCode (..), - ) -where - -import Control.Concurrent (isCurrentThreadBound, myThreadId) -import Control.Concurrent.Async (asyncBound, wait) -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, 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 -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 (lmap) -import Streamly.Internal.Data.Unfold.Type (Unfold (Unfold)) - -newtype Environment mode = Environment (Ptr MDB_env) - -isReadOnlyEnvironment :: Mode mode => Environment mode -> Bool -isReadOnlyEnvironment = isReadOnlyMode . mode - where - mode :: Environment mode -> mode - mode = undefined - --- | LMDB environments have various limits on the size and number of databases and concurrent --- readers. -data Limits = Limits - { -- | 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 '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. -defaultLimits :: Limits -defaultLimits = - Limits - { mapSize = 1024 * 1024, -- 1 MiB. - maxDatabases = 0, - maxReaders = 126 - } - --- | 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. -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. --- --- 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). -openEnvironment :: Mode mode => FilePath -> Limits -> IO (Environment mode) -openEnvironment path limits = do - 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) - - -- 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 - - 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 - --- | Closes the given environment. --- --- If you have merely a few dozen environments at most, there should be no need for this. (It is a --- common practice with LMDB to create one’s environments once and reuse them for the remainder of --- the program’s execution.) If you find yourself needing this, it is your responsibility to heed --- the [documented --- caveats](https://github.com/LMDB/lmdb/blob/8d0cbbc936091eb85972501a9b31a8f86d4c51a7/libraries/liblmdb/lmdb.h#L787). --- --- In particular, you will probably, before calling this function, want to (a) use 'closeDatabase', --- and (b) pass in precreated transactions and cursors to 'readLMDB' and 'unsafeReadLMDB' to make --- sure there are no transactions or cursors still left to be cleaned up by the garbage collector. --- (As an alternative to (b), one could try manually triggering the garbage collector.) -closeEnvironment :: (Mode mode) => Environment mode -> IO () -closeEnvironment (Environment penv) = - c_mdb_env_close penv - --- | 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 - --- | 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 - --- | Closes the given database. --- --- If you have merely a few dozen databases at most, there should be no need for this. (It is a --- common practice with LMDB to create one’s databases once and reuse them for the remainder of the --- program’s execution.) If you find yourself needing this, it is your responsibility to heed the --- [documented --- caveats](https://github.com/LMDB/lmdb/blob/8d0cbbc936091eb85972501a9b31a8f86d4c51a7/libraries/liblmdb/lmdb.h#L1200). -closeDatabase :: (Mode mode) => Database mode -> IO () -closeDatabase (Database penv dbi) = - c_mdb_dbi_close penv dbi - --- | Creates an unfold with which we can stream key-value pairs from the given database. --- --- If an existing read-only transaction and cursor are not provided, a read-only transaction and --- cursor are 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 and cursors --- are not being garbage collected fast enough), consider precreating a transaction and cursor using --- 'beginReadOnlyTxn' and 'openCursor'. --- --- 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 -> - Maybe (ReadOnlyTxn, Cursor) -> - ReadOptions -> - Unfold m Void (ByteString, ByteString) -readLMDB db mtxncurs ropts = unsafeReadLMDB db mtxncurs ropts packCStringLen packCStringLen - --- | 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'. -{-# INLINE unsafeReadLMDB #-} -unsafeReadLMDB :: - (MonadIO m, Mode mode) => - Database mode -> - Maybe (ReadOnlyTxn, Cursor) -> - ReadOptions -> - (CStringLen -> IO k) -> - (CStringLen -> IO v) -> - Unfold m Void (k, v) -unsafeReadLMDB (Database penv dbi) mtxncurs 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) - (txn_begin, cursor_open, cursor_get, cursor_close, txn_abort) = - if readUnsafeFFI ropts - then - ( mdb_txn_begin_unsafe, - mdb_cursor_open_unsafe, - c_mdb_cursor_get_unsafe, - c_mdb_cursor_close_unsafe, - c_mdb_txn_abort_unsafe - ) - else - ( mdb_txn_begin, - mdb_cursor_open, - c_mdb_cursor_get, - c_mdb_cursor_close, - c_mdb_txn_abort - ) - - supply = lmap . const - in supply firstOp $ - Unfold - ( \(op, pcurs, pk, pv, ref) -> do - rc <- - liftIO $ - if op == mdb_set_range && subsequentOp == mdb_prev - then do - -- A “reverse MDB_SET_RANGE” (i.e., a “less than or equal to”) is not - -- available in LMDB, so we simulate it ourselves. - kfst' <- peek pk - kfst <- packCStringLen (mv_data kfst', fromIntegral $ mv_size kfst') - rc <- cursor_get pcurs pk pv op - if rc /= 0 && rc == mdb_notfound - then 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 cursor_get pcurs pk pv mdb_prev - else return rc - else return rc - else 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 - - 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, pcurs) <- case mtxncurs of - Nothing -> do - ptxn <- txn_begin penv nullPtr mdb_rdonly - pcurs <- cursor_open ptxn dbi - return (ptxn, pcurs) - Just (ReadOnlyTxn ptxn, Cursor pcurs) -> - return (ptxn, pcurs) - pk <- malloc - pv <- malloc - - _ <- case readStart ropts of - Nothing -> return () - Just k -> unsafeUseAsCStringLen k $ \(kp, kl) -> - poke pk (MDB_val (fromIntegral kl) kp) - - ref <- newIOFinalizer $ do - free pv >> free pk - when (isNothing mtxncurs) $ - -- With LMDB, there is ordinarily no need to commit read-only transactions. - -- (The exception is when we want to make databases that were opened during - -- the transaction available later, but that’s not applicable here.) We can - -- therefore abort ptxn, both for failure (exceptions) and success. - cursor_close pcurs >> 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 - -newtype Cursor = Cursor (Ptr MDB_cursor) - --- | Opens a cursor for use with 'readLMDB' or 'unsafeReadLMDB'. It is your responsibility to (a) --- make sure the cursor only gets used by a single 'readLMDB' or 'unsafeReadLMDB' @Unfold@ at the --- same time (to be safe, one can open a new cursor for every 'readLMDB' or 'unsafeReadLMDB' call), --- (b) make sure the provided database is within the environment on which the provided transaction --- was begun, and (c) dispose of the cursor with 'closeCursor' (logically before 'abortReadOnlyTxn', --- although the order doesn’t really matter for read-only transactions). -openCursor :: ReadOnlyTxn -> Database mode -> IO Cursor -openCursor (ReadOnlyTxn ptxn) (Database _ dbi) = - Cursor <$> mdb_cursor_open ptxn dbi - --- | Disposes of a cursor created with 'openCursor'. -closeCursor :: Cursor -> IO () -closeCursor (Cursor pcurs) = - c_mdb_cursor_close pcurs - -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), - -- | Use @unsafe@ FFI calls under the hood. This can increase iteration speed, but one should - -- bear in mind that @unsafe@ FFI calls can have an adverse impact on the performance of the - -- rest of the program (e.g., its ability to effectively spawn green threads). - readUnsafeFFI :: !Bool - } - deriving (Show) - --- | By default, we start reading from the beginning of the database (i.e., from the smallest key), --- and we don’t use unsafe FFI calls. -defaultReadOptions :: ReadOptions -defaultReadOptions = - ReadOptions - { readDirection = Forward, - readStart = Nothing, - readUnsafeFFI = False - } - --- | Direction of key iteration. -data ReadDirection = Forward | Backward deriving (Show) - -data OverwriteOptions - = -- | When a key reoccurs, overwrite the value. - OverwriteAllow - | -- | When a key reoccurs, throw an exception except when the value is the same. - OverwriteAllowSame - | -- | When a key reoccurs, throw an exception. - OverwriteDisallow - deriving (Eq) - -data WriteOptions = WriteOptions - { -- | The number of key-value pairs per write transaction. - writeTransactionSize :: !Int, - writeOverwriteOptions :: !OverwriteOptions, - -- | Assume the input data is already ordered. This allows the use of @MDB_APPEND@ under the - -- hood and substantially improves write performance. An exception will be thrown if the - -- assumption about the ordering is not true. - writeAppend :: !Bool, - -- | Use @unsafe@ FFI calls under the hood. This can increase write performance, but one should - -- bear in mind that @unsafe@ FFI calls can have an adverse impact on the performance of the - -- rest of the program (e.g., its ability to effectively spawn green threads). - writeUnsafeFFI :: !Bool - } + -- *** Accumulators --- | By default, we use a write transaction size of 1 (one write transaction for each key-value --- pair), allow overwriting, don’t assume that the input data is already ordered, and don’t use --- unsafe FFI calls. -defaultWriteOptions :: WriteOptions -defaultWriteOptions = - WriteOptions - { writeTransactionSize = 1, - writeOverwriteOptions = OverwriteAllow, - writeAppend = False, - writeUnsafeFFI = False - } + -- | Accumulator types for 'OverwriteDisallow' and 'OverwriteAppend'. Various commonly used + -- accumulators are provided as well. + WriteAccum, + WriteAccumWithOld, + ShowKey, + ShowValue, + writeAccumThrow, + writeAccumThrowAllowSameValue, + writeAccumIgnore, + writeAccumStop, -newtype ExceptionString = ExceptionString String deriving (Show) + -- ** Chunked writing #chunkedwriting# + chunkPairs, + chunkPairsFold, + writeLMDBChunk, -instance Exception ExceptionString + -- * Deleting + deleteLMDB, + clearDatabase, --- | 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.) --- --- 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) wopts = - let txnSize = max 1 (writeTransactionSize wopts) - overwriteOpt = writeOverwriteOptions wopts - flags = - combineOptions $ - [mdb_nooverwrite | overwriteOpt `elem` [OverwriteAllowSame, OverwriteDisallow]] - ++ [mdb_append | writeAppend wopts] - (txn_begin, txn_commit, put_, get) = - if writeUnsafeFFI wopts - then (mdb_txn_begin_unsafe, mdb_txn_commit_unsafe, mdb_put_unsafe_, c_mdb_get_unsafe) - else (mdb_txn_begin, mdb_txn_commit, mdb_put_, c_mdb_get) - 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 + -- * Mode + Mode, + ReadWrite, + ReadOnly, + SubMode, - currChunkSz' <- - liftIO $ - if currChunkSz >= txnSize - then do - let (_, ref) = fromJust mtxn - runIOFinalizer ref - return 0 - else return currChunkSz + -- * Error types + LMDB_Error (..), + MDB_ErrCode (..), - (ptxn, ref) <- - if currChunkSz' == 0 - then liftIO $ - mask_ $ do - ptxn <- txn_begin penv nullPtr 0 - ref <- newIOFinalizer $ txn_commit ptxn - return (ptxn, ref) - else return $ fromJust mtxn + -- * Miscellaneous + ChunkSize (..), + MaybeTxn (..), + EitherTxn (..), + kibibyte, + mebibyte, + gibibyte, + tebibyte, + ) +where - liftIO $ - unsafeUseAsCStringLen k $ \(kp, kl) -> unsafeUseAsCStringLen v $ \(vp, vl) -> - catch - (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 <- 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 - ) +import Streamly.External.LMDB.Internal +import Streamly.External.LMDB.Internal.Foreign diff --git a/src/Streamly/External/LMDB/Internal.hs b/src/Streamly/External/LMDB/Internal.hs --- a/src/Streamly/External/LMDB/Internal.hs +++ b/src/Streamly/External/LMDB/Internal.hs @@ -1,19 +1,1572 @@ -module Streamly.External.LMDB.Internal where - -import Foreign (Ptr) -import Streamly.External.LMDB.Internal.Foreign (MDB_dbi_t, MDB_env) - --- This is in a separate internal module because the tests make use of the Database constructor. - -class Mode a where - isReadOnlyMode :: a -> Bool - -data ReadWrite - -data ReadOnly - -instance Mode ReadWrite where isReadOnlyMode _ = False - -instance Mode ReadOnly where isReadOnlyMode _ = True - -data Database mode = Database (Ptr MDB_env) MDB_dbi_t +{-# LANGUAGE AllowAmbiguousTypes #-} +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE NamedFieldPuns #-} +{-# LANGUAGE NumericUnderscores #-} +{-# LANGUAGE OverloadedRecordDot #-} +{-# LANGUAGE PolyKinds #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE StandaloneKindSignatures #-} +{-# LANGUAGE TupleSections #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TypeFamilies #-} + +module Streamly.External.LMDB.Internal where + +import Control.Concurrent +import qualified Control.Concurrent.Lifted as LI +import Control.Concurrent.STM +import qualified Control.Exception as E +import qualified Control.Exception.Lifted as LI +import Control.Exception.Safe +import Control.Monad +import Control.Monad.IO.Class +import Control.Monad.Trans.Control +import Data.Bifunctor +import Data.ByteString (ByteString) +import qualified Data.ByteString as B +import qualified Data.ByteString.Unsafe as B +import Data.Foldable +import Data.Kind +import Data.Sequence (Seq) +import qualified Data.Sequence as Seq +import Foreign hiding (void) +import Foreign.C +import GHC.TypeLits +import Streamly.Data.Fold (Fold) +import qualified Streamly.Data.Fold as F +import qualified Streamly.Data.Stream.Prelude as S +import Streamly.External.LMDB.Internal.Error +import Streamly.External.LMDB.Internal.Foreign +import qualified Streamly.Internal.Data.Fold as F +import Streamly.Internal.Data.IOFinalizer +import Streamly.Internal.Data.Stream (Stream) +import Streamly.Internal.Data.Unfold (Unfold) +import qualified Streamly.Internal.Data.Unfold as U +import System.Directory +import System.Mem +import Text.Printf + +isReadOnlyEnvironment :: forall emode. (Mode emode) => Bool +isReadOnlyEnvironment = + isReadOnlyMode @emode (error "isReadOnlyEnvironment: unreachable") + +-- | LMDB environments have various limits on the size and number of databases and concurrent +-- readers. +data Limits = Limits + { -- | 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 (see [Databases](#g: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. +-- +-- 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 = mebibyte, + maxDatabases = 0, + maxReaders = 126 + } + +-- | 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; when creating a new +-- environment, one should create an empty file or directory beforehand. If a regular file, an +-- additional file with "-lock" appended to the name is automatically created 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. +-- +-- 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). +-- +-- To satisfy certain low-level LMDB requirements, please do not have opened the same environment +-- (i.e., the same 'FilePath') more than once in the same process at the same time. Furthermore, +-- please use the environment in the process that opened it (not after forking a new process). +openEnvironment :: forall emode. (Mode emode) => FilePath -> Limits -> IO (Environment emode) +openEnvironment path limits = mask_ $ do + -- Low-level requirements: + -- https://github.com/LMDB/lmdb/blob/8d0cbbc936091eb85972501a9b31a8f86d4c51a7/libraries/liblmdb/lmdb.h#L100, + -- https://github.com/LMDB/lmdb/blob/8d0cbbc936091eb85972501a9b31a8f86d4c51a7/libraries/liblmdb/lmdb.h#L102 + + penv <- mdb_env_create + onException + ( do + mdb_env_set_mapsize penv limits.mapSize + let maxDbs = limits.maxDatabases in when (maxDbs /= 0) $ mdb_env_set_maxdbs penv maxDbs + mdb_env_set_maxreaders penv limits.maxReaders + + exists <- doesPathExist path + unless exists $ + throwError + "openEnvironment" + ( "no file or directory found at the specified path; " + ++ "please create an empty file or directory beforehand" + ) + + isDir <- doesDirectoryExist path + + -- Always use MDB_NOTLS, which is crucial for Haskell applications; see + -- https://github.com/LMDB/lmdb/blob/8d0cbbc936091eb85972501a9b31a8f86d4c51a7/libraries/liblmdb/lmdb.h#L615 + let isRo = isReadOnlyEnvironment @emode + flags = mdb_notls : ([mdb_rdonly | isRo] ++ [mdb_nosubdir | not isDir]) + + catchJust + ( \case + LMDB_Error {e_code = Left code} + | Errno (fromIntegral code) == eNOENT && isRo -> Just () + _ -> Nothing + ) + (mdb_env_open penv path (combineOptions flags)) + ( \() -> + -- Provide a friendlier error for a presumably common user mistake. + throwError + "openEnvironment" + ( "mdb_env_open returned 2 (ENOENT); " + ++ "one possibility is that a new empty environment " + ++ "wasn't first opened in ReadWrite mode" + ) + ) + ) + -- In particular if mdb_env_open fails, the environment must be closed; see + -- https://github.com/LMDB/lmdb/blob/8d0cbbc936091eb85972501a9b31a8f86d4c51a7/libraries/liblmdb/lmdb.h#L546 + (c_mdb_env_close penv) + + mvars <- + (,,,) + <$> newTMVarSIO 0 + <*> (WriteLock <$> newEmptyMVar) + <*> (WriteThread <$> newEmptyMVar) + <*> (CloseDbLock <$> newMVar ()) + + return $ Environment penv mvars + +-- | Closes the given environment. +-- +-- If you have merely a few dozen environments at most, there should be no need for this. (It is a +-- common practice with LMDB to create one’s environments once and reuse them for the remainder of +-- the program’s execution.) +-- +-- To satisfy certain low-level LMDB requirements: +-- +-- * Before calling this function, please call 'closeDatabase' on all databases in the environment. +-- * Before calling this function, close all cursors and commit\/abort all transactions on the +-- environment. To make sure this requirement is satisified for read-only transactions, either (a) +-- call 'waitReaders' or (b) pass precreated cursors/transactions to 'readLMDB' and +-- 'unsafeReadLMDB'. +-- * After calling this function, do not use the environment or any related databases, transactions, +-- and cursors. +closeEnvironment :: forall emode. (Mode emode) => Environment emode -> IO () +closeEnvironment (Environment penv _) = do + -- An environment should only be closed once, so the low-level concurrency requirements should be + -- fulfilled: + -- https://github.com/LMDB/lmdb/blob/8d0cbbc936091eb85972501a9b31a8f86d4c51a7/libraries/liblmdb/lmdb.h#L787 + c_mdb_env_close penv + +-- | Gets a database with the given name. +-- +-- 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. +-- +-- /Warning/: When getting a named database for the first time (i.e., creating it), one must do so +-- in the 'ReadWrite' environment mode. (This restriction does not apply for the unnamed database.) +-- In this case, this function spawns a bound thread and creates a temporary read-write transaction +-- under the hood; see [Transactions](#g:transactions). +getDatabase :: + forall emode. + (Mode emode) => + Environment emode -> + Maybe String -> + IO (Database emode) +getDatabase env@(Environment penv mvars) mName = mask_ $ do + -- To satisfy the lower-level concurrency requirements mentioned at + -- https://github.com/LMDB/lmdb/blob/8d0cbbc936091eb85972501a9b31a8f86d4c51a7/libraries/liblmdb/lmdb.h#L1118 + -- we imagine, for simplicity, that everything below is a read-write transaction. Thusly, we also + -- satisfy the MDB_NOTLS read-write transaction serialization requirement (for the case where a + -- read-write transaction actually occur below). This simplification shouldn’t cause any problems, + -- esp. since this function is presumably called relatively rarely in practice. + + let (_, WriteLock lock, _, _) = mvars + putMVar lock () -- Interruptible when waiting for other read-write transactions. + let disclaimWriteOwnership = takeMVar lock + + dbi <- + finally + ( case mName of + Nothing -> do + -- Use a read-only transaction to get the unnamed database. (MDB_CREATE is never needed + -- for the unnamed database.) + ptxn <- mdb_txn_begin penv nullPtr mdb_rdonly + onException + (mdb_dbi_open ptxn Nothing 0 <* mdb_txn_commit ptxn) + (c_mdb_txn_abort ptxn) + Just name -> do + mdbi <- getNamedDb env name + case mdbi of + Just dbi -> + return dbi + Nothing -> + -- The named database was not found. + if isReadOnlyEnvironment @emode + then + throwError + "getDatabase" + ( "please use the ReadWrite environment mode for getting a named database " + ++ "for the first time (i.e., for creating a named database)" + ) + else + -- Use a read-write transaction to create the named database. + -- + -- We run this in a bound thread to make sure the read-write transaction doesn’t + -- cross OS threads. (We do this ourselves instead of putting this burden on the + -- user because this function is presumably called relatively rarely in + -- practice.) + runInBoundThread $ do + ptxn <- mdb_txn_begin penv nullPtr 0 + onException + (mdb_dbi_open ptxn (Just name) mdb_create <* mdb_txn_commit ptxn) + (c_mdb_txn_abort ptxn) + ) + disclaimWriteOwnership + + return $ Database env dbi + +-- | Closes the given database. +-- +-- If you have merely a few dozen databases at most, there should be no need for this. (It is a +-- common practice with LMDB to create one’s databases once and reuse them for the remainder of the +-- program’s execution.) +-- +-- To satisfy certain low-level LMDB requirements: +-- +-- * Before calling this function, please make sure all read-write transactions that have modified +-- the database have already been committed or aborted. +-- * After calling this function, do not use the database or any of its cursors again. To make sure +-- this requirement is satisfied for cursors on read-only transactions, either (a) call +-- 'waitReaders' or (b) pass precreated cursors/transactions to 'readLMDB' and 'unsafeReadLMDB'. +closeDatabase :: forall emode. (Mode emode) => Database emode -> IO () +closeDatabase (Database (Environment penv mvars) dbi) = do + -- We need to serialize the closing; see + -- https://github.com/LMDB/lmdb/blob/8d0cbbc936091eb85972501a9b31a8f86d4c51a7/libraries/liblmdb/lmdb.h#L1200 + let (_, _, _, CloseDbLock lock) = mvars + withMVar lock $ \() -> + c_mdb_dbi_close penv dbi + +-- | A type for an optional thing where we want to fix the transaction mode to @ReadOnly@ in the +-- nothing case. (@Maybe@ isn’t powerful enough for this.) +data MaybeTxn tmode a where + NoTxn :: MaybeTxn ReadOnly a + JustTxn :: a -> MaybeTxn tmode a + +-- | A type for an @Either@-like choice where we want to fix the transaction mode to @ReadOnly@ in +-- the @Left@ case. (@Either@ isn’t powerful enough for this.) +data EitherTxn tmode a b where + LeftTxn :: a -> EitherTxn ReadOnly a b + RightTxn :: b -> EitherTxn tmode a b + +-- | Use @unsafe@ FFI calls under the hood. This can increase iteration speed, but one should +-- bear in mind that @unsafe@ FFI calls, since they block all other threads, can have an adverse +-- impact on the performance of the rest of the program. +-- +-- /Internal/. +newtype UseUnsafeFFI = UseUnsafeFFI Bool deriving (Show) + +-- | Creates an unfold with which we can stream key-value pairs from the given database. +-- +-- If an existing transaction and cursor are not provided, there are two possibilities: (a) If a +-- chunk size is not provided, a read-only transaction and cursor are automatically created for the +-- entire duration of the unfold. (b) Otherwise, new transactions and cursors are automatically +-- created according to the desired chunk size. In this case, each transaction (apart from the first +-- one) starts as expected at the key next to (i.e., the largest\/smallest key less\/greater than) +-- the previously encountered key. +-- +-- If you want to iterate through a large database while avoiding a long-lived transaction (see +-- [Transactions](#g:transactions)), it is your responsibility to either chunk up your usage of +-- 'readLMDB' (with which 'readStart' can help) or specify a chunk size as described above. +-- +-- Runtime consideration: If you call 'readLMDB' very frequently without a precreated transaction +-- and cursor, you might find upon profiling that a significant time is being spent at +-- @mdb_txn_begin@, or find yourself having to increase 'maxReaders' in the environment’s limits +-- because the transactions and cursors are not being garbage collected fast enough. In this case, +-- please consider precreating a transaction and cursor. +-- +-- 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 :: + forall m emode tmode. + (MonadIO m, Mode emode, Mode tmode, SubMode emode tmode) => + Unfold + m + ( ReadOptions, + Database emode, + EitherTxn tmode (Maybe ChunkSize) (Transaction tmode emode, Cursor) + ) + (ByteString, ByteString) +readLMDB = + U.lmap + (\(ropts, db, etxncurs) -> (ropts, UseUnsafeFFI False, db, etxncurs)) + readLMDB' + +-- | Similar to 'readLMDB', except that it has an extra 'UseUnsafeFFI' parameter. +-- +-- /Internal/. +{-# INLINE readLMDB' #-} +readLMDB' :: + forall m emode tmode. + (MonadIO m, Mode emode, Mode tmode, SubMode emode tmode) => + Unfold + m + ( ReadOptions, + UseUnsafeFFI, + Database emode, + EitherTxn tmode (Maybe ChunkSize) (Transaction tmode emode, Cursor) + ) + (ByteString, ByteString) +readLMDB' = + U.lmap + (\(ropts, us, db, etxncurs) -> (ropts, us, db, etxncurs, B.packCStringLen, B.packCStringLen)) + unsafeReadLMDB' + +-- | Similar to 'readLMDB', except that the keys and values are not automatically converted into +-- Haskell @ByteString@s. +-- +-- To ensure safety, please 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 :: + forall m k v emode tmode. + (MonadIO m, Mode emode, Mode tmode, SubMode emode tmode) => + Unfold + m + ( ReadOptions, + Database emode, + EitherTxn tmode (Maybe ChunkSize) (Transaction tmode emode, Cursor), + CStringLen -> IO k, + CStringLen -> IO v + ) + (k, v) +unsafeReadLMDB = + U.lmap + (\(ropts, db, etxncurs, kmap, vmap) -> (ropts, UseUnsafeFFI False, db, etxncurs, kmap, vmap)) + unsafeReadLMDB' + +-- | Similar to 'unsafeReadLMDB', except that it has an extra 'UseUnsafeFFI' parameter. +-- +-- /Internal/. +{-# INLINE unsafeReadLMDB' #-} +unsafeReadLMDB' :: + forall m k v emode tmode. + (MonadIO m, Mode emode, Mode tmode, SubMode emode tmode) => + Unfold + m + ( ReadOptions, + UseUnsafeFFI, + Database emode, + EitherTxn tmode (Maybe ChunkSize) (Transaction tmode emode, Cursor), + CStringLen -> IO k, + CStringLen -> IO v + ) + (k, v) +unsafeReadLMDB' = + -- Performance notes: + -- + Unfortunately, introducing ChunkSize support increased overhead compared to C from around 50 + -- ns/pair to around 90 ns/s (for the safe FFI case). + -- + We mention below what things helped with performance. + -- + We tried various other things (e.g., using [0] or [4] phase control for the inlined + -- subfunctions, wrapping changing state in IORefs, etc.), but to no avail. + -- + For now, we presume that there is simply not much more gain left to achieve, given the extra + -- workload needed for the chunking support. (TODO: We noticed that removing the unsafe FFI + -- support altogether seems to give around 5 ns/pair speedup, but for now we don’t bother.) + let {-# INLINE newTxnCurs #-} + newTxnCurs :: + (Ptr MDB_env -> Ptr MDB_txn -> CUInt -> IO (Ptr MDB_txn)) -> + (Ptr MDB_txn -> IO ()) -> + (Ptr MDB_txn -> MDB_dbi_t -> IO (Ptr MDB_cursor)) -> + (Ptr MDB_cursor -> IO ()) -> + Ptr MDB_env -> + MDB_dbi_t -> + TMVarS NumReaders -> + IO (Ptr MDB_cursor, IOFinalizer) + newTxnCurs + txn_begin + txn_abort + cursor_open + cursor_close + penv + dbi + numReadersT = + mask_ $ do + atomically $ do + n <- takeTMVarS numReadersT + putTMVarS numReadersT $ n + 1 + + let decrReaders = atomically $ do + -- This should, when this is called, never be interruptible because we are, of + -- course, finishing up a reader that is still known to exist. + n <- takeTMVarS numReadersT + putTMVarS numReadersT $ n - 1 + + ptxn <- + onException + (txn_begin penv nullPtr mdb_rdonly) + decrReaders + + pcurs <- + onException + (cursor_open ptxn dbi) + (txn_abort ptxn >> decrReaders) + + txnCursRef <- + onException + ( do + newIOFinalizer . mask_ $ do + -- With LMDB, there is ordinarily no need to commit read-only transactions. + -- (The exception is when we want to make databases that were opened during + -- the transaction available later, but that’s not applicable here.) We can + -- therefore abort ptxn, both for failure (exceptions) and success. + -- + -- Note furthermore that this should be sound in the face of asynchronous + -- exceptions (where this finalizer could get called from a different thread) + -- because LMDB with MDB_NOTLS allows for read-only transactions being used + -- from multiple threads; see + -- https://github.com/LMDB/lmdb/blob/8d0cbbc936091eb85972501a9b31a8f86d4c51a7/libraries/liblmdb/lmdb.h#L984 + cursor_close pcurs >> txn_abort ptxn >> decrReaders + ) + (cursor_close pcurs >> txn_abort ptxn >> decrReaders) + + return (pcurs, txnCursRef) + + {-# INLINE positionCurs #-} + positionCurs :: + (Ptr MDB_cursor -> Ptr MDB_val -> Ptr MDB_val -> MDB_cursor_op_t -> IO CInt) -> + ReadStart -> + Ptr MDB_cursor -> + Ptr MDB_val -> + Ptr MDB_val -> + -- Three possibilities (cursor_get return value): 0 (found), mdb_notfound (not found), other + -- non-zero (error). + IO CInt + positionCurs cursor_get readStart pcurs pk pv = + case readStart of + ReadBeg -> cursor_get pcurs pk pv mdb_first + ReadEnd -> cursor_get pcurs pk pv mdb_last + ReadGE k -> B.unsafeUseAsCStringLen k $ \(kp, kl) -> do + poke pk (MDB_val (fromIntegral kl) kp) + cursor_get pcurs pk pv mdb_set_range + -- For the other cases, LMDB has no built-in operators; so we simulate them ourselves. + ReadGT k -> B.unsafeUseAsCStringLen k $ \(kp, kl) -> do + poke pk (MDB_val (fromIntegral kl) kp) + rc <- cursor_get pcurs pk pv mdb_set_range + if rc == 0 + then do + k' <- peek pk >>= \x -> B.unsafePackCStringLen (x.mv_data, fromIntegral x.mv_size) + if k' == k + then cursor_get pcurs pk pv mdb_next + else return rc + else + -- Error; or not found (if GE is not found, GT doesn’t exist either). + return rc + ReadLE k -> B.unsafeUseAsCStringLen k $ \(kp, kl) -> do + poke pk (MDB_val (fromIntegral kl) kp) + rc <- cursor_get pcurs pk pv mdb_set_range + if rc == 0 + then do + k' <- peek pk >>= \x -> B.unsafePackCStringLen (x.mv_data, fromIntegral x.mv_size) + if k' == k + then return rc + else cursor_get pcurs pk pv mdb_prev + else do + if rc == mdb_notfound + then cursor_get pcurs pk pv mdb_last + else return rc + ReadLT k -> B.unsafeUseAsCStringLen k $ \(kp, kl) -> do + poke pk (MDB_val (fromIntegral kl) kp) + rc <- cursor_get pcurs pk pv mdb_set_range + if rc == 0 + then + -- In both GE and GT cases, find the previous one (to reach LT). + cursor_get pcurs pk pv mdb_prev + else do + if rc == mdb_notfound + then cursor_get pcurs pk pv mdb_last + else return rc + in U.Unfold + ( \( (rc, chunkSz, pcurs, txnCursRef), + rf@ReadLMDBFixed_ {r_db = Database (Environment !penv !mvars) !dbi} + ) -> + liftIO $ do + let (numReadersT, _, _, _) = mvars + + if rc == 0 + then do + -- + pk and pv now contain the data we want to yield; prepare p and v for + -- yielding. + -- + Note: pk will remain important below in the case where the desired maximum + -- ChunkSize is exceeded. + -- + (Avoiding the extra byte size things in the non-ChunkBytes cases seemed to + -- improve performance by around 10 ns/pair.) + -- + (These bang patterns and/or the below bang pattern seemed to improve + -- performance by around 20 ns/pair.) + (!k, !v, !chunkSz') <- + if rf.r_chunkSzInc > 1 + then do + (kSz, k) <- + peek rf.r_pk >>= \x -> + let sz = fromIntegral x.mv_size in (sz,) <$> rf.r_kmap (x.mv_data, sz) + (vSz, v) <- + peek rf.r_pv >>= \x -> + let sz = fromIntegral x.mv_size in (sz,) <$> rf.r_vmap (x.mv_data, sz) + return (k, v, chunkSz + kSz + vSz) + else do + k <- peek rf.r_pk >>= \x -> rf.r_kmap (x.mv_data, fromIntegral x.mv_size) + v <- peek rf.r_pv >>= \x -> rf.r_vmap (x.mv_data, fromIntegral x.mv_size) + return (k, v, chunkSz + rf.r_chunkSzInc) + + -- If the chunk size has exceeded a desired limit, dispose of the existing + -- read-only transaction and cursor and create new ones. + !x <- + if chunkSz' < rf.r_chunkSzMax + then do + -- Staying on the same chunk. + rc' <- rf.r_cursor_get pcurs rf.r_pk rf.r_pv rf.r_nextPrevOp + return ((rc', chunkSz', pcurs, txnCursRef), rf) + else do + -- We make a copy of pk before aborting the current read-only transaction + -- (which makes the data in pk unavailable). + prevk <- + peek rf.r_pk + >>= \x -> B.packCStringLen (x.mv_data, fromIntegral x.mv_size) + runIOFinalizer txnCursRef + (pcurs', txnCursRef') <- rf.r_newtxncurs penv dbi numReadersT + rc' <- + positionCurs + rf.r_cursor_get + (rf.r_nextChunkOp prevk) + pcurs' + rf.r_pk + rf.r_pv + return ((rc', 0, pcurs', txnCursRef'), rf) + + return $ U.Yield (k, v) x + else + if rc == mdb_notfound + then do + runIOFinalizer txnCursRef >> runIOFinalizer rf.r_pref + return U.Stop + else do + runIOFinalizer txnCursRef >> runIOFinalizer rf.r_pref + throwLMDBErrNum "mdb_cursor_get" rc + ) + ( \(ropts, us, db@(Database (Environment penv mvars) dbi), etxncurs, kmap, vmap) -> + liftIO $ do + let useInternalTxnCurs = case etxncurs of + LeftTxn _ -> True + RightTxn _ -> False + + -- Type-level guarantee. + when (useInternalTxnCurs && not (isReadOnlyEnvironment @tmode)) $ error "unreachable" + + case etxncurs of + LeftTxn Nothing -> return () + RightTxn _ -> return () + LeftTxn (Just (ChunkNumPairs maxPairs)) -> + unless (maxPairs > 0) $ + throwError "readLMDB" "please specify positive ChunkNumPairs" + LeftTxn (Just (ChunkBytes maxBytes)) -> + unless (maxBytes > 0) $ + throwError "readLMDB" "please specify positive ChunkBytes" + + (pk, pv, pref) <- mask_ $ do + pk <- malloc + pv <- onException malloc (free pk) + pref <- onException (newIOFinalizer $ free pv >> free pk) (free pv >> free pk) + return (pk, pv, pref) + + let (numReadersT, _, _, _) = mvars + chunkSz :: Int = 0 + + -- + Avoid case lookups in each iteration. + -- + (This seemed to improve performance by over 200 ns/pair.) + (nextPrevOp, nextChunkOp) = case ropts.readDirection of + Forward -> (mdb_next, ReadGT) + Backward -> (mdb_prev, ReadLT) + (chunkSzInc, chunkSzMax) :: (Int, Int) = case etxncurs of + -- chunkSz stays zero. + LeftTxn Nothing -> (0, maxBound) + RightTxn _ -> (0, maxBound) + -- chunkSz increments by 1. + LeftTxn (Just (ChunkNumPairs maxPairs)) -> (1, maxPairs) + -- “chunkSzInc > 1” means we should increment by bytes. (2 is meaningless.) + LeftTxn (Just (ChunkBytes maxBytes)) -> (2, maxBytes) + + (txn_begin, cursor_open, cursor_get, cursor_close, txn_abort) = + case us of + UseUnsafeFFI True -> + ( mdb_txn_begin_unsafe, + mdb_cursor_open_unsafe, + c_mdb_cursor_get_unsafe, + c_mdb_cursor_close_unsafe, + c_mdb_txn_abort_unsafe + ) + UseUnsafeFFI False -> + ( mdb_txn_begin, + mdb_cursor_open, + c_mdb_cursor_get, + c_mdb_cursor_close, + c_mdb_txn_abort + ) + + newtxncurs = newTxnCurs txn_begin txn_abort cursor_open cursor_close + + (pcurs, txnCursRef) <- case etxncurs of + LeftTxn _ -> do + -- Create first transaction and cursor. + newtxncurs penv dbi numReadersT + RightTxn (_, Cursor pcurs) -> do + -- Transaction and cursor are provided by the user. + f <- newIOFinalizer $ return () + return (pcurs, f) + + rc <- positionCurs cursor_get ropts.readStart pcurs pk pv + + return + ( -- State that can change in iterations. + (rc, chunkSz, pcurs, txnCursRef), + ReadLMDBFixed_ + { r_db = db, + r_kmap = kmap, + r_vmap = vmap, + r_newtxncurs = newtxncurs, + r_cursor_get = cursor_get, + r_nextPrevOp = nextPrevOp, + r_nextChunkOp = nextChunkOp, + r_chunkSzInc = chunkSzInc, + r_chunkSzMax = chunkSzMax, + r_pk = pk, + r_pv = pv, + r_pref = pref + } + ) + ) + +-- | State that stays fixed in 'readLMDB' iterations. +-- +-- /Internal/. +data ReadLMDBFixed_ emode k v = ReadLMDBFixed_ + { -- (Keeping the records lazy seemed to improve performance by around 5-10 ns/pair.) + r_db :: Database emode, + r_kmap :: CStringLen -> IO k, + r_vmap :: CStringLen -> IO v, + r_newtxncurs :: + Ptr MDB_env -> + MDB_dbi_t -> + TMVarS NumReaders -> + IO (Ptr MDB_cursor, IOFinalizer), + r_cursor_get :: + Ptr MDB_cursor -> + Ptr MDB_val -> + Ptr MDB_val -> + MDB_cursor_op_t -> + IO CInt, + r_nextPrevOp :: MDB_cursor_op_t, + r_nextChunkOp :: ByteString -> ReadStart, + r_chunkSzInc :: Int, + r_chunkSzMax :: Int, + r_pk :: Ptr MDB_val, + r_pv :: Ptr MDB_val, + r_pref :: IOFinalizer + } + +-- | Looks up the value for the given key in the given database. +-- +-- If an existing transaction is not provided, a read-only transaction is automatically created +-- internally. +-- +-- Runtime consideration: If you call 'getLMDB' very frequently without a precreated transaction, +-- you might find upon profiling that a significant time is being spent at @mdb_txn_begin@, or find +-- yourself having to increase 'maxReaders' in the environment’s limits because the transactions are +-- not being garbage collected fast enough. In this case, please consider precreating a transaction. +{-# INLINE getLMDB #-} +getLMDB :: + forall emode tmode. + (Mode emode, Mode tmode, SubMode emode tmode) => + Database emode -> + MaybeTxn tmode (Transaction tmode emode) -> + ByteString -> + IO (Maybe ByteString) +getLMDB (Database env@(Environment _ _) dbi) mtxn k = + let {-# INLINE brack #-} + brack io = case mtxn of + NoTxn -> withReadOnlyTransaction env $ \(Transaction _ ptxn) -> io ptxn + JustTxn (Transaction _ ptxn) -> io ptxn + in B.unsafeUseAsCStringLen k $ \(kp, kl) -> + with (MDB_val (fromIntegral kl) kp) $ \pk -> alloca $ \pv -> + brack $ \ptxn -> do + rc <- c_mdb_get ptxn dbi pk pv + if rc == 0 + then do + v' <- peek pv + Just <$> B.packCStringLen (v'.mv_data, fromIntegral v'.mv_size) + else + if rc == mdb_notfound + then return Nothing + else throwLMDBErrNum "mdb_get" rc + +-- | A read-only (@tmode@: 'ReadOnly') or read-write (@tmode@: 'ReadWrite') transaction. +-- +-- @emode@: the environment’s mode. Note: 'ReadOnly' environments can only have 'ReadOnly' +-- transactions; we enforce this at the type level. +data Transaction tmode emode = Transaction !(Environment emode) !(Ptr MDB_txn) + +-- | Begins an LMDB read-only transaction on the given environment. +-- +-- For read-only transactions returned from this function, it is your responsibility to (a) make +-- sure the transaction only gets used by a single 'readLMDB', 'unsafeReadLMDB', or 'getLMDB' at the +-- same time, (b) use the transaction only on databases in the environment on which the transaction +-- was begun, (c) make sure that those databases were already obtained before the transaction was +-- begun, (d) dispose of the transaction with 'abortReadOnlyTransaction', and (e) be aware of the +-- caveats regarding long-lived transactions; see [Transactions](#g:transactions). +-- +-- To easily manage a read-only transaction’s lifecycle, we suggest using 'withReadOnlyTransaction'. +{-# INLINE beginReadOnlyTransaction #-} +beginReadOnlyTransaction :: + forall emode. + (Mode emode) => + Environment emode -> + IO (Transaction ReadOnly emode) +beginReadOnlyTransaction env@(Environment penv mvars) = mask_ $ do + -- The non-concurrency requirement: + -- https://github.com/LMDB/lmdb/blob/mdb.master/libraries/liblmdb/lmdb.h#L614 + let (numReadersT, _, _, _) = mvars + + -- Similar comments for NumReaders as in unsafeReadLMDB. + atomically $ do + n <- takeTMVarS numReadersT + putTMVarS numReadersT $ n + 1 + + onException + (Transaction @ReadOnly env <$> mdb_txn_begin penv nullPtr mdb_rdonly) + ( atomically $ do + n <- takeTMVarS numReadersT + putTMVarS numReadersT $ n - 1 + ) + +-- | Disposes of a read-only transaction created with 'beginReadOnlyTransaction'. +-- +-- It is your responsibility to not use the transaction or any of its cursors afterwards. +{-# INLINE abortReadOnlyTransaction #-} +abortReadOnlyTransaction :: forall emode. (Mode emode) => Transaction ReadOnly emode -> IO () +abortReadOnlyTransaction (Transaction (Environment _ mvars) ptxn) = mask_ $ do + let (numReadersT, _, _, _) = mvars + c_mdb_txn_abort ptxn + -- Similar comments for NumReaders as in unsafeReadLMDB. + atomically $ do + n <- takeTMVarS numReadersT + putTMVarS numReadersT $ n - 1 + +-- | Creates a temporary read-only transaction on which the provided action is performed, after +-- which the transaction gets aborted. The transaction also gets aborted upon exceptions. +-- +-- You have the same responsibilities as documented for 'beginReadOnlyTransaction' (apart from the +-- transaction disposal). +{-# INLINE withReadOnlyTransaction #-} +withReadOnlyTransaction :: + forall m a emode. + (Mode emode, MonadBaseControl IO m, MonadIO m) => + Environment emode -> + (Transaction ReadOnly emode -> m a) -> + m a +withReadOnlyTransaction env = + LI.bracket + (liftIO $ beginReadOnlyTransaction env) + -- Aborting a transaction should never fail (as it merely frees a pointer), so any potential + -- issues solved by safe-exceptions (in particular “swallowing asynchronous exceptions via + -- failing cleanup handlers”) shouldn’t apply here. + (liftIO . abortReadOnlyTransaction) + +-- | A cursor. +newtype Cursor = Cursor (Ptr MDB_cursor) + +-- | Opens a cursor for use with 'readLMDB' or 'unsafeReadLMDB'. It is your responsibility to (a) +-- make sure the cursor only gets used by a single 'readLMDB' or 'unsafeReadLMDB' at the same time, +-- (b) make sure the provided database is within the environment on which the provided transaction +-- was begun, and (c) dispose of the cursor with 'closeCursor'. +-- +-- To easily manage a cursor’s lifecycle, we suggest using 'withCursor'. +{-# INLINE openCursor #-} +openCursor :: + forall emode tmode. + (Mode emode, Mode tmode, SubMode emode tmode) => + Transaction tmode emode -> + Database emode -> + IO Cursor +openCursor (Transaction _ ptxn) (Database _ dbi) = + Cursor <$> mdb_cursor_open ptxn dbi + +-- | Disposes of a cursor created with 'openCursor'. +{-# INLINE closeCursor #-} +closeCursor :: Cursor -> IO () +closeCursor (Cursor pcurs) = + -- (Sidenote: Although a cursor will, at least for users who use brackets, usually be called + -- before the transaction gets aborted (read-only/read-write) or committed (read-write), the order + -- doesn’t really matter for read-only transactions.) + c_mdb_cursor_close pcurs + +-- | Creates a temporary cursor on which the provided action is performed, after which the cursor +-- gets closed. The cursor also gets closed upon exceptions. +-- +-- You have the same responsibilities as documented for 'openCursor' (apart from the cursor +-- disposal). +{-# INLINE withCursor #-} +withCursor :: + forall m a emode tmode. + (MonadBaseControl IO m, MonadIO m, Mode emode, Mode tmode, SubMode emode tmode) => + Transaction tmode emode -> + Database emode -> + (Cursor -> m a) -> + m a +withCursor txn db = + LI.bracket + (liftIO $ openCursor txn db) + -- Closing a cursor should never fail (as it merely frees a pointer), so any potential issues + -- solved by safe-exceptions (in particular “swallowing asynchronous exceptions via failing + -- cleanup handlers”) shouldn’t apply here. + (liftIO . closeCursor) + +data ReadOptions = ReadOptions + -- It might seem strange to allow, e.g., ReadBeg and Backward together. However, this simplifies + -- things in the sense that it separates the initial position concept from the iteration + -- (next/prev) concept. + { readStart :: !ReadStart, + readDirection :: !ReadDirection + } + deriving (Show) + +-- | By default, we start reading from the beginning of the database (i.e., from the smallest key) +-- and iterate in forward direction. +defaultReadOptions :: ReadOptions +defaultReadOptions = + ReadOptions + { readStart = ReadBeg, + readDirection = Forward + } + +-- | Direction of key iteration. +data ReadDirection = Forward | Backward deriving (Show) + +-- | The key from which an iteration should start. +data ReadStart + = -- | Start from the smallest key. + ReadBeg + | -- | Start from the largest key. + ReadEnd + | -- | Start from the smallest key that is greater than or equal to the given key. + ReadGE !ByteString + | -- | Start from the smallest key that is greater than the given key. + ReadGT !ByteString + | -- | Start from the largest key that is less than or equal to the given key. + ReadLE !ByteString + | -- | Start from the largest key that is less than the given key. + ReadLT !ByteString + deriving (Show) + +-- | Begins an LMDB read-write transaction on the given environment. +-- +-- Unlike read-only transactions, a given read-write transaction is not allowed to stray from the OS +-- thread on which it was begun, and it is your responsibility to make sure of this. You can achieve +-- this with, e.g., 'Control.Concurrent.runInBoundThread'. +-- +-- Additionally, for read-write transactions returned from this function, it is your responsibility +-- to (a) use the transaction only on databases in the environment on which the transaction was +-- begun, (b) make sure that those databases were already obtained before the transaction was begun, +-- (c) commit\/abort the transaction with 'commitReadWriteTransaction'\/'abortReadWriteTransaction', +-- and (d) be aware of the caveats regarding long-lived transactions; see +-- [Transactions](#g:transactions). +-- +-- To easily manage a read-write transaction’s lifecycle, we suggest using +-- 'withReadWriteTransaction'. +{-# INLINE beginReadWriteTransaction #-} +beginReadWriteTransaction :: Environment ReadWrite -> IO (Transaction ReadWrite ReadWrite) +beginReadWriteTransaction env@(Environment penv mvars) = mask_ $ do + isCurrentThreadBound + >>= flip + unless + (throwError "beginReadWriteTransaction" "please call on a bound thread") + threadId <- myThreadId + + let (_, WriteLock lock, WriteThread writeThread, _) = mvars + putMVar lock () -- Interruptible when waiting for other read-write transactions. + tryPutMVar writeThread threadId >>= flip unless (error "unreachable") + let disclaimWriteOwnership = mask_ $ tryTakeMVar writeThread >> tryTakeMVar lock + + onException + (Transaction @ReadWrite env <$> mdb_txn_begin penv nullPtr 0) + disclaimWriteOwnership + +-- | Aborts a read-write transaction created with 'beginReadWriteTransaction'. +-- +-- It is your responsibility to not use the transaction afterwards. +{-# INLINE abortReadWriteTransaction #-} +abortReadWriteTransaction :: Transaction ReadWrite ReadWrite -> IO () +abortReadWriteTransaction (Transaction (Environment _ mvars) ptxn) = mask_ $ do + let throwErr = throwError "abortReadWriteTransaction" + let (_, WriteLock lock, WriteThread writeThread, _) = mvars + detectUserErrors True writeThread throwErr + c_mdb_txn_abort ptxn + void $ tryTakeMVar lock + +-- | Commits a read-write transaction created with 'beginReadWriteTransaction'. +-- +-- It is your responsibility to not use the transaction afterwards. +{-# INLINE commitReadWriteTransaction #-} +commitReadWriteTransaction :: Transaction ReadWrite ReadWrite -> IO () +commitReadWriteTransaction (Transaction (Environment _ mvars) ptxn) = mask_ $ do + let throwErr = throwError "commitReadWriteTransaction" + let (_, WriteLock lock, WriteThread writeThread, _) = mvars + detectUserErrors True writeThread throwErr + onException + (mdb_txn_commit ptxn) + (c_mdb_txn_abort ptxn >> tryTakeMVar lock) + void $ tryTakeMVar lock + +-- | Spawns a new bound thread and creates a temporary read-write transaction on which the provided +-- action is performed, after which the transaction gets committed. The transaction gets aborted +-- upon exceptions. +-- +-- You have the same responsibilities as documented for 'beginReadWriteTransaction' (apart from +-- running it on a bound thread and committing/aborting it). +{-# INLINE withReadWriteTransaction #-} +withReadWriteTransaction :: + forall m a. + (MonadBaseControl IO m, MonadIO m) => + Environment ReadWrite -> + (Transaction ReadWrite ReadWrite -> m a) -> + m a +withReadWriteTransaction env io = + LI.runInBoundThread $ + -- We need an enhanced bracket. Using the normal bracket and simply committing after running io + -- in the “in-between” computation is incorrect because this entire “in-between” computation + -- falls under “restore” in bracket’s implementation, so an asynchronous exception can cause a + -- commit to be followed by an abort. (Our 'testAsyncExceptionsConcurrent' test exposed this.) + liftedBracket2 + (liftIO $ beginReadWriteTransaction env) + -- + Aborting a transaction should never fail (as it merely frees a pointer), so any potential + -- issues solved by safe-exceptions (in particular “swallowing asynchronous exceptions via + -- failing cleanup handlers”) shouldn’t apply here. + -- + Note: We have convinced ourselves that both the abort and commit are uninterruptible. (In + -- particular, we presume 'myThreadId' (in 'detectUserErrors') is uninterruptible.) + (liftIO . abortReadWriteTransaction) + (liftIO . commitReadWriteTransaction) + io + +-- | +-- * @OverwriteAllow@: When a key reoccurs, overwrite the value. +-- * @OverwriteDisallow@: When a key reoccurs, don’t overwrite and hand the maladaptive key-value +-- pair to the accumulator. +-- * @OverwriteAppend@: Assume the input data is already increasing, which allows the use of +-- @MDB_APPEND@ under the hood and substantially improves write performance. Hand arriving +-- key-value pairs in a maladaptive order to the accumulator. +data OverwriteOptions m a where + OverwriteAllow :: OverwriteOptions m () + OverwriteDisallow :: Either (WriteAccum m a) (WriteAccumWithOld m a) -> OverwriteOptions m a + OverwriteAppend :: WriteAccum m a -> OverwriteOptions m a + +-- | A fold for @(key, new value)@. +type WriteAccum m a = Fold m (ByteString, ByteString) a + +-- | A fold for @(key, new value, old value)@. This has the overhead of getting the old value. +type WriteAccumWithOld m a = Fold m (ByteString, ByteString, ByteString) a + +newtype WriteOptions m a = WriteOptions + { writeOverwriteOptions :: OverwriteOptions m a + } + +-- | A function that shows a database key. +type ShowKey = ByteString -> String + +-- | A function that shows a database value. +type ShowValue = ByteString -> String + +-- | Throws upon the first maladaptive key. If desired, shows the maladaptive key-value pair in the +-- exception. +{-# INLINE writeAccumThrow #-} +writeAccumThrow :: (Monad m) => Maybe (ShowKey, ShowValue) -> WriteAccum m () +writeAccumThrow mshow = + F.foldlM' + ( \() (k, v) -> + throwError "writeLMDB" $ + "Maladaptive key encountered" + ++ maybe + "" + (\(showk, showv) -> printf "; (key,value)=(%s,%s)" (showk k) (showv v)) + mshow + ) + (return ()) + +-- | Throws upon the first maladaptive key where the old value differs from the new value. If +-- desired, shows the maladaptive key-value pair with the old value in the exception. +{-# INLINE writeAccumThrowAllowSameValue #-} +writeAccumThrowAllowSameValue :: (Monad m) => Maybe (ShowKey, ShowValue) -> WriteAccumWithOld m () +writeAccumThrowAllowSameValue mshow = + F.foldlM' + ( \() (k, v, oldv) -> + when (v /= oldv) $ + throwError "writeLMDB" $ + "Maladaptive key encountered" + ++ maybe + "" + ( \(showk, showv) -> + printf "; (key,value,oldValue)=(%s,%s,%s)" (showk k) (showv v) (showv oldv) + ) + mshow + ) + (return ()) + +-- | Silently ignores maladaptive keys. +{-# INLINE writeAccumIgnore #-} +writeAccumIgnore :: (Monad m) => WriteAccum m () +writeAccumIgnore = F.drain + +-- | Gracefully stops upon the first maladaptive key. +{-# INLINE writeAccumStop #-} +writeAccumStop :: (Monad m) => WriteAccum m () +writeAccumStop = void F.one + +-- | By default, we allow overwriting. +defaultWriteOptions :: WriteOptions m () +defaultWriteOptions = + WriteOptions + { writeOverwriteOptions = OverwriteAllow + } + +-- | A chunk size. +data ChunkSize + = -- | Chunk up key-value pairs by number of pairs. The final chunk can have a fewer number of + -- pairs. + ChunkNumPairs !Int + | -- | Chunk up key-value pairs by number of bytes. As soon as the byte count for the keys and + -- values is reached, a new chunk is created (such that each chunk has at least one key-value + -- pair and can end up with more than the desired number of bytes). The final chunk can have + -- less than the desired number of bytes. + ChunkBytes !Int + deriving (Show) + +-- | Chunks up the incoming stream of key-value pairs using the desired chunk size. One can try, +-- e.g., @ChunkBytes mebibyte@ (1 MiB chunks) and benchmark from there. +-- +-- The chunks are processed using the desired fold. +{-# INLINE chunkPairsFold #-} +chunkPairsFold :: + forall m a. + (Monad m) => + ChunkSize -> + Fold m (Seq (ByteString, ByteString)) a -> + Fold m (ByteString, ByteString) a +chunkPairsFold chunkSz (F.Fold astep ainit aextr afinal) = + let {-# INLINE final #-} + final sequ as = + case sequ of + Seq.Empty -> afinal as + _ -> do + astep' <- astep as sequ + case astep' of + F.Done b -> return b + F.Partial as' -> afinal as' + in case chunkSz of + ChunkNumPairs numPairs -> + F.Fold + ( \(!sequ, !as) (k, v) -> + let sequ' = sequ Seq.|> (k, v) + in if Seq.length sequ' == numPairs + then + -- The (user-supplied) astep could already be Done here. + first (Seq.empty,) <$> astep as sequ' + else return . F.Partial $ (sequ', as) + ) + -- The (user-supplied) ainit could already be Done here. + (first (Seq.empty,) <$> ainit) + -- If driven with a scan, the collection fold is assumed to also be compatible with + -- scans and will result in the same output repeatedly for a chunk being built. (This + -- should already be clear to the user.) + (\(_, as) -> aextr as) + -- This is the only direct exit point of this outer fold (since elsewhere it yields a + -- partial). This is therefore the only place where afinal needs to be called. + (uncurry final) + ChunkBytes bytes -> + -- All the comments for the above case hold here too. + F.Fold + ( \(!sequ, !byt, !as) (k, v) -> + let sequ' = sequ Seq.|> (k, v) + byt' = byt + B.length k + B.length v + in if byt' >= bytes + then + first (Seq.empty,0,) <$> astep as sequ' + else + -- For long streams of empty keys and values, sequ' can also get long; but + -- this should be expected behavior (and is an irrelevant edge case for most + -- users anyway). + return . F.Partial $ (sequ', byt', as) + ) + (first (Seq.empty,0,) <$> ainit) + (\(_, _, as) -> aextr as) + (\(sequ, _, as) -> final sequ as) + +-- | Chunks up the incoming stream of key-value pairs using the desired chunk size. One can try, +-- e.g., @ChunkBytes mebibyte@ (1 MiB chunks) and benchmark from there. +{-# INLINE chunkPairs #-} +chunkPairs :: + (Monad m) => + ChunkSize -> + Stream m (ByteString, ByteString) -> + Stream m (Seq (ByteString, ByteString)) +chunkPairs chunkSz = + S.foldMany $ + case chunkSz of + ChunkNumPairs numPairs -> + F.Fold + ( \(!sequ) (k, v) -> + let sequ' = sequ Seq.|> (k, v) + in return $ + if Seq.length sequ' == numPairs + then F.Done sequ' + else F.Partial sequ' + ) + (return $ F.Partial Seq.empty) + (error "unreachable") + return + ChunkBytes bytes -> + F.Fold + ( \(!sequ, !byt) (k, v) -> + let sequ' = sequ Seq.|> (k, v) + byt' = byt + B.length k + B.length v + in return $ + if byt' >= bytes + then F.Done sequ' + else F.Partial (sequ', byt') + ) + (return $ F.Partial (Seq.empty, 0)) + (error "unreachable") + (\(sequ, _) -> return sequ) + +-- | Writes a chunk of key-value pairs to the given database. Under the hood, it uses 'writeLMDB' +-- surrounded with a 'withReadWriteTransaction'. +{-# INLINE writeLMDBChunk #-} +writeLMDBChunk :: + forall m a. + (MonadBaseControl IO m, MonadIO m, MonadCatch m) => + WriteOptions m a -> + Database ReadWrite -> + Seq (ByteString, ByteString) -> + m a +writeLMDBChunk = + writeLMDBChunk' (UseUnsafeFFI False) + +-- | Similar to 'writeLMDBChunk', except that it has an extra 'UseUnsafeFFI' parameter. +-- +-- /Internal/. +{-# INLINE writeLMDBChunk' #-} +writeLMDBChunk' :: + forall m a. + (MonadBaseControl IO m, MonadIO m, MonadCatch m) => + UseUnsafeFFI -> + WriteOptions m a -> + Database ReadWrite -> + Seq (ByteString, ByteString) -> + m a +writeLMDBChunk' useUnsafeFFI wopts db@(Database env _) sequ = + withReadWriteTransaction env $ \txn -> + S.fold (writeLMDB' useUnsafeFFI wopts db txn) . S.fromList . toList $ sequ + +-- | Creates a fold that writes a stream of key-value pairs to the provided database using the +-- provided transaction. +-- +-- If you have a long stream of key-value pairs that you want to write to an LMDB database while +-- avoiding a long-lived transaction (see [Transactions](#g:transactions)), you can use the +-- functions for [chunked writing](#g:chunkedwriting). +{-# INLINE writeLMDB #-} +writeLMDB :: + forall m a. + (MonadIO m, MonadCatch m, MonadThrow m) => + WriteOptions m a -> + Database ReadWrite -> + Transaction ReadWrite ReadWrite -> + Fold m (ByteString, ByteString) a +writeLMDB = + writeLMDB' (UseUnsafeFFI False) + +-- | Similar to 'writeLMDB', except that it has an extra 'UseUnsafeFFI' parameter. +-- +-- /Internal/. +{-# INLINE writeLMDB' #-} +writeLMDB' :: + forall m a. + (MonadIO m, MonadCatch m, MonadThrow m) => + UseUnsafeFFI -> + WriteOptions m a -> + Database ReadWrite -> + Transaction ReadWrite ReadWrite -> + Fold m (ByteString, ByteString) a +writeLMDB' + (UseUnsafeFFI us) + wopts + (Database env@(Environment _ mvars) dbi) + txn@(Transaction _ ptxn) = + -- Notes on why writeLMDB relies on the user creating read-write transactions up front, as + -- opposed to writeLMDB itself maintaining the write transactions internally (as was the case in + -- versions <=0.7.0): + -- * The old way was not safe because LMDB read-write transactions are (unless MDB_NOLOCK is + -- used) not allowed to cross OS threads; but upon asynchronous exceptions, the read-write + -- transaction aborting would happen upon garbage collection (GC), which can occur on a + -- different OS thread (even if the user ran the original writeLMDB on a bound thread). + -- * We see no way around this but to wrap every read-write transaction in a bona fide bracket + -- managed by the user (not a streamly-type bracket, which, again, relies on GC). + -- * Two things we investigated: (a) Channels allow us to pass all writing to a specific OS + -- thread, but doing this one-by-one for every mdb_put is way too slow; for channels to + -- become performant, they need chunking. (b) We can use MDB_NOLOCK to avoid the + -- same-OS-thread requirement, but this means other processes can no longer safely interact + -- with the LMDB environment. + -- * Two benefits of the new way: (a) A stream can be demuxed into writeLMDB folds on the same + -- environment. (b) The writeLMDB fold works with scans. + let put_ = + if us + then mdb_put_unsafe_ + else mdb_put_ + + throwErr = throwError "writeLMDB" + + {-# INLINE validate #-} + validate = do + let (_, _, WriteThread writeThread, _) = mvars + liftIO $ detectUserErrors False writeThread throwErr + + {-# INLINE putCatchKeyExists #-} + putCatchKeyExists :: + ByteString -> + ByteString -> + s -> -- State of the Accum fold (for the failures). + CUInt -> + (() -> m (F.Step s d)) -> + m (F.Step s d) + putCatchKeyExists k v s op = + catchJust + ( \case + LMDB_Error {e_code = Right MDB_KEYEXIST} -> Just () + _ -> Nothing + ) + ( do + liftIO $ + B.unsafeUseAsCStringLen k $ \(kp, kl) -> + B.unsafeUseAsCStringLen v $ \(vp, vl) -> + put_ ptxn dbi kp (fromIntegral kl) vp (fromIntegral vl) op + return $ F.Partial s + ) + + {-# INLINE commonFold #-} + commonFold (F.Fold fstep finit fextr ffinal) op = + F.Fold @m + (\s (k, v) -> putCatchKeyExists k v s op (\() -> fstep s (k, v))) + (validate >> finit) + fextr + ffinal + in case writeOverwriteOptions wopts of + OverwriteAllow -> + F.foldlM' @m @() + ( \() (k, v) -> liftIO $ + B.unsafeUseAsCStringLen k $ \(kp, kl) -> B.unsafeUseAsCStringLen v $ \(vp, vl) -> + put_ ptxn dbi kp (fromIntegral kl) vp (fromIntegral vl) 0 + ) + validate + OverwriteDisallow (Left f) -> + commonFold f mdb_nooverwrite + OverwriteDisallow (Right (F.Fold fstep finit fextr ffinal)) -> + F.Fold @m + ( \s (k, v) -> + putCatchKeyExists k v s mdb_nooverwrite $ \_ -> do + mVold <- liftIO $ getLMDB (Database env dbi) (JustTxn txn) k + vold <- case mVold of + Nothing -> throwErr "getLMDB; old value not found; this should never happen" + Just vold -> return vold + fstep s (k, v, vold) + ) + (validate >> finit) + fextr + ffinal + OverwriteAppend f -> + commonFold f mdb_append + +-- | Waits for active read-only transactions on the given environment to finish. Note: This triggers +-- garbage collection. +waitReaders :: (Mode emode) => Environment emode -> IO () +waitReaders (Environment _ mvars) = do + let (numReadersT, _, _, _) = mvars + performGC -- Complete active readers as soon as possible. + numReaders <- + atomically $ do + numReaders <- takeTMVarS numReadersT + check $ numReaders <= 0 -- Sanity check: use <=0 to catch unexpected negative readers. + return numReaders + when (numReaders /= 0) $ throwError "waitReaders" "zero numReaders expected" + +-- | Clears, i.e., removes all key-value pairs from, the given database. +-- +-- /Warning/: Under the hood, this function spawns a bound thread and creates a potentially +-- long-lived read-write transaction; see [Transactions](#g:transactions). +clearDatabase :: Database ReadWrite -> IO () +clearDatabase (Database (Environment penv mvars) dbi) = mask $ \restore -> do + let (_, WriteLock lock, _, _) = mvars + putMVar lock () -- Interruptible when waiting for other read-write transactions. + let disclaimWriteOwnership = takeMVar lock + + finally + ( runInBoundThread $ do + ptxn <- mdb_txn_begin penv nullPtr 0 + onException + -- Unmask a potentially long-running operation. + ( restore $ do + mdb_clear ptxn dbi + mdb_txn_commit ptxn + ) + (c_mdb_txn_abort ptxn) -- TODO: Could abort be long-running? + ) + disclaimWriteOwnership + +-- | Deletes the given key from the given database using the given transaction. +{-# INLINE deleteLMDB #-} +deleteLMDB :: + DeleteOptions -> + Database emode -> + Transaction ReadWrite emode -> + ByteString -> + IO () +deleteLMDB dopts (Database (Environment _ _) dbi) (Transaction _ ptxn) k = + B.unsafeUseAsCStringLen k $ \(kp, kl) -> + with (MDB_val (fromIntegral kl) kp) $ \pk -> + c_mdb_del ptxn dbi pk nullPtr >>= \rc -> + when (rc /= 0) $ + unless (rc == mdb_notfound && not dopts.deleteAssumeExists) $ + throwLMDBErrNum "mdb_del" rc + +newtype DeleteOptions = DeleteOptions + { -- | Assume that the key being deleted already exists in the database and throw if it doesn’t. + deleteAssumeExists :: Bool + } + deriving (Show) + +-- | By default, we do /not/ assume the key being deleted already exists in the database. +defaultDeleteOptions :: DeleteOptions +defaultDeleteOptions = + DeleteOptions {deleteAssumeExists = False} + +-- | A convenience constant for obtaining 1 KiB. +kibibyte :: (Num a) => a +kibibyte = 1_024 + +-- | A convenience constant for obtaining 1 MiB. +mebibyte :: (Num a) => a +mebibyte = 1_024 * 1_024 + +-- | A convenience constant for obtaining 1 GiB. +gibibyte :: (Num a) => a +gibibyte = 1_024 * 1_024 * 1_024 + +-- | A convenience constant for obtaining 1 TiB. +tebibyte :: (Num a) => a +tebibyte = 1_024 * 1_024 * 1_024 * 1_024 + +-- | A type class for 'ReadOnly' and 'ReadWrite' environments and transactions. +class Mode a where + isReadOnlyMode :: a -> Bool + +data ReadWrite + +data ReadOnly + +instance Mode ReadWrite where isReadOnlyMode _ = False + +instance Mode ReadOnly where isReadOnlyMode _ = True + +-- | Enforces at the type level that @ReadWrite@ environments support both @ReadWrite@ and +-- @ReadOnly@ transactions, but @ReadOnly@ environments support only @ReadOnly@ transactions. +type SubMode :: k -> k -> Constraint +type family SubMode emode tmode where + SubMode ReadWrite _ = () + SubMode ReadOnly ReadOnly = () + SubMode ReadOnly ReadWrite = + TypeError ('Text "ReadOnly environments only support ReadOnly transactions") + +data Environment emode + = Environment + !(Ptr MDB_env) + !(TMVarS NumReaders, WriteLock, WriteThread, CloseDbLock) + +newtype TMVarS a = TMVarS (TMVar a) + +{-# INLINE newTMVarSIO #-} +newTMVarSIO :: a -> IO (TMVarS a) +newTMVarSIO a = + TMVarS <$> newTMVarIO a + +{-# INLINE takeTMVarS #-} +takeTMVarS :: TMVarS a -> STM a +takeTMVarS (TMVarS tmVar) = + takeTMVar tmVar + +-- Same as putTMVar except it makes sure the value is evaluated to WHNF. (For now we only use this +-- to prevent NumReaders thunks, for which WHNF is enough.) +{-# INLINE putTMVarS #-} +putTMVarS :: TMVarS a -> a -> STM () +putTMVarS (TMVarS tmVar) a = + putTMVar tmVar $! a + +-- The number of current readers. This needs to be kept track of due to MDB_NOLOCK; see comments in +-- writeLMDB. +newtype NumReaders = NumReaders Int deriving (Eq, Num, Ord) + +-- An increasing counter for various write-related functions using the same environment. +newtype WriteCounter = WriterCounter Int deriving (Bounded, Eq, Num, Ord) + +-- | Keeps track of the 'ThreadId' of the current read-write transaction. +newtype WriteThread = WriteThread (MVar ThreadId) + +-- For read-write transaction serialization. +newtype WriteLock = WriteLock (MVar ()) + +-- For closeDatabase serialization. +newtype CloseDbLock = CloseDbLock (MVar ()) + +data Database emode = Database !(Environment emode) !MDB_dbi_t + +-- | Utility function for getting a named database with a read-only transaction, returning 'Nothing' +-- if it was not found. +-- +-- /Internal/. +getNamedDb :: + forall emode. + (Mode emode) => + Environment emode -> + String -> + IO (Maybe MDB_dbi_t) +getNamedDb (Environment penv _) name = mask_ $ do + -- Use a read-only transaction to try to get the named database. + ptxn <- mdb_txn_begin penv nullPtr mdb_rdonly + onException + ( catchJust + ( \case + -- Assumption: mdb_txn_commit never returns MDB_NOTFOUND. + LMDB_Error {e_code} | e_code == Right MDB_NOTFOUND -> Just () + _ -> Nothing + ) + (Just <$> mdb_dbi_open ptxn (Just name) 0 <* mdb_txn_commit ptxn) + ( \() -> do + -- The named database was not found. + c_mdb_txn_abort ptxn + return Nothing + ) + ) + (c_mdb_txn_abort ptxn) + +-- | A utility function for detecting a few user errors. +-- +-- /Internal/. +detectUserErrors :: Bool -> MVar ThreadId -> (String -> IO ()) -> IO () +detectUserErrors shouldTake writeThread throwErr = do + let info = "LMDB transactions might now be in a mangled state" + inappr ctx = printf "inappropriately called (%s); %s" ctx info + unexpThread = "called on unexpected thread; " ++ info + caseShouldTake = "before aborting/committing read-write transaction" + caseNotShouldTake = "before starting writeLMDB" + threadId <- myThreadId + if shouldTake + then + -- Before aborting/committing read-write transactions. + tryTakeMVar writeThread >>= \case + Nothing -> throwErr $ inappr caseShouldTake + Just tid + | tid /= threadId -> throwErr unexpThread + | otherwise -> return () + else do + -- Before starting a writeLMDB. + isEmptyMVar writeThread >>= flip when (throwErr $ inappr caseNotShouldTake) + void $ withMVarMasked writeThread $ \tid -> + if tid /= threadId + then throwErr unexpThread + else void $ return tid + +-- | +-- @liftedBracket2 acquire failure success thing@ +-- +-- Same as @Control.Exception.Lifted.bracket@ (from @lifted-base@) except it distinguishes between +-- failure and success. +-- +-- Notes: +-- +-- * When @acquire@, @success@, or @failure@ throw exceptions, any monadic side effects in @m@ will +-- be discarded. +-- * When @thing@ throws an exception, any monadic side effects in @m@ produced by @thing@ will be +-- discarded, but the side effects of @acquire@ and (non-excepting) @failure@ will be retained. +-- * When (following a @thing@ success) @success@ throws an exception, any monadic side effects in +-- @m@ produced by @success@ will be discarded, but the side effects of @acquire@, @thing@, and +-- (non-excepting) @failure@ will be retained. +-- +-- /Internal/. +{-# INLINE liftedBracket2 #-} +liftedBracket2 :: + (MonadBaseControl IO m) => + m a -> + (a -> m b) -> + (a -> m b) -> + (a -> m c) -> + m c +liftedBracket2 acquire failure success thing = control $ \runInIO -> + bracket2 + (runInIO acquire) + (\st -> runInIO $ restoreM st >>= failure) + (\st -> runInIO $ restoreM st >>= success) + (\st -> runInIO $ restoreM st >>= thing) + +-- | Same as @Control.Exception.bracket@ except it distinguishes between failure and success. (If +-- the success action throws an exception, the failure action gets called.) +-- +-- /Internal/. +{-# INLINE bracket2 #-} +bracket2 :: + IO a -> + (a -> IO b) -> + (a -> IO b) -> + (a -> IO c) -> + IO c +bracket2 acquire failure success thing = mask $ \restore -> do + a <- acquire + r <- restore (thing a) `E.onException` failure a + _ <- success a `E.onException` failure a + return r diff --git a/src/Streamly/External/LMDB/Internal/Error.hs b/src/Streamly/External/LMDB/Internal/Error.hs new file mode 100644 --- /dev/null +++ b/src/Streamly/External/LMDB/Internal/Error.hs @@ -0,0 +1,14 @@ +module Streamly.External.LMDB.Internal.Error where + +import Control.Exception +import Text.Printf + +data Error = Error !String !String + +instance Show Error where + show (Error ctx msg) = printf "streamly-lmdb; %s; %s" ctx msg + +instance Exception Error + +throwError :: String -> String -> m a +throwError ctx msg = throw $ Error ctx msg diff --git a/src/Streamly/External/LMDB/Internal/Foreign.hsc b/src/Streamly/External/LMDB/Internal/Foreign.hsc --- a/src/Streamly/External/LMDB/Internal/Foreign.hsc +++ b/src/Streamly/External/LMDB/Internal/Foreign.hsc @@ -4,12 +4,11 @@ #include -import Control.Exception (Exception, throwIO) -import Control.Monad (when) -import Foreign ((.|.), Ptr, Storable (alignment, peek, peekByteOff, - poke, pokeByteOff, sizeOf), Word16, Word32, alloca, nullPtr) -import Foreign.C.String (CString, peekCString, withCString) -import Foreign.C.Types (CChar, CInt (CInt), CSize (CSize), CUInt (CUInt)) +import Control.Exception +import Control.Monad +import Foreign +import Foreign.C.String +import Foreign.C.Types import qualified Data.List as L @@ -121,6 +120,9 @@ foreign import ccall safe "lmdb.h mdb_drop" c_mdb_drop :: Ptr MDB_txn -> MDB_dbi_t -> CInt -> IO CInt +foreign import ccall safe "lmdb.h mdb_del" + c_mdb_del :: Ptr MDB_txn -> MDB_dbi_t -> Ptr MDB_val -> Ptr MDB_val -> IO CInt + data LMDB_Error = LMDB_Error { e_context :: String , e_description :: String @@ -200,6 +202,9 @@ mdb_notls :: CUInt mdb_notls = #const MDB_NOTLS +mdb_nolock :: CUInt +mdb_nolock = #const MDB_NOLOCK + mdb_nosubdir :: CUInt mdb_nosubdir = #const MDB_NOSUBDIR @@ -266,33 +271,31 @@ alloca $ \pptxn -> c_mdb_txn_begin_unsafe penv parent flags pptxn >>= \rc -> if rc /= 0 then throwLMDBErrNum "mdb_txn_begin" rc else peek pptxn --- If the commit fails, aborts the transaction. mdb_txn_commit :: Ptr MDB_txn -> IO () mdb_txn_commit ptxn = c_mdb_txn_commit ptxn >>= \rc -> - when (rc /= 0) $ c_mdb_txn_abort ptxn >> throwLMDBErrNum "mdb_txn_commit" rc + when (rc /= 0) $ throwLMDBErrNum "mdb_txn_commit" rc --- If the commit fails, aborts the transaction. mdb_txn_commit_unsafe :: Ptr MDB_txn -> IO () mdb_txn_commit_unsafe ptxn = c_mdb_txn_commit_unsafe ptxn >>= \rc -> - when (rc /= 0) $ c_mdb_txn_abort_unsafe ptxn >> throwLMDBErrNum "mdb_txn_commit" rc + when (rc /= 0) $ throwLMDBErrNum "mdb_txn_commit" rc mdb_cursor_open :: Ptr MDB_txn -> MDB_dbi_t -> IO (Ptr MDB_cursor) mdb_cursor_open ptxn dbi = alloca $ \ppcurs -> c_mdb_cursor_open ptxn dbi ppcurs >>= \rc -> - if rc /= 0 then c_mdb_txn_abort ptxn >> throwLMDBErrNum "mdb_cursor_open" rc else peek ppcurs + if rc /= 0 then throwLMDBErrNum "mdb_cursor_open" rc else peek ppcurs mdb_cursor_open_unsafe :: Ptr MDB_txn -> MDB_dbi_t -> IO (Ptr MDB_cursor) mdb_cursor_open_unsafe ptxn dbi = alloca $ \ppcurs -> c_mdb_cursor_open_unsafe ptxn dbi ppcurs >>= \rc -> - if rc /= 0 then c_mdb_txn_abort ptxn >> throwLMDBErrNum "mdb_cursor_open" rc else peek ppcurs + if rc /= 0 then throwLMDBErrNum "mdb_cursor_open" rc else peek ppcurs mdb_dbi_open :: Ptr MDB_txn -> Maybe String -> CUInt -> IO MDB_dbi_t mdb_dbi_open ptxn name flags = do withCStringMaybe name $ \cname -> alloca $ \pdbi -> c_mdb_dbi_open ptxn cname flags pdbi >>= \rc -> - if rc /= 0 then c_mdb_txn_abort ptxn >> throwLMDBErrNum "mdb_dbi_open" rc else peek pdbi + if rc /= 0 then throwLMDBErrNum "mdb_dbi_open" rc else peek pdbi {-# INLINE mdb_put #-} mdb_put :: Ptr MDB_txn -> MDB_dbi_t -> Ptr MDB_val -> Ptr MDB_val -> CUInt -> IO () diff --git a/streamly-lmdb.cabal b/streamly-lmdb.cabal --- a/streamly-lmdb.cabal +++ b/streamly-lmdb.cabal @@ -1,6 +1,6 @@ cabal-version: 3.0 name: streamly-lmdb -version: 0.7.0 +version: 0.8.0 synopsis: Stream data to or from LMDB databases using the streamly library. description: Please see the README on GitHub at category: Database, Streaming, Streamly @@ -8,7 +8,7 @@ bug-reports: https://github.com/shlok/streamly-lmdb/issues author: Shlok Datye maintainer: sd-haskell@quant.is -copyright: 2023 Shlok Datye +copyright: 2024 Shlok Datye license: BSD-3-Clause license-file: LICENSE build-type: Simple @@ -24,6 +24,7 @@ exposed-modules: Streamly.External.LMDB Streamly.External.LMDB.Internal + Streamly.External.LMDB.Internal.Error Streamly.External.LMDB.Internal.Foreign other-modules: Paths_streamly_lmdb @@ -41,17 +42,23 @@ extra-libraries: lmdb build-depends: - async >=2.2.2 && <2.3 - , base >=4.7 && <5 + base >=4.7 && <5 , bytestring >=0.10.10.0 && <0.12 - , streamly ==0.9.* - , streamly-core ==0.1.* + , containers >=0.6.5.1 && <0.7 + , directory >=1.3.6.0 && <1.4 + , lifted-base >=0.2.3.12 && <0.3 + , monad-control >=1.0.3.1 && <1.1 + , safe-exceptions >=0.1.7.3 && <0.2 + , stm >=2.5.0.2 && <2.6 + , streamly >=0.10.0 && <0.11 + , streamly-core >=0.2.0 && <0.3 default-language: Haskell2010 test-suite streamly-lmdb-test type: exitcode-stdio-1.0 main-is: TestSuite.hs other-modules: + ReadmeMain Streamly.External.LMDB.Tests Paths_streamly_lmdb autogen-modules: @@ -62,15 +69,22 @@ extra-libraries: lmdb build-depends: - QuickCheck >=2.13.2 && <2.15 - , async >=2.2.2 && <2.3 + async >=2.2.2 && <2.3 , base >=4.7 && <5 , bytestring >=0.10.10.0 && <0.12 + , cereal >=0.5.8.3 && <0.6 + , containers >=0.6.5.1 && <0.7 , directory >=1.3.6.0 && <1.4 - , streamly ==0.9.* - , streamly-core ==0.1.* + , mtl >=2.3.1 && <2.4 + , QuickCheck >=2.13.2 && <2.15 + , random >=1.2.1.2 && <1.3 + , streamly >=0.10.0 && <0.11 + , streamly-core >=0.2.0 && <0.3 , streamly-lmdb , tasty >=1.2.3 && <1.5 + , tasty-hunit >=0.10.0.3 && <0.11 , tasty-quickcheck >=0.10.1.1 && <0.11 - , temporary ==1.3.* + , temporary >=1.3 && <1.4 + , transformers >=0.6.0.2 && <0.7 + , vector >=0.12.3.1 && <0.14 default-language: Haskell2010 diff --git a/test/ReadmeMain.hs b/test/ReadmeMain.hs new file mode 100644 --- /dev/null +++ b/test/ReadmeMain.hs @@ -0,0 +1,38 @@ +{-# LANGUAGE OverloadedStrings #-} + +module ReadmeMain where + +import Data.Function +import qualified Streamly.Data.Fold as F +import qualified Streamly.Data.Stream.Prelude as S +import Streamly.External.LMDB + +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} + + -- 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. + withReadWriteTransaction env $ \txn -> + [("baz", "a"), ("foo", "b"), ("bar", "c")] + & S.fromList + & S.fold (writeLMDB defaultWriteOptions db txn) + + -- Stream key-value pairs out of the + -- database, printing them along the way. + -- Output: + -- ("bar","c") + -- ("baz","a") + -- ("foo","b") + S.unfold readLMDB (defaultReadOptions, db, LeftTxn Nothing) + & S.mapM print + & S.fold F.drain diff --git a/test/Streamly/External/LMDB/Tests.hs b/test/Streamly/External/LMDB/Tests.hs --- a/test/Streamly/External/LMDB/Tests.hs +++ b/test/Streamly/External/LMDB/Tests.hs @@ -1,362 +1,1031 @@ -{-# LANGUAGE TypeApplications #-} - -module Streamly.External.LMDB.Tests (tests) where - -import Control.Concurrent.Async (asyncBound, wait) -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.Data.Stream.Prelude (fromList, toList, unfold) -import qualified Streamly.Data.Stream.Prelude as S -import Streamly.External.LMDB -import Streamly.External.LMDB.Internal (Database (..)) -import Streamly.External.LMDB.Internal.Foreign -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) - -tests :: IO (Database ReadWrite, Environment ReadWrite) -> [TestTree] -tests dbenv = - [ testReadLMDB dbenv, - testUnsafeReadLMDB dbenv, - testWriteLMDB dbenv, - testWriteLMDB_2 dbenv, - testWriteLMDB_3 dbenv, - testBetween - ] - -withReadOnlyTxnAndCurs :: - (Mode mode) => - Environment mode -> - Database mode -> - ((ReadOnlyTxn, Cursor) -> IO r) -> - IO r -withReadOnlyTxnAndCurs env db = - bracket - (beginReadOnlyTxn env >>= \txn -> openCursor txn db >>= \curs -> return (txn, curs)) - (\(txn, curs) -> closeCursor curs >> abortReadOnlyTxn txn) - --- | 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, Environment mode) -> TestTree -testReadLMDB res = testProperty "readLMDB" . monadicIO $ do - (db, env) <- run res - keyValuePairs <- arbitraryKeyValuePairs'' - run $ clearDatabase db - - run $ writeChunk db False keyValuePairs - let keyValuePairsInDb = sort . removeDuplicateKeys $ keyValuePairs - - (readOpts, expectedResults) <- pick $ readOptionsAndResults keyValuePairsInDb - let unf txn = toList $ unfold (readLMDB db txn readOpts) undefined - results <- run $ unf Nothing - resultsTxn <- run $ withReadOnlyTxnAndCurs env db (unf . Just) - - 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, Environment mode) -> TestTree -testUnsafeReadLMDB res = testProperty "unsafeReadLMDB" . monadicIO $ do - (db, env) <- run res - keyValuePairs <- arbitraryKeyValuePairs'' - run $ clearDatabase db - - 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 - let unf txn = - toList $ - unfold (unsafeReadLMDB db txn readOpts (return . snd) (return . snd)) undefined - lengths <- run $ unf Nothing - lengthsTxn <- run $ withReadOnlyTxnAndCurs env db (unf . Just) - - 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, Environment ReadWrite) -> TestTree -testWriteLMDB res = testProperty "writeLMDB" . monadicIO $ do - (db, _) <- run res - keyValuePairs <- arbitraryKeyValuePairs - run $ clearDatabase db - - chunkSz <- pick arbitrary - unsafeFFI <- pick arbitrary - - let fol' = - writeLMDB db $ - defaultWriteOptions - { writeTransactionSize = chunkSz, - writeOverwriteOptions = OverwriteAllow, - writeUnsafeFFI = unsafeFFI - } - - -- 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 Nothing defaultReadOptions) undefined - - 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, Environment ReadWrite) -> TestTree -testWriteLMDB_2 res = testProperty "writeLMDB_2" . monadicIO $ do - (db, _) <- run res - keyValuePairs <- arbitraryKeyValuePairs' - run $ clearDatabase db - - chunkSz <- pick arbitrary - unsafeFFI <- pick arbitrary - - -- TODO: Run with new "bound" functionality in streamly. - let fol' = - writeLMDB db $ - defaultWriteOptions - { writeTransactionSize = chunkSz, - writeOverwriteOptions = OverwriteDisallow, - writeUnsafeFFI = unsafeFFI - } - 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 Nothing defaultReadOptions) undefined - let pairsAsExpected = keyValuePairsInDb == readPairsAll - - 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, Environment ReadWrite) -> TestTree -testWriteLMDB_3 res = testProperty "writeLMDB_3" . monadicIO $ do - (db, _) <- run res - keyValuePairs <- arbitraryKeyValuePairs' - run $ clearDatabase db - - chunkSz <- pick arbitrary - unsafeFFI <- pick arbitrary - - -- TODO: Run with new "bound" functionality in streamly. - let fol' = - writeLMDB db $ - defaultWriteOptions - { writeTransactionSize = chunkSz, - writeOverwriteOptions = OverwriteAllowSame, - writeUnsafeFFI = unsafeFFI - } - 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 Nothing defaultReadOptions) undefined - let pairsAsExpected = keyValuePairsInDb == readPairsAll - - return $ exceptionAsExpected && pairsAsExpected - -arbitraryKeyValuePairs :: PropertyM IO [(ByteString, ByteString)] -arbitraryKeyValuePairs = - map (\(ws1, ws2) -> (pack ws1, pack ws2)) - . filter (\(ws1, _) -> not (null ws1)) -- LMDB does not allow empty keys. - <$> 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 - --- 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' - ) - ] - --- | 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 - -hasDuplicateKeys :: (Eq a) => [(a, b)] -> Bool -hasDuplicateKeys l = - 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 - -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 - 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 - Nothing -> xs - Just i -> take i xs - --- Assumes first < second. -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] - -testBetween :: TestTree -testBetween = testProperty "testBetween" $ \ws1 ws2 -> - (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 - -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 - unsafeFFI <- arbitrary - let len = length pairsInDb - readAll <- frequency [(1, return True), (3, return False)] - let ropts = defaultReadOptions {readDirection = dir, readUnsafeFFI = unsafeFFI} - 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 - -- 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 - (EQ, Forward) -> forwEq - (EQ, Backward) -> backwEq - (GT, Forward) -> case nextKey of - 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) - (LT, Forward) -> case prevKey of - 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) - --- Writes the given key-value pairs to the given database. -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 - -{-# INLINE marshalOut #-} -marshalOut :: ByteString -> (MDB_val -> IO ()) -> IO () -marshalOut bs f = - unsafeUseAsCStringLen bs $ \(ptr, len) -> f $ MDB_val (fromIntegral len) (castPtr ptr) +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE InstanceSigs #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TupleSections #-} +{-# LANGUAGE TypeApplications #-} +{-# OPTIONS_GHC -Wno-unrecognised-pragmas #-} + +{-# HLINT ignore "Use <$>" #-} + +module Streamly.External.LMDB.Tests where + +import Control.Concurrent +import Control.Concurrent.Async +import Control.Exception hiding (assert) +import Control.Monad +import Control.Monad.IO.Class +import Control.Monad.Identity +import Control.Monad.Trans.Class +import Data.Bifunctor +import Data.ByteString (ByteString) +import qualified Data.ByteString as B +import Data.ByteString.Unsafe +import Data.Either +import Data.Foldable +import Data.Function +import Data.Functor +import Data.List +import qualified Data.Map.Strict as M +import Data.Maybe +import qualified Data.Sequence as Seq +import Data.Serialize.Get +import Data.Serialize.Put +import qualified Data.Set as Set +import qualified Data.Vector as V +import Data.Word +import Foreign +import GHC.Conc +import qualified Streamly.Data.Fold as F +import qualified Streamly.Data.Stream.Prelude as S +import qualified Streamly.Data.Unfold as U +import Streamly.External.LMDB +import Streamly.External.LMDB.Internal +import Streamly.External.LMDB.Internal.Error +import Streamly.External.LMDB.Internal.Foreign +import qualified Streamly.Internal.Data.Fold as F +import System.Directory +import System.IO.Temp +import System.Random +import Test.QuickCheck hiding (mapSize) +import Test.QuickCheck.Monadic +import Test.Tasty (TestTree) +import Test.Tasty.HUnit +import Test.Tasty.QuickCheck hiding (mapSize) +import Text.Printf + +tests :: [TestTree] +tests = + [ testGetDatabase, + testReadLMDB, + testUnsafeReadLMDB, + testWriteLMDBChunked (ModeSerial (FoldBasedChunk False)), + testWriteLMDBChunked (ModeSerial (FoldBasedChunk True)), + testWriteLMDBChunked (ModeParallel (ClearDatabase False)), + testWriteLMDBChunked (ModeParallel (ClearDatabase True)), + testWriteLMDBOneChunk OverwrDisallowThrow, + testWriteLMDBOneChunk OverwrDisallowThrowAllowSameVal, + testWriteLMDBOneChunk OverwrDisallowIgnore, + testWriteLMDBOneChunk OverwrDisallowStop, + testWriteLMDBOneChunk OverwrAppendThrow, + testWriteLMDBOneChunk OverwrAppendIgnore, + testWriteLMDBOneChunk OverwrAppendStop, + testWriteInPlace WriteInPlaceBefore, + testWriteInPlace WriteInPlaceAfter, + testAsyncExceptionsConcurrent, + testDeleteClear, + testBetween + ] + +testGetDatabase :: TestTree +testGetDatabase = + testCase "getDatabase" $ do + -- Make sure things work concurrently on separate environments. + let nenvs = 20 + replicateConcurrently_ nenvs $ do + tmpParent <- getCanonicalTemporaryDirectory + tmpDir <- createTempDirectory tmpParent "streamly-lmdb-tests" + + -- Make sure things work concurrently in the same environment. + let ndbs = 20 + + let limits = defaultLimits {mapSize = tebibyte, maxDatabases = ndbs} + openRoEnv = openEnvironment @ReadOnly tmpDir limits + openRwEnv = openEnvironment @ReadWrite tmpDir limits + + -- Prepare environment for the first time using read-write mode. + openRwEnv >>= closeEnvironment + + -- Open environment in ReadOnly mode. + env1 <- openRoEnv + + -- Can get unnamed database in ReadOnly mode. + replicateConcurrently_ ndbs $ do + getDatabase @ReadOnly env1 Nothing >>= closeDatabase + + -- Cannot create named databases in ReadOnly mode. + forConcurrently_ [1 .. ndbs] $ \i -> do + e <- try (getDatabase @ReadOnly env1 (Just $ "name" ++ show i)) + case e of + Left (_ :: SomeException) -> return () + Right _ -> error "getDatabase: should not succeed" + + -- Close read-only environment. + closeEnvironment env1 + + -- Can create named databases in ReadWrite mode. + env2 <- openRwEnv + forConcurrently_ [1 .. ndbs] $ \i -> + getDatabase @ReadWrite env2 (Just $ "name" ++ show i) >>= closeDatabase + closeEnvironment env2 + + -- Can get existing named databases in ReadOnly mode. + env3 <- openRoEnv + forConcurrently_ [1 .. ndbs] $ \i -> + getDatabase @ReadOnly env3 (Just $ "name" ++ show i) >>= closeDatabase + closeEnvironment env3 + + removeDirectoryRecursive tmpDir + +-- | Creates a single environment with one unnamed database ('Nothing') or >=1 ('Just') named +-- databases for which the provided action is run. +withEnvDbs :: + (MonadTrans m) => + Maybe Int -> + ((Environment ReadWrite, V.Vector (Database ReadWrite)) -> m IO a) -> + m IO a +withEnvDbs mNumDbs f = do + -- TODO: bracket for MonadTrans (or PropertyM)? + tmpParent <- lift getCanonicalTemporaryDirectory + tmpDir <- lift $ createTempDirectory tmpParent "streamly-lmdb-tests" + env <- + lift $ + openEnvironment tmpDir $ + defaultLimits + { mapSize = tebibyte, + maxDatabases = fromMaybe 0 mNumDbs + } + dbs <- lift $ case mNumDbs of + Nothing -> V.singleton <$> getDatabase env Nothing + Just numDbs + | numDbs >= 1 -> forM (V.fromList [1 .. numDbs]) $ \i -> + getDatabase env (Just $ "name" ++ show i) + | otherwise -> error "invalid mNumDbs" + + a <- f (env, dbs) + + lift $ do + waitReaders env + forM_ dbs $ \db -> closeDatabase db + closeEnvironment env + removeDirectoryRecursive tmpDir + return a + +genChunkSz :: Int -> Gen (Maybe ChunkSize) +genChunkSz numKvPairs = + oneof + [ do + numPairs <- max 1 <$> chooseInt (1, 2 * numKvPairs) + return . Just $ ChunkNumPairs numPairs, + do + numBytes <- max 1 <$> chooseInt (1, 2 * 2 * kShortBsMaxLen * numKvPairs) + return . Just $ ChunkBytes numBytes, + return Nothing + ] + +-- | Write key-value pairs to a database in a normal manner, read them back using our library, and +-- make sure the result is what we wrote. Concurrent read transactions are covered. +testReadLMDB :: TestTree +testReadLMDB = + testProperty "readLMDB" . monadicIO . withEnvDbs Nothing $ \(env, dbs) -> do + let db = case V.toList dbs of [x] -> x; _ -> error "unreachable" + keyValuePairs <- toByteStrings <$> arbitraryKeyValuePairs 200 + run $ writeChunk db False keyValuePairs + let keyValuePairsInDb = sort . removeDuplicateKeysRetainLast $ keyValuePairs + + mChunkSz <- pick $ genChunkSz (length keyValuePairs) + let nread = 20 + vec <- V.replicateM nread . pick $ readOptionsAndResults keyValuePairsInDb + bs <- run $ forConcurrently vec $ \(us, readOpts, expectedResults) -> + forM readOpts $ \(ropts, ShouldSucceed ss) -> do + let unf txn = S.toList $ S.unfold readLMDB' (ropts, us, db, txn) + results <- unf $ LeftTxn mChunkSz + resultsTxn <- withReadOnlyTransaction env $ \t -> + withCursor t db $ \c -> unf $ RightTxn (t, c) + return $ boolToOp ss results expectedResults && boolToOp ss resultsTxn expectedResults + + return $ all and bs + +-- | Similar to 'testReadLMDB', except that it tests the unsafe function in a different manner. +testUnsafeReadLMDB :: TestTree +testUnsafeReadLMDB = + testProperty "unsafeReadLMDB" . monadicIO . withEnvDbs Nothing $ \(env, dbs) -> do + let db = case V.toList dbs of [x] -> x; _ -> error "unreachable" + keyValuePairs <- toByteStrings <$> arbitraryKeyValuePairs 200 + run $ writeChunk db False keyValuePairs + let keyValuePairsInDb = sort . removeDuplicateKeysRetainLast $ keyValuePairs + + mChunkSz <- pick $ genChunkSz (length keyValuePairs) + let nread = 20 + vec <- V.replicateM nread . pick $ readOptionsAndResults keyValuePairsInDb + bs <- run $ forConcurrently vec $ \(us, readOpts, expectedResults) -> + forM readOpts $ \(ropts, ShouldSucceed ss) -> do + let expectedLengths = map (bimap B.length B.length) expectedResults + let unf txn = + S.toList $ + S.unfold unsafeReadLMDB' (ropts, us, db, txn, return . snd, return . snd) + lengths <- unf $ LeftTxn mChunkSz + lengthsTxn <- withReadOnlyTransaction env $ \t -> + withCursor t db $ \c -> unf $ RightTxn (t, c) + return $ boolToOp ss lengths expectedLengths && boolToOp ss lengthsTxn expectedLengths + + return $ all and bs + +newtype ClearDatabase = ClearDatabase Bool + +-- | Note: Fold-based chunking is not meant for parallelization (as folds are stateful one-by-one +-- computations). +newtype FoldBasedChunk = FoldBasedChunk Bool + +data ChunkedMode = ModeSerial !FoldBasedChunk | ModeParallel !ClearDatabase + +instance Show ChunkedMode where + show (ModeSerial (FoldBasedChunk False)) = "serial (stream-based chunking)" + show (ModeSerial (FoldBasedChunk True)) = "serial (fold-based chunking)" + show (ModeParallel (ClearDatabase clear)) = + printf + "parallel (%s)" + (if clear then "with clearDatabase" :: String else "without clearDatabase") + +-- | Write key-value pairs to a database using our library in a chunked manner with overwriting +-- allowed, read all key-value pairs back from the database using our library (already covered by +-- 'testReadLMDB'), and make sure they are as expected. +-- +-- When parallelization is enabled: No duplicate keys are written to the database (as we can’t +-- predict the value order). This test makes sure the read-write transaction serialization mechanism +-- is working properly. While we’re at it, we also sprinkle in read-only transactions alongside the +-- read-write ones using 'readLMDB' and 'getLMDB'; and more read-write transactions from +-- 'clearDatabase'. +-- +-- In both cases, we use the opportunity to test writing the same data to a single unnamed database +-- and one or more named databases. +testWriteLMDBChunked :: ChunkedMode -> TestTree +testWriteLMDBChunked mode = + testProperty (printf "writeLMDBChunked (%s)" (show mode)) . monadicIO $ do + numDbs <- + (\x -> if x == 0 then Nothing else Just x) + <$> pick (chooseInt (0, 3)) + withEnvDbs numDbs $ \(_, dbs) -> do + -- These options should have no effect on the end-result. + us <- UseUnsafeFFI <$> pick arbitrary + let wopts = defaultWriteOptions + + -- The chunk size should also have no effect on the end-result. Note: Low chunk sizes (e.g., 1 + -- pair or 1 byte) normally result in bad performance. We therefore base the number of pairs + -- on the chunk size. + chunkByNumPairs <- pick arbitrary + (chunkSize, maxPairsBase) <- -- “base”: Not yet multiplied by numThreads (for parallel case). + if chunkByNumPairs + then do + chunkNumPairs <- pick $ chooseInt (1, 50) + + -- Assures both less and more pairs than just a single chunk. + let maxPairs = chunkNumPairs * 5 + + return (ChunkNumPairs chunkNumPairs, maxPairs) + else do + -- Compare with the ShortBs max length of 50. (For the longest possible key-value pairs + -- (100 bytes), we still have the possibility of up to two pairs per chunk.) + chunkBytes <- pick $ chooseInt (1, 200) + + -- Average ShortBs length: 50/2=25; average key-value pair length: 2*25=50. + let avgNumKeyValuePairsPerChunk = max 1 (chunkBytes `quot` 50) + let maxPairs = avgNumKeyValuePairsPerChunk * 5 + + return (ChunkBytes chunkBytes, maxPairs) + + case mode of + ModeSerial fbc -> do + -- Write key-value pairs to the database. No parallelization is desired, so duplicate keys + -- are allowed. + keyValuePairs <- toByteStrings <$> arbitraryKeyValuePairs maxPairsBase + case fbc of + FoldBasedChunk False -> + run $ + S.fromList @IO keyValuePairs + & chunkPairs chunkSize + -- Assign each chunk to each database and flatten. + & fmap (\sequ -> V.toList $ V.map (sequ,) dbs) + & S.unfoldMany U.fromList + & S.mapM (\(sequ, db) -> writeLMDBChunk' us wopts db sequ) + & S.fold F.drain + FoldBasedChunk True -> + run $ + S.fromList @IO keyValuePairs + & S.fold + ( chunkPairsFold + chunkSize + ( F.foldlM' + ( \() sequ -> + forM_ dbs $ \db -> + writeLMDBChunk' us wopts db sequ + ) + (return ()) + ) + ) + + -- Read all key-value pairs back from the databases. + readPairss <- forM dbs $ \db -> + run . S.toList $ S.unfold readLMDB (defaultReadOptions, db, LeftTxn Nothing) + + -- And make sure they are as expected. + let expectedPairsInEachDb = sort $ removeDuplicateKeysRetainLast keyValuePairs + return $ V.all (== expectedPairsInEachDb) readPairss + ModeParallel (ClearDatabase clearDb) -> do + -- Write key-value pairs to the database. Parallelization is desired, so duplicate keys + -- are eliminated up-front. + assertEnoughCapabilities + numThreads <- pick $ chooseInt (2, 2 * numCapabilities) + keyValuePairs <- + toByteStrings . removeDuplicateKeysRetainLast + <$> arbitraryKeyValuePairs (maxPairsBase * numThreads) + + chunks <- + S.fromList keyValuePairs + & chunkPairs chunkSize + & S.fold F.toList + + -- The chunk we don’t write out but at which we instead clear the database. + mClearDbChunkIdx <- + if clearDb && not (null chunks) + then Just <$> pick (chooseInt (0, length chunks - 1)) + else return Nothing + + mChunkSz <- pick $ genChunkSz (length keyValuePairs) + + readPairsAlts <- -- Alternative way of reading pairs back from the databases. + run $ + S.fromList @IO chunks + & S.indexed + -- Assign each chunk to each database and flatten. + & fmap (\(chunkIdx, sequ) -> map (chunkIdx,sequ,) [0 .. V.length dbs - 1]) + & S.unfoldMany U.fromList + & S.parMapM + (S.maxThreads numThreads . S.maxBuffer numThreads) + ( \(chunkIdx, sequ, dbIdx) -> do + let db = dbs V.! dbIdx + + -- Clear or write. + case mClearDbChunkIdx of + Just clearDbChunkIdx | clearDbChunkIdx == chunkIdx -> clearDatabase db + _ -> writeLMDBChunk' us wopts db sequ + + -- Sprinkle in read-only transactions to make sure those can coexist with + -- read-write transactions. This is also an alternative (inefficient) way of + -- reading back from the database. + b <- randomIO + (chunkIdx,dbIdx,) + <$> if b + then + -- Read entire database using 'readLMDB'. (In the Nothing ChunkSz case, + -- While the database grows as the chunks get written, these read-only + -- transactions get longer lived.) + S.toList $ S.unfold readLMDB (defaultReadOptions, db, LeftTxn mChunkSz) + else + -- Read using 'getLMDB' (not the entire database but only this chunk). + toList sequ + & S.fromList @IO + & S.mapM + ( \(k, _) -> + -- The database could have been cleared above, in which case the + -- key might not exist in the database. + getLMDB db NoTxn k >>= \mv -> return $ (k,) <$> mv + ) + & S.catMaybes + & S.fold F.toList + ) + & S.fold F.toList + + -- Read all key-value pairs back from the databases. + readPairss <- forM dbs $ \db -> + run . S.toList $ S.unfold readLMDB (defaultReadOptions, db, LeftTxn Nothing) + + -- Make sure they are as expected. (Recall that for readPairsAlts, we sometimes read the + -- whole database and sometimes only one chunk.) + return $ case mClearDbChunkIdx of + Nothing -> + -- No clearing took place. + let readPairssExpected = + V.fromList + . M.elems + -- Make sure there is at least empty data for each database. (This is for the + -- special case of no key-value pairs.) + . M.unionWith (++) (M.fromList . map (,[]) $ [0 .. V.length dbs - 1]) + -- Needed because of the parallel reading. + . M.map (sort . removeDuplicateKeysRetainLast) + . M.fromListWith (++) + $ map (\(_, dbIdx, pairs) -> (dbIdx, pairs)) readPairsAlts + in readPairss == readPairssExpected + && all (== V.head readPairss) readPairss + Just clearDbChunkIdx -> + -- For one of the chunks (clearDbChunkIdx), clearing took place. We know far less + -- about the expected data in the database. Data that was written beyond a certain + -- chunk should exist in there at least. + let readPairssExpectedSubsets = + V.fromList + . M.elems + . M.map Set.fromList + -- Make sure there is at least empty data for each database. + . M.unionWith (++) (M.fromList . map (,[]) $ [0 .. V.length dbs - 1]) + -- (No sorting or duplication removal needed because we convert to sets.) + . M.fromListWith (++) + . map (\(_, dbIdx, pairs) -> (dbIdx, pairs)) + . filter + ( \(chunkIdx, dbIdx, _) -> + -- For this database (dbIdx), find the chunkIdx beyond which data is + -- known to exist. + let (endChunkIdxTmp, endDbIdx) = + (clearDbChunkIdx * V.length dbs + dbIdx + numThreads - 1) + `quotRem` V.length dbs + endChunkIdx = + if endDbIdx < dbIdx + then endChunkIdxTmp - 1 + else endChunkIdxTmp + in chunkIdx > endChunkIdx + ) + $ readPairsAlts + in flip all (V.indexed readPairss) $ \(dbIdx, readPairs) -> + (readPairssExpectedSubsets V.! dbIdx) + `Set.isSubsetOf` Set.fromList readPairs + +-- | Choices for the standard writing accumulators. +data OverwriteOpts + = OverwrDisallowThrow + | OverwrDisallowThrowAllowSameVal + | OverwrDisallowIgnore + | OverwrDisallowStop + | OverwrAppendThrow + | OverwrAppendIgnore + | OverwrAppendStop + deriving (Show) + +-- | Write key-value pairs to a database using our library non-concurrently in one chunk (chunking +-- and concurrency is already tested with 'testWriteLMDBChunked') with the standard overwriting +-- options other than “allow,” and make sure the exceptions and written key-value pairs are as +-- expected. (We see no reason to intermingle this with chunking and concurrency.) +testWriteLMDBOneChunk :: OverwriteOpts -> TestTree +testWriteLMDBOneChunk owOpts = + testProperty (printf "writeLMDBOneChunk (%s)" (show owOpts)) . monadicIO . withEnvDbs Nothing $ + \(_, dbs) -> do + let db = case V.toList dbs of [x] -> x; _ -> error "unreachable" + + wopts = + defaultWriteOptions + { writeOverwriteOptions = case owOpts of + OverwrDisallowThrow -> OverwriteDisallow . Left $ writeAccumThrow @IO Nothing + OverwrDisallowThrowAllowSameVal -> + OverwriteDisallow . Right $ writeAccumThrowAllowSameValue Nothing + OverwrDisallowIgnore -> OverwriteDisallow $ Left writeAccumIgnore + OverwrDisallowStop -> OverwriteDisallow $ Left writeAccumStop + OverwrAppendThrow -> OverwriteAppend $ writeAccumThrow Nothing + OverwrAppendIgnore -> OverwriteAppend writeAccumIgnore + OverwrAppendStop -> OverwriteAppend writeAccumStop + } + maxPairs = 200 + + keyValuePairs <- toByteStrings <$> arbitraryKeyValuePairsDupsMoreLikely maxPairs + let hasDuplicateKeys' = hasDuplicateKeys keyValuePairs + hasDuplicateKeysWithDiffVals' = hasDuplicateKeysWithDiffVals keyValuePairs + isSorted = sort keyValuePairs == keyValuePairs + isStrictlySorted = isSorted && not hasDuplicateKeys' + + e <- run . try @SomeException $ writeLMDBChunk wopts db (Seq.fromList keyValuePairs) + + -- Make sure exceptions occurred as expected from the written key-value pairs. + let assert1 s = assertMsg $ "assert (first checks) failure " ++ s + case owOpts of + OverwrDisallowThrow -> case e of + Left _ -> assert1 "(1)" hasDuplicateKeys' + Right _ -> assert1 "(2)" $ not hasDuplicateKeys' + OverwrDisallowThrowAllowSameVal -> case e of + Left _ -> assert1 "(3)" hasDuplicateKeysWithDiffVals' + Right _ -> assert1 "(4)" $ not hasDuplicateKeysWithDiffVals' + OverwrDisallowIgnore -> assert1 "(5)" (isRight e) + OverwrDisallowStop -> assert1 "(6)" (isRight e) + OverwrAppendThrow -> case e of + Left _ -> assert1 "(7)" $ not isStrictlySorted + Right _ -> assert1 "(8)" isStrictlySorted + OverwrAppendIgnore -> assert1 "(9)" (isRight e) + OverwrAppendStop -> assert1 "(10)" (isRight e) + + -- Regardless of whether an exception occurred, read all key-value pairs back from the + -- database and make sure they are as expected. + readPairs <- run . S.toList $ S.unfold readLMDB (defaultReadOptions, db, LeftTxn Nothing) + assertMsg "assert (second checks) failure" $ + readPairs == case owOpts of + OverwrDisallowThrow -> + if hasDuplicateKeys' then [] else sort keyValuePairs + OverwrDisallowThrowAllowSameVal -> + if hasDuplicateKeysWithDiffVals' + then [] + else sort $ removeDuplicateKeysRetainLast keyValuePairs + OverwrDisallowIgnore -> sort $ filterOutReoccurringKeys keyValuePairs + OverwrDisallowStop -> sort $ prefixBeforeDuplicate keyValuePairs + OverwrAppendThrow -> + if keysAreStrictlyIncreasing keyValuePairs + then keyValuePairs + else [] + OverwrAppendIgnore -> fst $ filterGreaterThan keyValuePairs + OverwrAppendStop -> prefixBeforeStrictlySortedKeysEnd keyValuePairs + +data WriteInPlaceMode + = -- | Write new keys before existing keys. + WriteInPlaceBefore + | -- | Write new keys after existing keys. + WriteInPlaceAfter + deriving (Show) + +-- | Write key-value pairs to a database, update the values of these keys while writing more +-- key-value pairs (using the same read-write transaction), and make sure things are as expected. +testWriteInPlace :: WriteInPlaceMode -> TestTree +testWriteInPlace mode = + testCase (printf "writeInPlace (%s)" (show mode)) $ + runIdentityT $ + withEnvDbs @IdentityT Nothing $ \(env, dbs) -> do + let db = case V.toList dbs of [x] -> x; _ -> error "unreachable" + + -- Write 100,101,...200 (big-endian Word64s) keys to the database. + lift $ withReadWriteTransaction env $ \txn -> + [100 :: Word64 .. 200] + & S.fromList @IO + & fmap (\w -> (runPut . putWord64be $ w, "a")) + & S.fold (writeLMDB defaultWriteOptions db txn) + + (iteratedKeys, ()) <- + lift $ withReadWriteTransaction env $ \txn -> withCursor txn db $ \curs -> + -- Read all keys back from the database (starting at 100). + S.unfold @IO readLMDB (defaultReadOptions, db, RightTxn (txn, curs)) + & S.mapM + ( \(k, _) -> do + w <- either (error "unexpected") return $ runGet getWord64be k + return + [ Left $ Just w, -- Left: collection of iterated keys. + Right $ Just w, -- Right: updated and new pairs. + Right $ case mode of + WriteInPlaceBefore -> Just $ w - 100 -- [0 .. 100] + WriteInPlaceAfter -> + if w <= 200 -- Prevent infinite database growth. + then Just $ w + 100 -- [200 .. 300] + else Nothing + ] + ) + & S.unfoldMany U.fromList + & S.fold + ( F.partition + ( F.lmap + (fromMaybe (error "unreachable")) + F.toList + ) + ( F.catMaybes . F.lmap (\w -> (runPut . putWord64be $ w, "b")) $ + writeLMDB defaultWriteOptions db txn + ) + ) + + -- Make sure the keys that were iterated through are as expected. (Keys <=100 should not be + -- grabbed during the iteration but keys >=200 should because we start the iteration at 100 + -- and “cursor next” should grab the next word.) + lift $ + assertEqual + "iterated pairs as expected" + ( case mode of + WriteInPlaceBefore -> [100 .. 200] + WriteInPlaceAfter -> [100 .. 300] + ) + iteratedKeys + + -- Make sure all key-value pairs are as expected. + kvs <- + lift $ + S.unfold @IO readLMDB (defaultReadOptions, db, LeftTxn Nothing) + & S.fold F.toList + lift $ + assertEqual + "final pairs as expected" + ( map (\w -> (runPut . putWord64be $ w, "b")) $ case mode of + WriteInPlaceBefore -> [0 .. 200] + WriteInPlaceAfter -> [100 .. 300] + ) + kvs + +-- | Perform reads and writes on a database concurrently, throwTo threads at random, and read all +-- key-value pairs back from the database using our library (already covered by 'testReadLMDB') and +-- make sure they are as expected. +testAsyncExceptionsConcurrent :: TestTree +testAsyncExceptionsConcurrent = + testProperty "asyncExceptionsConcurrent" . monadicIO . withEnvDbs Nothing $ \(env, dbs) -> do + let db = case V.toList dbs of [x] -> x; _ -> error "unreachable" + assertEnoughCapabilities + numThreads <- pick $ chooseInt (1, 2 * numCapabilities) + + pairss <- generateConcurrentPairs numThreads 0 100 + + -- Whether to kill the thread. + shouldKills :: [Bool] <- replicateM numThreads $ pick arbitrary + + -- Delay transactions to increase possibility they will get killed while active. + delayss :: [Int] <- replicateM numThreads . pick $ chooseInt (0, 10) + + threads <- run $ + forConcurrently (zip pairss delayss) $ + \((threadIdx, pairs), delay) -> + (threadIdx,) + <$> asyncBound + -- Add a surrounding read-only transaction to test it too gets cleaned. (We could be + -- adding more possibilities, randomness, etc. here; but the idea is those things were + -- already tested elsewhere, esp. by 'testWriteLMDBChunked'.) + ( withReadOnlyTransaction env $ \_ -> withReadWriteTransaction env $ \txn -> do + S.fromList @IO pairs + & S.indexed + & S.mapM + ( \(idx, pair) -> do + when (idx == 0) $ threadDelay delay + return pair + ) + & S.fold (writeLMDB @IO defaultWriteOptions db txn) + ) + + delay <- pick $ chooseInt (0, 20) + run $ threadDelay delay + + run $ forConcurrently_ (zip shouldKills threads) $ \(shouldKill, (_, as)) -> + if shouldKill + then cancel as + else wait as + + let expectedSubset = + Set.fromList + . concatMap (\(_, (_, pairs)) -> pairs) + . filter (\(shouldKill, _) -> not shouldKill) + $ zip shouldKills pairss + + readPairs <- + Set.fromList + <$> (run . S.toList $ S.unfold readLMDB (defaultReadOptions, db, LeftTxn Nothing)) + + assertMsg "assert failure" $ expectedSubset `Set.isSubsetOf` readPairs + +testDeleteClear :: TestTree +testDeleteClear = + testCase "deleteLMDB/clearDatabase" $ + runIdentityT $ + withEnvDbs Nothing $ \(env, dbs) -> do + let db = V.head dbs + writeLMDBChunk defaultWriteOptions db $ Seq.fromList [("a", "0"), ("b", "0"), ("c", "0")] + + e1 <- liftIO . try $ withReadWriteTransaction env $ \roTxn -> + deleteLMDB defaultDeleteOptions db roTxn "foo" + case e1 of + Left (_ :: SomeException) -> error "should not throw" + Right _ -> return () + + e2 <- liftIO . try $ withReadWriteTransaction env $ \roTxn -> + deleteLMDB defaultDeleteOptions {deleteAssumeExists = True} db roTxn "foo" + case e2 of + Left (_ :: SomeException) -> return () + Right _ -> error "should throw" + + liftIO $ withReadWriteTransaction env $ \roTxn -> + deleteLMDB defaultDeleteOptions db roTxn "a" + let readPairs = + liftIO $ + S.unfold readLMDB (defaultReadOptions, db, LeftTxn Nothing) + & S.fold F.toList + kvps1 <- readPairs + unless (kvps1 == [("b", "0"), ("c", "0")]) $ error "unexpected pairs" + + liftIO $ clearDatabase db + kvps2 <- readPairs + unless (null kvps2) $ error "unexpected pairs" + +-- | A bytestring with a limited random length. (We don’t see much value in testing with longer +-- bytestrings.) +newtype ShortBS = ShortBS ByteString deriving (Eq, Ord, Show) + +kShortBsMaxLen :: Int +kShortBsMaxLen = 50 + +instance Arbitrary ShortBS where + arbitrary :: Gen ShortBS + arbitrary = do + len <- chooseInt (0, kShortBsMaxLen) + ShortBS . B.pack <$> vector len + +toByteStrings :: [(ShortBS, ShortBS)] -> [(ByteString, ByteString)] +toByteStrings = map (\(ShortBS k, ShortBS v) -> (k, v)) + +arbitraryKeyValuePairs :: Int -> PropertyM IO [(ShortBS, ShortBS)] +arbitraryKeyValuePairs maxLen = do + len <- pick $ chooseInt (0, maxLen) + filter (\(ShortBS k, _) -> not $ B.null k) -- LMDB does not allow empty keys. + <$> pick (vector len) + +-- | A variation that makes duplicate keys more likely. (The idea is to generate more relevant data +-- for testing writeLMDB with the various writeOverwriteOptions.) +arbitraryKeyValuePairsDupsMoreLikely :: Int -> PropertyM IO [(ShortBS, ShortBS)] +arbitraryKeyValuePairsDupsMoreLikely maxLen = do + arb <- arbitraryKeyValuePairs maxLen + 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 pick arbitrary + i <- pick $ chooseInt (negate $ length arb, 2 * length arb) + let (arb1, arb2) = splitAt i arb + let arb' = arb1 ++ [(k, v')] ++ arb2 + return arb' + else return arb + +-- | This function retains the last encountered value for each key. +removeDuplicateKeysRetainLast :: (Eq a) => [(a, b)] -> [(a, b)] +removeDuplicateKeysRetainLast = + 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 + +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 + +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 + 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 + Nothing -> xs + Just i -> take i xs + +-- | +-- @> filterOutReoccurringKeys [(1,"a"),(2,"b"),(2,"c"),(10,"d")] = [(1,"a"),(2,"b"),(10,"d")]@ +filterOutReoccurringKeys :: (Eq a) => [(a, b)] -> [(a, b)] +filterOutReoccurringKeys = nubBy (\(a1, _) (a2, _) -> a1 == a2) + +-- | +-- @> prefixBeforeStrictlySortedKeysEnd [(1,"a"),(2,"b"),(3,"c"),(2, "d")] = [(1,"a"),(2,"b"),(3,"c")]@ +-- @> prefixBeforeStrictlySortedKeysEnd [(1,"a"),(2,"b"),(3,"c"),(3, "d")] = [(1,"a"),(2,"b"),(3,"c")]@ +prefixBeforeStrictlySortedKeysEnd :: (Ord a) => [(a, b)] -> [(a, b)] +prefixBeforeStrictlySortedKeysEnd xs = + map fst + . takeWhile + ( \((a, _), mxs) -> + case mxs of + Nothing -> True + Just (aPrev, _) -> a > aPrev + ) + $ zip xs (Nothing : map Just xs) + +keysAreStrictlyIncreasing :: (Ord a) => [(a, b)] -> Bool +keysAreStrictlyIncreasing xs = + all + ( \((a, _), mxs) -> + case mxs of + Nothing -> True + Just (aPrev, _) -> a > aPrev + ) + $ zip xs (Nothing : map Just xs) + +-- | +-- @filterGreaterThan [(1,"a"),(2,"b"),(1,"c"),(3, "d")] = ([(1,"a"),(2,"b"),(3,"d")],[(1,"c")])@ +-- @filterGreaterThan [(1,"a"),(2,"b"),(2,"c"),(3, "d")] = [(1,"a"),(2,"b"),(3,"d"),[(2,"c")]]@ +filterGreaterThan :: (Ord a) => [(a, b)] -> ([(a, b)], [(a, b)]) +filterGreaterThan = + (\(xs, ys, _) -> (reverse xs, reverse ys)) + . foldl' + ( \(accList, accOffending, mLastIncluded) (a, b) -> case mLastIncluded of + Nothing -> ((a, b) : accList, accOffending, Just a) + Just lastIncluded -> + if a > lastIncluded + then ((a, b) : accList, accOffending, Just a) + else (accList, (a, b) : accOffending, Just lastIncluded) + ) + ([], [], Nothing) + +-- | +-- * @between first second []@ (should be called with the third parameter at @[]@). +-- * Assumes @first < second@. +-- * Returns a list in between those two. +between :: [Word8] -> [Word8] -> [Word8] -> Maybe [Word8] +-- Idea: Compare the elements from left to right, and collect the equal ones into commonPrefixRev +-- (in reverse order), until we hit inequality. +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] + +testBetween :: TestTree +testBetween = testProperty "testBetween" $ \ws1 ws2 -> + (ws1 == ws2) + || let (smaller, bigger) = if ws1 < ws2 then (ws1, ws2) else (ws2, ws1) + in case between smaller bigger [] of + Nothing -> drop (length smaller) bigger == replicate (length bigger - length smaller) 0 + Just betw -> smaller < betw && betw < bigger + +betweenBs :: ByteString -> ByteString -> Maybe ByteString +betweenBs bs1 bs2 = between (B.unpack bs1) (B.unpack bs2) [] <&> B.pack + +type PairsInDatabase = [(ByteString, ByteString)] + +type ExpectedReadResult = [(ByteString, ByteString)] + +newtype ShouldSucceed = ShouldSucceed Bool deriving (Show) + +-- | Given database pairs, randomly generates read options and corresponding expected results. +readOptionsAndResults :: + PairsInDatabase -> + Gen (UseUnsafeFFI, [(ReadOptions, ShouldSucceed)], ExpectedReadResult) +readOptionsAndResults pairsInDb = do + forw <- arbitrary + let dir = if forw then Forward else Backward + us <- UseUnsafeFFI <$> arbitrary + let len = length pairsInDb + readAll <- frequency [(1, return True), (3, return False)] + let ropts = defaultReadOptions {readDirection = dir} + if readAll + then + return + ( us, + [(ropts {readStart = if forw then ReadBeg else ReadEnd}, ShouldSucceed True)], + (if forw then id else reverse) pairsInDb + ) + else + if len == 0 + then do + bs <- arbitrary >>= \(NonEmpty ws) -> return $ B.pack ws + return + ( us, + map + (\x -> (ropts {readStart = x}, ShouldSucceed True)) + [ReadBeg, ReadEnd, ReadGE bs, ReadGT bs, ReadLE bs, ReadLT bs], + [] + ) + else do + idx <- + if len < 3 + then chooseInt (0, len - 1) + else frequency [(1, chooseInt (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 + -- 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 = + ( us, + [ (ropts {readStart = ReadGE $ keyAt idx}, ShouldSucceed True), + (ropts {readStart = ReadGT $ keyAt idx}, ShouldSucceed False) + ], + drop idx pairsInDb + ) + let backwEq = + ( us, + [ (ropts {readStart = ReadLE $ keyAt idx}, ShouldSucceed True), + (ropts {readStart = ReadLT $ keyAt idx}, ShouldSucceed False) + ], + reverse $ take (idx + 1) pairsInDb + ) + + -- Test with three desired possibilities: The chosen readStart is equal to, less than, or + -- greater than one of the keys in the database. + ord <- arbitrary @Ordering + + return $ case (ord, dir) of + (EQ, Forward) -> forwEq + (EQ, Backward) -> backwEq + (GT, Forward) -> case nextKey of + Nothing -> (us, [], []) + Just nextKey' -> + ( us, + [ (ropts {readStart = ReadGE nextKey'}, ShouldSucceed True), + (ropts {readStart = ReadGT nextKey'}, ShouldSucceed True) + ], + drop (idx + 1) pairsInDb + ) + (GT, Backward) -> case nextKey of + Nothing -> (us, [], []) + Just nextKey' -> + ( us, + [ (ropts {readStart = ReadLE nextKey'}, ShouldSucceed True), + (ropts {readStart = ReadLT nextKey'}, ShouldSucceed True) + ], + reverse $ take (idx + 1) pairsInDb + ) + (LT, Forward) -> case prevKey of + Nothing -> (us, [], []) + Just prevKey' -> + ( us, + [ (ropts {readStart = ReadGE prevKey'}, ShouldSucceed True), + (ropts {readStart = ReadGT prevKey'}, ShouldSucceed True) + ], + drop idx pairsInDb + ) + (LT, Backward) -> case prevKey of + Nothing -> (us, [], []) + Just prevKey' -> + ( us, + [ (ropts {readStart = ReadLE prevKey'}, ShouldSucceed True), + (ropts {readStart = ReadLT prevKey'}, ShouldSucceed True) + ], + reverse $ take idx pairsInDb + ) + +boolToOp :: (Eq a) => Bool -> (a -> a -> Bool) +boolToOp True = (==) +boolToOp False = (/=) + +generatePairsWithUniqueKeys :: + (Monad m) => Int -> Int -> S.Stream (PropertyM m) [(ByteString, ByteString)] +generatePairsWithUniqueKeys minPairs maxPairs = + S.fromList [0 :: Word64 ..] -- Increasing Word64s for uniqueness. + & S.foldMany + ( F.Fold + ( \(totalNumPairs, pairsSoFar, !accPairs) w -> + if pairsSoFar >= totalNumPairs + then -- Reverse to make the keys increasing. + return . F.Done $ reverse accPairs + else do + let k = runPut $ putWord64be w + ShortBS v <- pick arbitrary + return $ F.Partial (totalNumPairs, pairsSoFar + 1, (k, v) : accPairs) + ) + ( do + totalNumPairs <- pick $ chooseInt (minPairs, maxPairs) + return $ F.Partial (totalNumPairs, 0, []) + ) + (\_ -> error "unreachable") + (\_ -> error "unreachable") + ) + +newtype ThreadIdx = ThreadIdx Int deriving (Show) + +-- | Generates pairs meant to be originating from multiple threads. +-- +-- We make the keys unique; otherwise the result would be unpredictable (since we wouldn’t know +-- which values for duplicate keys would end up in the final database). +generateConcurrentPairs :: + (Monad m) => + Int -> + Int -> + Int -> + PropertyM m [(ThreadIdx, [(ByteString, ByteString)])] +generateConcurrentPairs numThreads minPairs maxPairs = + generatePairsWithUniqueKeys minPairs maxPairs + & S.take numThreads + & S.indexed -- threadIdx. + & fmap (first ThreadIdx) + & S.fold F.toList + +assertEnoughCapabilities :: (Monad m) => PropertyM m () +assertEnoughCapabilities = + when + (numCapabilities <= 1) + ( run $ + throwError + "enoughCapabilities" + "available threads <= 1; machine / Haskell RTS settings not useful for testing" + ) + +-- Writes the given key-value pairs to the given database in an ordinary way. +writeChunk :: + (Foldable t, Mode mode) => + Database mode -> + Bool -> + t (ByteString, ByteString) -> + IO () +writeChunk (Database (Environment 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 + +{-# INLINE marshalOut #-} +marshalOut :: ByteString -> (MDB_val -> IO ()) -> IO () +marshalOut bs f = + unsafeUseAsCStringLen bs $ \(ptr, len) -> f $ MDB_val (fromIntegral len) (castPtr ptr) + +assertMsg :: (Monad m) => String -> Bool -> PropertyM m () +assertMsg _ True = return () +assertMsg msg False = fail $ "Assert failed: " ++ msg diff --git a/test/TestSuite.hs b/test/TestSuite.hs --- a/test/TestSuite.hs +++ b/test/TestSuite.hs @@ -1,40 +1,16 @@ 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 qualified Streamly.External.LMDB.Tests +import Test.Tasty 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, env)) - ) - (\(tmpDir, _) -> removeDirectoryRecursive tmpDir) - (\io -> tests $ snd <$> io) +main = defaultMain tests -tests :: IO (Database ReadWrite, Environment ReadWrite) -> TestTree -tests res = +tests :: TestTree +tests = testGroup "Tests" - [ testGroup "Streamly.External.LMDB.Tests" $ - Streamly.External.LMDB.Tests.tests res + [ testGroup + "Streamly.External.LMDB.Tests" + Streamly.External.LMDB.Tests.tests ]