diff --git a/ChangeLog.md b/ChangeLog.md --- a/ChangeLog.md +++ b/ChangeLog.md @@ -1,3 +1,7 @@ +## 0.2.0 + +* Added read options for backward reading and starting from a specific key. + ## 0.1.0 * Added a write option that disallows overwriting but does not throw an error when attempting to write a key-value pair that already exists in the database. diff --git a/README.md b/README.md --- a/README.md +++ b/README.md @@ -20,9 +20,9 @@ module Main where import Streamly.External.LMDB - (Limits (mapSize), WriteOptions (chunkSize), defaultLimits, - defaultWriteOptions, getDatabase, openEnvironment, readLMDB, - tebibyte, writeLMDB) + (Limits (mapSize), WriteOptions (writeTransactionSize), defaultLimits, + defaultReadOptions, defaultWriteOptions, getDatabase, openEnvironment, + readLMDB, tebibyte, writeLMDB) import qualified Streamly.Prelude as S main :: IO () @@ -38,7 +38,7 @@ db <- getDatabase env Nothing -- Stream key-value pairs into the database. - let fold' = writeLMDB db defaultWriteOptions { chunkSize = 1 } + let fold' = writeLMDB db defaultWriteOptions { writeTransactionSize = 1 } let writeStream = S.fromList [("baz", "a"), ("foo", "b"), ("bar", "c")] _ <- S.fold fold' writeStream @@ -48,7 +48,7 @@ -- ("bar","c") -- ("baz","a") -- ("foo","b") - let unfold' = readLMDB db + let unfold' = readLMDB db defaultReadOptions let readStream = S.unfold unfold' undefined S.mapM_ print readStream ``` 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 @@ -13,6 +13,8 @@ Mode, ReadWrite, ReadOnly, + ReadDirection(..), + ReadOptions(..), OverwriteOptions(..), WriteOptions(..), @@ -28,6 +30,7 @@ clearDatabase, -- ** Reading + defaultReadOptions, readLMDB, unsafeReadLMDB, @@ -48,6 +51,7 @@ import Foreign.C (Errno (Errno), eNOTDIR) import Foreign.C.String (CStringLen) import Foreign.Marshal.Utils (with) +import Foreign.Storable (poke) import Streamly.Internal.Data.Fold (Fold (Fold)) import Streamly.Internal.Data.Stream.StreamD (newFinalizedIORef, runIORefFinalizer) import Streamly.Internal.Data.Stream.StreamD.Type (Step (Stop, Yield)) @@ -147,7 +151,21 @@ mdb_clear ptxn dbi mdb_txn_commit ptxn) >>= wait --- | Creates an unfold with which we can stream all key-value pairs from the given database in increasing key order. +-- | Direction of key iteration. +data ReadDirection = Forward | Backward deriving (Show) + +data ReadOptions = ReadOptions + { readDirection :: !ReadDirection + -- | If 'Nothing', a forward [backward] iteration starts at the beginning [end] of the database. + -- Otherwise, it starts at the first key that is greater [less] than or equal to the 'Just' key. + , readStart :: !(Maybe ByteString) } deriving (Show) + +defaultReadOptions :: ReadOptions +defaultReadOptions = ReadOptions + { readDirection = Forward + , readStart = Nothing } + +-- | Creates an unfold with which we can stream key-value pairs from the given database. -- -- A read transaction is kept open for the duration of the unfold; one should therefore -- bear in mind LMDB's [caveats regarding long-lived transactions](https://git.io/JJZE6). @@ -155,10 +173,10 @@ -- 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 -> Unfold m Void (ByteString, ByteString) -readLMDB db = unsafeReadLMDB db packCStringLen packCStringLen +readLMDB :: (MonadIO m, Mode mode) => Database mode -> ReadOptions -> Unfold m Void (ByteString, ByteString) +readLMDB db ropts = unsafeReadLMDB db ropts packCStringLen packCStringLen --- | Creates an unfold with which we can stream all key-value pairs from the given database in increasing key order. +-- | Creates an unfold with which we can stream key-value pairs from the given database. -- -- A read transaction is kept open for the duration of the unfold; one should therefore -- bear in mind LMDB's [caveats regarding long-lived transactions](https://git.io/JJZE6). @@ -168,20 +186,50 @@ -- 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 -> (CStringLen -> IO k) -> (CStringLen -> IO v) -> Unfold m Void (k, v) -unsafeReadLMDB (Database penv dbi) kmap vmap = - flip supply mdb_first $ Unfold + => Database mode + -> ReadOptions + -> (CStringLen -> IO k) + -> (CStringLen -> IO v) + -> Unfold m Void (k, v) +unsafeReadLMDB (Database penv dbi) ropts kmap vmap = + let (firstOp, subsequentOp) = case (readDirection ropts, readStart ropts) of + (Forward, Nothing) -> (mdb_first, mdb_next) + (Forward, Just _) -> (mdb_set_range, mdb_next) + (Backward, Nothing) -> (mdb_last, mdb_prev) + (Backward, Just _) -> (mdb_set_range, mdb_prev) + in flip supply firstOp $ Unfold (\(op, pcurs, pk, pv, ref) -> do - found <- liftIO $ c_mdb_cursor_get pcurs pk pv op >>= \rc -> + rc <- liftIO $ + if op == mdb_set_range && subsequentOp == mdb_prev then do + -- Reverse MDB_SET_RANGE. + kfst' <- peek pk + kfst <- packCStringLen (mv_data kfst', fromIntegral $ mv_size kfst') + rc <- c_mdb_cursor_get pcurs pk pv op + if rc /= 0 && rc == mdb_notfound then + c_mdb_cursor_get pcurs pk pv mdb_last + else if rc == 0 then do + k' <- peek pk + k <- unsafePackCStringLen (mv_data k', fromIntegral $ mv_size k') + if k /= kfst then + c_mdb_cursor_get pcurs pk pv mdb_prev + else + return rc + else + return rc + else + c_mdb_cursor_get pcurs pk pv op + + found <- liftIO $ if rc /= 0 && rc /= mdb_notfound then do runIORefFinalizer 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) (mdb_next, pcurs, pk, pv, ref) + return $ Yield (k, v) (subsequentOp, pcurs, pk, pv, ref) else do runIORefFinalizer ref return Stop) @@ -191,6 +239,12 @@ pcurs <- liftIO $ mdb_cursor_open ptxn dbi pk <- liftIO malloc pv <- liftIO malloc + + _ <- case readStart ropts of + Nothing -> return () + Just k -> unsafeUseAsCStringLen k $ \(kp, kl) -> + poke pk (MDB_val (fromIntegral kl) kp) + ref <- liftIO . newFinalizedIORef $ do free pv >> free pk c_mdb_cursor_close pcurs 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 @@ -188,8 +188,17 @@ mdb_first :: MDB_cursor_op_t mdb_first = #const MDB_FIRST +mdb_last :: MDB_cursor_op_t +mdb_last = #const MDB_LAST + mdb_next :: MDB_cursor_op_t mdb_next = #const MDB_NEXT + +mdb_prev :: MDB_cursor_op_t +mdb_prev = #const MDB_PREV + +mdb_set_range :: MDB_cursor_op_t +mdb_set_range = #const MDB_SET_RANGE mdb_env_create :: IO (Ptr MDB_env) mdb_env_create = do diff --git a/streamly-lmdb.cabal b/streamly-lmdb.cabal --- a/streamly-lmdb.cabal +++ b/streamly-lmdb.cabal @@ -4,10 +4,10 @@ -- -- see: https://github.com/sol/hpack -- --- hash: 43ef967b54a793b5ad742a0e6d24b951d3961ac437567bc8756786917e48d678 +-- hash: 9ad1d377f90069c99f24bcc557c1ff43a9f075fdfead1071725615fe8f352f50 name: streamly-lmdb -version: 0.1.0 +version: 0.2.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 @@ -70,7 +70,7 @@ , bytestring >=0.10.10.0 && <0.11 , directory >=1.3.6.0 && <1.4 , streamly >=0.7.2 && <0.8 - , streamly-lmdb ==0.1.0 + , streamly-lmdb ==0.2.0 , tasty >=1.2.3 && <1.3 , tasty-quickcheck >=0.10.1.1 && <0.11 , temporary >=1.3 && <1.4 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 @@ -5,12 +5,13 @@ import Control.Concurrent.Async (asyncBound, wait) import Control.Exception (SomeException, onException, try) import Control.Monad (forM_) -import Data.ByteString (ByteString, pack) +import Data.ByteString (ByteString, pack, unpack) import Data.ByteString.Unsafe (unsafeUseAsCStringLen) import Data.List (find, foldl', nubBy, sort) +import Data.Word (Word8) import Foreign (castPtr, nullPtr, with) import Streamly.Prelude (fromList, toList, unfold) -import Test.QuickCheck (choose) +import Test.QuickCheck (NonEmptyList (..), Gen, choose, elements, frequency) import Test.QuickCheck.Monadic (PropertyM, monadicIO, pick, run) import Test.Tasty (TestTree) import Test.Tasty.QuickCheck (arbitrary, testProperty) @@ -18,8 +19,8 @@ import qualified Data.ByteString as B import qualified Streamly.Prelude as S -import Streamly.External.LMDB (Mode, ReadWrite, OverwriteOptions(..), WriteOptions(..), - clearDatabase, defaultWriteOptions, readLMDB, unsafeReadLMDB, writeLMDB) +import Streamly.External.LMDB (Mode, ReadWrite, OverwriteOptions (..), ReadDirection (..), ReadOptions (..), + WriteOptions(..), clearDatabase, defaultReadOptions, defaultWriteOptions, readLMDB, unsafeReadLMDB, writeLMDB) import Streamly.External.LMDB.Internal (Database (..)) import Streamly.External.LMDB.Internal.Foreign (MDB_val (..), combineOptions, mdb_nooverwrite, mdb_put, mdb_txn_begin, mdb_txn_commit) @@ -30,43 +31,40 @@ , testUnsafeReadLMDB db , testWriteLMDB db , testWriteLMDB_2 db - , testWriteLMDB_3 db ] + , testWriteLMDB_3 db + , testBetween ] -- | Clear the database, write key-value pairs to it in a normal manner, read -- them back using our library, and make sure the result is what we wrote. testReadLMDB :: (Mode mode) => IO (Database mode) -> TestTree testReadLMDB res = testProperty "readLMDB" . monadicIO $ do db <- run res - keyValuePairs <- arbitraryKeyValuePairs + keyValuePairs <- arbitraryKeyValuePairs'' run $ clearDatabase db run $ writeChunk db False keyValuePairs let keyValuePairsInDb = sort . removeDuplicateKeys $ keyValuePairs - let stream = unfold (readLMDB db) undefined - readPairsAll <- run . toList $ stream - let allAsExpected = readPairsAll == keyValuePairsInDb - - -- Also make sure the stream’s 'take' functionality works. - readPairsFirstFew <- run . toList $ S.take 3 stream - let firstFewAsExpected = readPairsFirstFew == take 3 keyValuePairsInDb + (readOpts, expectedResults) <- pick $ readOptionsAndResults keyValuePairsInDb + results <- run . toList $ unfold (readLMDB db readOpts) undefined - return $ allAsExpected && firstFewAsExpected + return $ results == expectedResults -- | Similar to 'testReadLMDB', except that it tests the unsafe function in a different manner. testUnsafeReadLMDB :: (Mode mode) => IO (Database mode) -> TestTree testUnsafeReadLMDB res = testProperty "unsafeReadLMDB" . monadicIO $ do db <- run res - keyValuePairs <- arbitraryKeyValuePairs + keyValuePairs <- arbitraryKeyValuePairs'' run $ clearDatabase db run $ writeChunk db False keyValuePairs - let lengthsInDb = map (\(k, v) -> (B.length k, B.length v)) . sort . removeDuplicateKeys $ keyValuePairs + let keyValuePairsInDb = sort . removeDuplicateKeys $ keyValuePairs - let stream = unfold (unsafeReadLMDB db (return . snd) (return . snd)) undefined - readLengthsAll <- run . toList $ stream + (readOpts, expectedResults) <- pick $ readOptionsAndResults keyValuePairsInDb + let expectedLengths = map (\(k, v) -> (B.length k, B.length v)) expectedResults + lengths <- run . toList $ unfold (unsafeReadLMDB db readOpts (return . snd) (return . snd)) undefined - return $ readLengthsAll == lengthsInDb + return $ lengths == 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. @@ -85,7 +83,7 @@ run $ asyncBound (S.fold fol' (fromList keyValuePairs)) >>= wait let keyValuePairsInDb = sort . removeDuplicateKeys $ keyValuePairs - readPairsAll <- run . toList $ unfold (readLMDB db) undefined + readPairsAll <- run . toList $ unfold (readLMDB db defaultReadOptions) undefined return $ keyValuePairsInDb == readPairsAll @@ -110,7 +108,7 @@ Right _ -> return . not $ hasDuplicateKeys keyValuePairs let keyValuePairsInDb = sort . prefixBeforeDuplicate $ keyValuePairs - readPairsAll <- run . toList $ unfold (readLMDB db) undefined + readPairsAll <- run . toList $ unfold (readLMDB db defaultReadOptions) undefined let pairsAsExpected = keyValuePairsInDb == readPairsAll return $ exceptionAsExpected && pairsAsExpected @@ -137,7 +135,7 @@ Right _ -> return . not $ hasDuplicateKeysWithDiffVals keyValuePairs let keyValuePairsInDb = sort . removeDuplicateKeys . prefixBeforeDuplicateWithDiffVal $ keyValuePairs - readPairsAll <- run . toList $ unfold (readLMDB db) undefined + readPairsAll <- run . toList $ unfold (readLMDB db defaultReadOptions) undefined let pairsAsExpected = keyValuePairsInDb == readPairsAll return $ exceptionAsExpected && pairsAsExpected @@ -164,6 +162,22 @@ 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 @@ -191,6 +205,76 @@ 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 + let len = length pairsInDb + readAll <- frequency [ (1, return True), (3, return False) ] + let ropts = defaultReadOptions { readDirection = dir } + if readAll then + return (ropts { readStart = Nothing }, (if forw then id else reverse) pairsInDb) + else if len == 0 then do + bs <- arbitrary >>= \(NonEmpty ws) -> return $ pack ws + return (ropts { readStart = Just bs }, []) + else do + idx <- if len < 3 then choose (0, len - 1) else + frequency [ (1, choose (1, len - 2)), (3, elements [0, len - 1]) ] + let keyAt i = fst $ pairsInDb !! i + let nextKey + | idx + 1 <= len - 1 = betweenBs (keyAt idx) (keyAt $ idx + 1) + | otherwise = Just $ keyAt (len - 1) `B.append` B.singleton 0 + let prevKey + | idx == 0 && keyAt idx /= B.singleton 0 = Just $ B.singleton 0 -- Keys are known to be non-empty. + | 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 ()