open and close database to use it without bracket

.gitignore

@@ -17,3 +17,4 @@ cabal.config
 .stack-work
 librocksdb.so
 *.liquid
+cabal.project.local

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
                               }

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,31 +62,33 @@ 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 =
         bracket
         (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 =

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"