diff --git a/.gitignore b/.gitignore index 06be10e..29d5880 100644 --- a/.gitignore +++ b/.gitignore @@ -17,3 +17,4 @@ cabal.config .stack-work librocksdb.so *.liquid +cabal.project.local diff --git a/src/Database/RocksDB/Base.hs b/src/Database/RocksDB/Base.hs index 183b6f8..e5342b2 100644 --- a/src/Database/RocksDB/Base.hs +++ b/src/Database/RocksDB/Base.hs @@ -25,6 +25,8 @@ module Database.RocksDB.Base -- * Basic Database Manipulations , withDB , withDBCF + , open + , close , put , putCF , delete @@ -75,14 +77,16 @@ data BatchOp = Put !ByteString !ByteString -- The returned handle will be automatically released with 'close' -- when the function exits. withDB :: MonadUnliftIO m => FilePath -> Config -> (DB -> m a) -> m a -withDB path config f = - withOptions config $ \opts_ptr -> - withReadOpts Nothing $ \read_opts -> - withWriteOpts $ \write_opts -> - bracket (create_db opts_ptr read_opts write_opts) destroy_db f +withDB path config = bracket (open path config) close + +open :: MonadIO m => FilePath -> Config -> m DB +open path config = liftIO $ do + opts_ptr <- c_rocksdb_options_create + setOptions config opts_ptr + read_opts <- createReadOpts Nothing + write_opts <- c_rocksdb_writeoptions_create + create_db opts_ptr read_opts write_opts where - destroy_db db = liftIO $ - c_rocksdb_close $ rocksDB db create_db opts_ptr read_opts write_opts = do when (createIfMissing config) $ createDirectoryIfMissing True path @@ -91,10 +95,18 @@ withDB path config f = c_rocksdb_open opts_ptr path_ptr return DB { rocksDB = db_ptr , columnFamilies = [] + , opts = opts_ptr , readOpts = read_opts , writeOpts = write_opts } +close :: MonadIO m => DB -> m () +close db = liftIO $ do + c_rocksdb_writeoptions_destroy $ writeOpts db + c_rocksdb_readoptions_destroy $ readOpts db + c_rocksdb_options_destroy $ opts db + c_rocksdb_close $ rocksDB db + withDBCF :: MonadUnliftIO m => FilePath -> Config @@ -131,6 +143,7 @@ withDBCF path config cf_cfgs f = db_ptr o n return DB { rocksDB = db_ptr , columnFamilies = cfs + , opts = opts_ptr , readOpts = read_opts , writeOpts = write_opts } @@ -164,6 +177,7 @@ withDBCF path config cf_cfgs f = cfs <- peekArray (length cf_cfgs + 1) cf_ptrs_array return DB { rocksDB = db_ptr , columnFamilies = tail cfs + , opts = opts_ptr , readOpts = read_opts , writeOpts = write_opts } diff --git a/src/Database/RocksDB/Internal.hs b/src/Database/RocksDB/Internal.hs index aaa51bb..52a0f57 100644 --- a/src/Database/RocksDB/Internal.hs +++ b/src/Database/RocksDB/Internal.hs @@ -18,6 +18,8 @@ module Database.RocksDB.Internal , withOptionsCF , withReadOpts , withWriteOpts + , createReadOpts + , setOptions -- * Utilities , freeCString @@ -37,6 +39,7 @@ import UnliftIO.Foreign data DB = DB { rocksDB :: !RocksDB , columnFamilies :: ![ColumnFamily] + , opts :: !Options , readOpts :: !ReadOpts , writeOpts :: !WriteOpts } @@ -59,24 +62,8 @@ instance Default Config where } withOptions :: MonadUnliftIO m => Config -> (Options -> m a) -> m a -withOptions Config {..} f = with_opts $ \opts -> do - liftIO $ do - when bloomFilter $ do - fp <- c_rocksdb_filterpolicy_create_bloom_full 10 - bo <- c_rocksdb_block_based_options_create - c_rocksdb_block_based_options_set_filter_policy bo fp - c_rocksdb_options_set_block_based_table_factory opts bo - forM_ prefixLength $ \l -> do - t <- c_rocksdb_slicetransform_create_fixed_prefix (intToCSize l) - c_rocksdb_options_set_prefix_extractor opts t - forM_ maxFiles $ - c_rocksdb_options_set_max_open_files opts . intToCInt - c_rocksdb_options_set_create_if_missing - opts (boolToCBool createIfMissing) - c_rocksdb_options_set_error_if_exists - opts (boolToCBool errorIfExists) - c_rocksdb_options_set_paranoid_checks - opts (boolToCBool paranoidChecks) +withOptions cfg f = with_opts $ \opts -> do + liftIO $ setOptions cfg opts f opts where with_opts = @@ -84,6 +71,24 @@ withOptions Config {..} f = with_opts $ \opts -> do (liftIO c_rocksdb_options_create) (liftIO . c_rocksdb_options_destroy) +setOptions :: Config -> Options -> IO () +setOptions Config {..} opts = do + when bloomFilter $ do + fp <- c_rocksdb_filterpolicy_create_bloom_full 10 + bo <- c_rocksdb_block_based_options_create + c_rocksdb_block_based_options_set_filter_policy bo fp + c_rocksdb_options_set_block_based_table_factory opts bo + forM_ prefixLength $ \l -> do + t <- c_rocksdb_slicetransform_create_fixed_prefix (intToCSize l) + c_rocksdb_options_set_prefix_extractor opts t + forM_ maxFiles $ + c_rocksdb_options_set_max_open_files opts . intToCInt + c_rocksdb_options_set_create_if_missing + opts (boolToCBool createIfMissing) + c_rocksdb_options_set_error_if_exists + opts (boolToCBool errorIfExists) + c_rocksdb_options_set_paranoid_checks + opts (boolToCBool paranoidChecks) withOptionsCF :: MonadUnliftIO m => [Config] -> ([Options] -> m a) -> m a withOptionsCF cfgs f = @@ -95,13 +100,15 @@ withOptionsCF cfgs f = withReadOpts :: MonadUnliftIO m => Maybe Snapshot -> (ReadOpts -> m a) -> m a withReadOpts maybe_snap_ptr = bracket - create_read_opts + (liftIO $ createReadOpts maybe_snap_ptr) (liftIO . c_rocksdb_readoptions_destroy) where - create_read_opts = liftIO $ do - read_opts_ptr <- c_rocksdb_readoptions_create - forM_ maybe_snap_ptr $ c_rocksdb_readoptions_set_snapshot read_opts_ptr - return read_opts_ptr + +createReadOpts :: Maybe Snapshot -> IO ReadOpts +createReadOpts maybe_snap_ptr = do + read_opts_ptr <- c_rocksdb_readoptions_create + forM_ maybe_snap_ptr $ c_rocksdb_readoptions_set_snapshot read_opts_ptr + return read_opts_ptr withWriteOpts :: MonadUnliftIO m => (WriteOpts -> m a) -> m a withWriteOpts = diff --git a/test/Spec.hs b/test/Spec.hs index 8e7e2e0..8dcd155 100644 --- a/test/Spec.hs +++ b/test/Spec.hs @@ -26,83 +26,111 @@ withTestDBCF cfs go = withSystemTempDirectory "rocksdb-tests-cf" $ \path -> withDBCF path conf (map (,conf) cfs) go +withTestDB :: (MonadUnliftIO m) => (DB -> m a) -> m a +withTestDB go = + withSystemTempDirectory "rocksdb-tests" $ \path -> + withDB path conf go + main :: IO () -main = do - hspec $ around (withTestDBCF ["one", "two", "tree"]) $ do - describe "Database" $ do - it "puts and gets an item" $ \db -> do - put db "aaa" "zzz" - get db "aaa" `shouldReturn` Just "zzz" - it "puts and gets from different type families" $ \db -> do - let two = head $ columnFamilies db - put db "aaa_key" "aaa_value" - get db "aaa_key" `shouldReturn` Just "aaa_value" - getCF db two "aaa_key" `shouldReturn` Nothing - putCF db two "two_key" "two_value" - getCF db two "two_key" `shouldReturn` Just "two_value" - get db "two_key" `shouldReturn` Nothing - describe "Multithreading" $ do - it "stores and retrieve items from multiple threads" $ \db -> do - let key i = C.pack $ printf "key_%04d" i - val i = C.pack $ printf "val_%04d" i - indices = [0 .. 9999] :: [Int] - keys = map key indices - vals = map val indices - kvs = zip keys vals - was <- mapM (\(k, v) -> async $ put db k v) kvs - mapM_ wait was - ras <- mapM (async . get db) keys - mapM wait ras `shouldReturn` map Just vals - describe "Iterators" $ do - it "retrieves entries using iterators" $ \db -> do - let key i = C.pack $ printf "key_%03d" i - val i = C.pack $ printf "val_%03d" i - indices = [0 .. 999] :: [Int] - keys = map key indices - vals = map val indices - kvs = zip keys vals - was <- mapM (\(k, v) -> async $ put db k v) kvs - mapM_ wait was - kvs' <- withIter db $ \itr -> do - let start = keys !! 500 - iterSeek itr start - fmap catMaybes $ replicateM 500 $ do - mkv <- iterEntry itr - iterNext itr - return mkv - kvs' `shouldBe` drop 500 kvs - it "walks back and forth" $ \db -> do - withIter db $ \itr -> do - iterSeek itr "hello" - iterKey itr `shouldReturn` Nothing - iterValue itr `shouldReturn` Nothing - iterEntry itr `shouldReturn` Nothing - put db "a" "aaa" - put db "b" "bbb" - put db "c" "ccc" - withIter db $ \itr -> do - iterSeek itr "b" - iterKey itr `shouldReturn` Just "b" - iterValue itr `shouldReturn` Just "bbb" +main = hspec $ do + describe "withDBCF" $ + around (withTestDBCF ["one", "two", "tree"]) $ do + describe "Database" $ do + testPutGet + testColumnFamilies + testMultithreading + testIterators + describe "withDB (open/close)" $ + around withTestDB $ do + describe "Database" $ + testPutGet + testMultithreading + testIterators + +testPutGet :: SpecWith DB +testPutGet = + it "puts and gets an item" $ \db -> do + put db "aaa" "zzz" + get db "aaa" `shouldReturn` Just "zzz" + +testColumnFamilies :: SpecWith DB +testColumnFamilies = + it "puts and gets from different type families" $ \db -> do + let two = head $ columnFamilies db + put db "aaa_key" "aaa_value" + get db "aaa_key" `shouldReturn` Just "aaa_value" + getCF db two "aaa_key" `shouldReturn` Nothing + putCF db two "two_key" "two_value" + getCF db two "two_key" `shouldReturn` Just "two_value" + get db "two_key" `shouldReturn` Nothing + +testMultithreading :: SpecWith DB +testMultithreading = + describe "Multithreading" $ do + it "stores and retrieve items from multiple threads" $ \db -> do + let key i = C.pack $ printf "key_%04d" i + val i = C.pack $ printf "val_%04d" i + indices = [0 .. 9999] :: [Int] + keys = map key indices + vals = map val indices + kvs = zip keys vals + was <- mapM (\(k, v) -> async $ put db k v) kvs + mapM_ wait was + ras <- mapM (async . get db) keys + mapM wait ras `shouldReturn` map Just vals + +testIterators :: SpecWith DB +testIterators = + describe "Iterators" $ do + it "retrieves entries using iterators" $ \db -> do + let key i = C.pack $ printf "key_%03d" i + val i = C.pack $ printf "val_%03d" i + indices = [0 .. 999] :: [Int] + keys = map key indices + vals = map val indices + kvs = zip keys vals + was <- mapM (\(k, v) -> async $ put db k v) kvs + mapM_ wait was + kvs' <- withIter db $ \itr -> do + let start = keys !! 500 + iterSeek itr start + fmap catMaybes $ replicateM 500 $ do + mkv <- iterEntry itr iterNext itr - iterKey itr `shouldReturn` Just "c" - iterValue itr `shouldReturn` Just "ccc" - iterNext itr -- After the last entry - iterKey itr `shouldReturn` Nothing - iterValue itr `shouldReturn` Nothing - iterPrev itr -- It can't come back - iterKey itr `shouldReturn` Nothing - iterValue itr `shouldReturn` Nothing - withIter db $ \itr -> do - iterSeek itr "b" - iterKey itr `shouldReturn` Just "b" - iterValue itr `shouldReturn` Just "bbb" - iterPrev itr - iterKey itr `shouldReturn` Just "a" - iterValue itr `shouldReturn` Just "aaa" - iterPrev itr -- Invalid before lowest key - iterKey itr `shouldReturn` Nothing - iterValue itr `shouldReturn` Nothing - iterNext itr -- But it remembers previous position - iterKey itr `shouldReturn` Just "b" - iterValue itr `shouldReturn` Just "bbb" + return mkv + kvs' `shouldBe` drop 500 kvs + it "walks back and forth" $ \db -> do + withIter db $ \itr -> do + iterSeek itr "hello" + iterKey itr `shouldReturn` Nothing + iterValue itr `shouldReturn` Nothing + iterEntry itr `shouldReturn` Nothing + put db "a" "aaa" + put db "b" "bbb" + put db "c" "ccc" + withIter db $ \itr -> do + iterSeek itr "b" + iterKey itr `shouldReturn` Just "b" + iterValue itr `shouldReturn` Just "bbb" + iterNext itr + iterKey itr `shouldReturn` Just "c" + iterValue itr `shouldReturn` Just "ccc" + iterNext itr -- After the last entry + iterKey itr `shouldReturn` Nothing + iterValue itr `shouldReturn` Nothing + iterPrev itr -- It can't come back + iterKey itr `shouldReturn` Nothing + iterValue itr `shouldReturn` Nothing + withIter db $ \itr -> do + iterSeek itr "b" + iterKey itr `shouldReturn` Just "b" + iterValue itr `shouldReturn` Just "bbb" + iterPrev itr + iterKey itr `shouldReturn` Just "a" + iterValue itr `shouldReturn` Just "aaa" + iterPrev itr -- Invalid before lowest key + iterKey itr `shouldReturn` Nothing + iterValue itr `shouldReturn` Nothing + iterNext itr -- But it remembers previous position + iterKey itr `shouldReturn` Just "b" + iterValue itr `shouldReturn` Just "bbb"