diff --git a/ChangeLog.md b/ChangeLog.md --- a/ChangeLog.md +++ b/ChangeLog.md @@ -1,3 +1,7 @@ +## 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. + ## 0.0.1.1 * Fixed `install-includes` and `include-dirs` in the Cabal file. 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,7 @@ Mode, ReadWrite, ReadOnly, + OverwriteOptions(..), WriteOptions(..), -- ** Environment and database @@ -37,15 +38,16 @@ import Control.Concurrent (isCurrentThreadBound, myThreadId) import Control.Concurrent.Async (asyncBound, wait) import Control.Exception (Exception, catch, tryJust, mask_, throw) -import Control.Monad (guard, when) +import Control.Monad (guard, unless, when) import Control.Monad.IO.Class (MonadIO, liftIO) import Data.ByteString (ByteString, packCStringLen) -import Data.ByteString.Unsafe (unsafeUseAsCStringLen) +import Data.ByteString.Unsafe (unsafePackCStringLen, unsafeUseAsCStringLen) import Data.Maybe (fromJust) import Data.Void (Void) -import Foreign (Ptr, free, malloc, nullPtr, peek) +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 Streamly.Internal.Data.Fold (Fold (Fold)) import Streamly.Internal.Data.Stream.StreamD (newFinalizedIORef, runIORefFinalizer) import Streamly.Internal.Data.Stream.StreamD.Type (Step (Stop, Yield)) @@ -197,15 +199,17 @@ return (pcurs, pk, pv, ref) return (op, pcurs, pk, pv, ref)) +data OverwriteOptions = OverwriteAllow | OverwriteAllowSame | OverwriteDisallow deriving (Eq) + data WriteOptions = WriteOptions { writeTransactionSize :: !Int - , noOverwrite :: !Bool + , overwriteOptions :: !OverwriteOptions , writeAppend :: !Bool } defaultWriteOptions :: WriteOptions defaultWriteOptions = WriteOptions { writeTransactionSize = 1 - , noOverwrite = False + , overwriteOptions = OverwriteAllow , writeAppend = False } @@ -225,7 +229,10 @@ writeLMDB :: (MonadIO m) => Database ReadWrite -> WriteOptions -> Fold m (ByteString, ByteString) () writeLMDB (Database penv dbi) options = let txnSize = max 1 (writeTransactionSize options) - flags = combineOptions $ [mdb_nooverwrite | noOverwrite options] ++ [mdb_append | writeAppend options] + overwriteOpt = overwriteOptions options + flags = combineOptions $ + [mdb_nooverwrite | overwriteOpt `elem` [OverwriteAllowSame, OverwriteDisallow]] + ++ [mdb_append | writeAppend options] in Fold (\(threadId, iter, currChunkSz, mtxn) (k, v) -> do -- In the first few iterations, ascertain that we are still on the same (bound) thread. iter' <- liftIO $ @@ -256,8 +263,21 @@ liftIO $ unsafeUseAsCStringLen k $ \(kp, kl) -> unsafeUseAsCStringLen v $ \(vp, vl) -> catch (mdb_put_ ptxn dbi kp (fromIntegral kl) vp (fromIntegral vl) flags) - (\(e :: LMDB_Error) -> runIORefFinalizer ref >> throw e) - + (\(e :: LMDB_Error) -> do + -- Discard error if OverwriteAllowSame was specified and the error from LMDB + -- was due to the exact same key-value pair already existing in the database. + ok <- with (MDB_val (fromIntegral kl) kp) $ \pk -> + alloca $ \pv -> do + rc <- c_mdb_get ptxn dbi pk pv + if rc == 0 then do + v' <- peek pv + vbs <- unsafePackCStringLen (mv_data v', fromIntegral $ mv_size v') + return $ overwriteOpt == OverwriteAllowSame + && e_code e == Right MDB_KEYEXIST + && vbs == v + else + return False + unless ok $ runIORefFinalizer ref >> throw e) return (threadId, iter', currChunkSz' + 1, Just (ptxn, ref))) (do isBound <- liftIO isCurrentThreadBound 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 @@ -79,6 +79,9 @@ foreign import ccall unsafe "lmdb.h mdb_cursor_close" c_mdb_cursor_close :: Ptr MDB_cursor -> IO () +foreign import ccall unsafe "lmdb.h mdb_get" + c_mdb_get :: Ptr MDB_txn -> MDB_dbi_t -> Ptr MDB_val -> Ptr MDB_val -> IO CInt + foreign import ccall unsafe "lmdb.h mdb_put" c_mdb_put :: Ptr MDB_txn -> MDB_dbi_t -> Ptr MDB_val -> Ptr MDB_val -> CUInt -> IO CInt @@ -116,7 +119,7 @@ | MDB_BAD_TXN | MDB_BAD_VALSIZE | MDB_BAD_DBI - deriving (Show) + deriving (Eq, Show) {-# INLINE errCodes #-} errCodes :: [(MDB_ErrCode, Int)] 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: d4449c158999353b1715fa4877104b424051746bec5dccfc6adaf535b84e0fba +-- hash: 43ef967b54a793b5ad742a0e6d24b951d3961ac437567bc8756786917e48d678 name: streamly-lmdb -version: 0.0.1.1 +version: 0.1.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.0.1.1 + , streamly-lmdb ==0.1.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 @@ -10,6 +10,7 @@ import Data.List (find, foldl', nubBy, sort) import Foreign (castPtr, nullPtr, with) import Streamly.Prelude (fromList, toList, unfold) +import Test.QuickCheck (choose) import Test.QuickCheck.Monadic (PropertyM, monadicIO, pick, run) import Test.Tasty (TestTree) import Test.Tasty.QuickCheck (arbitrary, testProperty) @@ -17,14 +18,19 @@ import qualified Data.ByteString as B import qualified Streamly.Prelude as S -import Streamly.External.LMDB (Mode, ReadWrite, WriteOptions(..), +import Streamly.External.LMDB (Mode, ReadWrite, OverwriteOptions(..), WriteOptions(..), clearDatabase, defaultWriteOptions, readLMDB, unsafeReadLMDB, writeLMDB) import Streamly.External.LMDB.Internal (Database (..)) import Streamly.External.LMDB.Internal.Foreign (MDB_val (..), combineOptions, mdb_nooverwrite, mdb_put, mdb_txn_begin, mdb_txn_commit) tests :: IO (Database ReadWrite) -> [TestTree] -tests db = [ testReadLMDB db, testUnsafeReadLMDB db, testWriteLMDB db, testWriteLMDB_2 db ] +tests db = + [ testReadLMDB db + , testUnsafeReadLMDB db + , testWriteLMDB db + , testWriteLMDB_2 db + , testWriteLMDB_3 db ] -- | 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. @@ -72,7 +78,8 @@ chunkSz <- pick arbitrary - let fol' = writeLMDB db $ defaultWriteOptions { writeTransactionSize = chunkSz, noOverwrite = False } + let fol' = writeLMDB db $ defaultWriteOptions { writeTransactionSize = chunkSz + , overwriteOptions = OverwriteAllow } -- TODO: Run with new "bound" functionality in streamly. run $ asyncBound (S.fold fol' (fromList keyValuePairs)) >>= wait @@ -88,19 +95,18 @@ testWriteLMDB_2 :: IO (Database ReadWrite) -> TestTree testWriteLMDB_2 res = testProperty "writeLMDB_2" . monadicIO $ do db <- run res - keyValuePairs <- arbitraryKeyValuePairs + keyValuePairs <- arbitraryKeyValuePairs' run $ clearDatabase db chunkSz <- pick arbitrary -- TODO: Run with new "bound" functionality in streamly. - let fol' = writeLMDB db $ defaultWriteOptions { writeTransactionSize = chunkSz, noOverwrite = True } + let fol' = writeLMDB db $ defaultWriteOptions { writeTransactionSize = chunkSz + , overwriteOptions = OverwriteDisallow } e <- run $ try @SomeException $ (asyncBound (S.fold fol' (fromList keyValuePairs)) >>= wait) exceptionAsExpected <- case e of - Left _ -> - -- This is somewhat rare but will be encountered for enough number of tests (e.g., 1000 instead of 100). - return $ hasDuplicateKeys keyValuePairs + Left _ -> return $ hasDuplicateKeys keyValuePairs Right _ -> return . not $ hasDuplicateKeys keyValuePairs let keyValuePairsInDb = sort . prefixBeforeDuplicate $ keyValuePairs @@ -109,12 +115,55 @@ return $ exceptionAsExpected && pairsAsExpected +-- | Clear the database, write key-value pairs to it using our library with key overwriting +-- disallowed except when attempting to replace an existing key-value pair, and make sure an +-- exception occurs iff we had a duplicate key with different values in our pairs. Furthermore +-- make sure that key-value pairs prior to a such a duplicate key are actually in the database. +testWriteLMDB_3 :: IO (Database ReadWrite) -> TestTree +testWriteLMDB_3 res = testProperty "writeLMDB_3" . monadicIO $ do + db <- run res + keyValuePairs <- arbitraryKeyValuePairs' + run $ clearDatabase db + + chunkSz <- pick arbitrary + + -- TODO: Run with new "bound" functionality in streamly. + let fol' = writeLMDB db $ defaultWriteOptions { writeTransactionSize = chunkSz + , overwriteOptions = OverwriteAllowSame } + e <- run $ try @SomeException $ (asyncBound (S.fold fol' (fromList keyValuePairs)) >>= wait) + exceptionAsExpected <- + case e of + Left _ -> return $ hasDuplicateKeysWithDiffVals keyValuePairs + Right _ -> return . not $ hasDuplicateKeysWithDiffVals keyValuePairs + + let keyValuePairsInDb = sort . removeDuplicateKeys . prefixBeforeDuplicateWithDiffVal $ keyValuePairs + readPairsAll <- run . toList $ unfold (readLMDB db) 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 length arb > 0 && b then do + let (k, v) = arb !! 0 + 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 + -- | 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 @@ -124,9 +173,21 @@ 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