Allow ResourceT for iterators and snapshots

Author: jprupp

Diff for: CHANGELOG.md

@@ -1,3 +1,7 @@
+## 2.1.4
+
+ * Add ResourceT support for snapshots and iterators.
+
 ## 2.1.3
 
  * Correct bug where database was ignoring prefix length config.

Diff for: package.yaml

@@ -1,5 +1,5 @@
 name: rocksdb-haskell-jprupp
-version: 2.1.3
+version: 2.1.4
 synopsis: Haskell bindings for RocksDB
 description: See README at <https://github.com/jprupp/rocksdb-haskell#readme>
 maintainer: Jean-Pierre Rupp <[email protected]>
@@ -19,6 +19,7 @@ dependencies:
   - bytestring
   - data-default
   - directory
+  - resourcet
   - unliftio
 
 library:

Diff for: rocksdb-haskell-jprupp.cabal

@@ -1,13 +1,13 @@
 cabal-version: 1.12
 
--- This file has been generated from package.yaml by hpack version 0.33.0.
+-- This file has been generated from package.yaml by hpack version 0.34.4.
 --
 -- see: https://github.com/sol/hpack
 --
--- hash: 094c0ad3842293bfb0c2b11fda87cf96076e2f497046a7e82023bfd4d96d4e0d
+-- hash: 48b49def634c00156ec5f8bc2f4a8c92caf0aca334fca3cef6507f51febe9867
 
 name:           rocksdb-haskell-jprupp
-version:        2.1.3
+version:        2.1.4
 synopsis:       Haskell bindings for RocksDB
 description:    See README at <https://github.com/jprupp/rocksdb-haskell#readme>
 category:       Database, FFI
@@ -45,6 +45,7 @@ library
     , bytestring
     , data-default
     , directory
+    , resourcet
     , unliftio
   default-language: Haskell2010
 
@@ -63,6 +64,7 @@ test-suite spec
     , data-default
     , directory
     , hspec
+    , resourcet
     , rocksdb-haskell-jprupp
     , string-conversions
     , unliftio

Diff for: src/Database/RocksDB/Base.hs

@@ -33,6 +33,9 @@ module Database.RocksDB.Base
     , get
     , getCF
     , withSnapshot
+    , snapshot
+    , createSnapshot
+    , releaseSnapshot
 
     -- * Administrative Functions
     , Property (..), getProperty
@@ -44,7 +47,7 @@ module Database.RocksDB.Base
     , module Database.RocksDB.Iterator
     ) where
 
-import           Control.Monad             (when, (>=>), forM)
+import           Control.Monad             (forM, when, (>=>))
 import           Data.ByteString           (ByteString)
 import qualified Data.ByteString           as BS
 import           Data.ByteString.Internal  (ByteString (..))
@@ -55,6 +58,7 @@ import           Database.RocksDB.Iterator
 import           UnliftIO
 import           UnliftIO.Directory
 import           UnliftIO.Foreign
+import           UnliftIO.Resource
 
 -- | Properties exposed by RocksDB
 data Property = NumFilesAtLevel Int | Stats | SSTables
@@ -166,14 +170,25 @@ withDBCF path config cf_cfgs f =
 -- | Run an action with a snapshot of the database.
 -- The 'DB' object is not valid after the action ends.
 withSnapshot :: MonadUnliftIO m => DB -> (DB -> m a) -> m a
-withSnapshot db@DB{rocksDB = db_ptr} f =
-    bracket create_snapshot release_snapshot (f . fst)
-  where
-    release_snapshot = liftIO . c_rocksdb_release_snapshot db_ptr . snd
-    create_snapshot = liftIO $ do
-        snap_ptr <- c_rocksdb_create_snapshot db_ptr
-        withReadOpts (Just snap_ptr) $ \read_opts ->
-            return (db{readOpts = read_opts}, snap_ptr)
+withSnapshot db f =
+    bracket (createSnapshot db) releaseSnapshot (f . fst)
+
+-- | The 'DB' snapshot is not valid outside of 'MonadResource'.
+snapshot :: (MonadIO m, MonadResource m) => DB -> m DB
+snapshot db =
+    fst . snd <$> allocate (createSnapshot db) releaseSnapshot
+
+-- | Manually create an unmanaged snapshot.
+createSnapshot :: MonadIO m => DB -> m (DB, Snapshot)
+createSnapshot db@DB{rocksDB = db_ptr} = liftIO $ do
+    snap_ptr <- c_rocksdb_create_snapshot db_ptr
+    withReadOpts (Just snap_ptr) $ \read_opts ->
+        return (db{readOpts = read_opts}, snap_ptr)
+
+-- | Function to release an unmanaged snapshot.
+releaseSnapshot :: MonadIO m => (DB, Snapshot) -> m ()
+releaseSnapshot (DB{rocksDB = db_ptr}, snap_ptr) =
+    liftIO $ c_rocksdb_release_snapshot db_ptr snap_ptr
 
 -- | Get a DB property.
 getProperty :: MonadIO m => DB -> Property -> m (Maybe ByteString)
@@ -340,5 +355,5 @@ withStrings :: MonadUnliftIO m => [String] -> ([CString] -> m a) -> m a
 withStrings ss f =
     go [] ss
   where
-    go acc [] = f (reverse acc)
+    go acc []     = f (reverse acc)
     go acc (x:xs) = withCString x $ \p -> go (p:acc) xs

Diff for: src/Database/RocksDB/Iterator.hs

@@ -15,6 +15,11 @@ module Database.RocksDB.Iterator
     ( Iterator
     , withIter
     , withIterCF
+    , iter
+    , iterCF
+    , iterator
+    , createIterator
+    , destroyIterator
     , iterEntry
     , iterFirst
     , iterGetError
@@ -39,6 +44,7 @@ import           Foreign.C.Error           (throwErrnoIfNull)
 import           Foreign.C.String          (CString, peekCString)
 import           Foreign.C.Types           (CSize)
 import           UnliftIO
+import           UnliftIO.Resource
 
 -- | Create 'Iterator' and use it.
 --
@@ -53,6 +59,13 @@ withIter db = withIterCommon db Nothing
 withIterCF :: MonadUnliftIO m => DB -> ColumnFamily -> (Iterator -> m a) -> m a
 withIterCF db cf = withIterCommon db (Just cf)
 
+-- | Variation on 'iterator' below.
+iter :: (MonadIO m, MonadResource m) => DB -> m Iterator
+iter db = iterator db Nothing
+
+iterCF :: (MonadIO m, MonadResource m) => DB -> ColumnFamily -> m Iterator
+iterCF db cf = iterator db (Just cf)
+
 withIterCommon :: MonadUnliftIO m
                => DB
                -> Maybe ColumnFamily
@@ -67,6 +80,23 @@ withIterCommon DB{rocksDB = rocks_db, readOpts = read_opts} mcf =
         Just cf -> c_rocksdb_create_iterator_cf rocks_db read_opts cf
         Nothing -> c_rocksdb_create_iterator rocks_db read_opts
 
+-- | Iterator is not valid outside of 'ResourceT' context.
+iterator :: (MonadIO m, MonadResource m)
+         => DB -> Maybe ColumnFamily -> m Iterator
+iterator db mcf =
+    snd <$> allocate (createIterator db mcf) destroyIterator
+
+-- | Manually create unmanaged iterator.
+createIterator :: MonadIO m => DB -> Maybe ColumnFamily -> m Iterator
+createIterator DB{rocksDB = rocks_db, readOpts = read_opts} mcf = liftIO $
+    throwErrnoIfNull "create_iterator" $ case mcf of
+        Just cf -> c_rocksdb_create_iterator_cf rocks_db read_opts cf
+        Nothing -> c_rocksdb_create_iterator rocks_db read_opts
+
+-- | Destroy unmanaged iterator.
+destroyIterator :: MonadIO m => Iterator -> m ()
+destroyIterator = liftIO . c_rocksdb_iter_destroy
+
 -- | An iterator is either positioned at a key/value pair, or not valid. This
 -- function returns /true/ iff the iterator is valid.
 iterValid :: MonadIO m => Iterator -> m Bool
@@ -127,9 +157,9 @@ iterValue = liftIO . flip iterString c_rocksdb_iter_value
 -- | Return the current entry as a pair, if the iterator is currently positioned
 -- at an entry, ie. 'iterValid'.
 iterEntry :: MonadIO m => Iterator -> m (Maybe (ByteString, ByteString))
-iterEntry iter = liftIO $ do
-    mkey <- iterKey iter
-    mval <- iterValue iter
+iterEntry it = liftIO $ do
+    mkey <- iterKey it
+    mval <- iterValue it
     return $ (,) <$> mkey <*> mval
 
 -- | Check for errors