cas batch calls for RocksDbTable

libs/chainweb-storage/chainweb-storage.cabal

@@ -74,6 +74,7 @@ test-suite rocksdb-tests
         , bytestring >=0.10
         , exceptions >=0.10
         , lens >=4.16
+        , vector >=0.12
 
 executable dedup-profiling
     hs-source-dirs: tools

libs/chainweb-storage/src/Data/CAS/RocksDB.hs

@@ -107,6 +107,7 @@ import qualified Data.ByteString.Char8 as B8
 import Data.CAS
 import Data.String
 import qualified Data.Text as T
+import qualified Data.Vector as V
 
 import qualified Database.RocksDB.Base as R
 import qualified Database.RocksDB.Iterator as I
@@ -576,9 +577,17 @@ instance (IsCasValue v, CasKeyType v ~ k) => IsCas (RocksDbTable k v) where
     casInsert db a = tableInsert db (casKey a) a
     casDelete = tableDelete
 
+    casInsertBatch db vs = updateBatch (mkOp <$> V.toList vs)
+      where
+        mkOp v = RocksDbInsert db (casKey v) v
+
+    casDeleteBatch db vs = updateBatch (RocksDbDelete db <$> V.toList vs)
+
     {-# INLINE casLookup #-}
     {-# INLINE casInsert #-}
     {-# INLINE casDelete #-}
+    {-# INLINE casInsertBatch #-}
+    {-# INLINE casDeleteBatch #-}
 
 -- | A newtype wrapper that takes only a single type constructor. This useful in
 -- situations where a Higher Order type constructor for a CAS is required. A
@@ -593,10 +602,14 @@ instance IsCasValue v => IsCas (RocksDbCas v) where
     casLookup (RocksDbCas x) = casLookup x
     casInsert (RocksDbCas x) = casInsert x
     casDelete (RocksDbCas x) = casDelete x
+    casInsertBatch (RocksDbCas x) = casInsertBatch x
+    casDeleteBatch (RocksDbCas x) = casDeleteBatch x
 
     {-# INLINE casLookup #-}
     {-# INLINE casInsert #-}
     {-# INLINE casDelete #-}
+    {-# INLINE casInsertBatch #-}
+    {-# INLINE casDeleteBatch #-}
 
 -- | Create a new 'RocksDbCas'.
 --

libs/chainweb-storage/test/Main.hs

@@ -1,5 +1,9 @@
+{-# LANGUAGE OverloadedLists #-}
 {-# LANGUAGE OverloadedStrings #-}
 {-# LANGUAGE TypeApplications #-}
+{-# LANGUAGE TypeFamilies #-}
+
+{-# OPTIONS_GHC -fno-warn-orphans #-}
 
 -- |
 -- Module: Main
@@ -8,7 +12,7 @@
 -- Maintainer: Lars Kuhtz <[email protected]>
 -- Stability: experimental
 --
--- A few basic tests for "Data.CAS.RocksDB"
+-- Tests for "Data.CAS.RocksDB"
 --
 module Main
 ( main
@@ -21,12 +25,17 @@ import Control.Monad
 import Control.Monad.Catch
 
 import qualified Data.ByteString.Char8 as B8
+import Data.Foldable
 import Data.List
+import Data.Vector ({- IsList Vector #-})
+
+import GHC.Stack
 
 import Text.Read
 
 -- internal modules
 
+import Data.CAS
 import Data.CAS.RocksDB
 
 -- -------------------------------------------------------------------------- --
@@ -35,14 +44,15 @@ import Data.CAS.RocksDB
 data RocksDbTableTestException
     = CodecException String
     | RocksDbTableTestFailure String
-    deriving (Show, Eq)
+    deriving (Show)
 
 instance Exception RocksDbTableTestException
 
-assertIO :: Eq a => Show a => IO a -> a -> IO ()
+assertIO :: HasCallStack => Eq a => Show a => IO a -> a -> IO ()
 assertIO f r = f >>= \a -> unless (a == r) $ throwM
     $ RocksDbTableTestFailure
-    $ "test failed:\n  expected: " <> show a <> "\n  actual: " <> show r
+    $ "test failed:\n  expected: " <> show r <> "\n  actual: " <> show a
+    <> " " <> prettyCallStack callStack
 
 -- -------------------------------------------------------------------------- --
 -- Test Table
@@ -58,12 +68,12 @@ intTable db tableName = newTable db intCodec intCodec [tableName]
 -- -------------------------------------------------------------------------- --
 -- Tests
 
-assertEmptyTable :: RocksDbTable Int Int -> IO ()
+assertEmptyTable :: HasCallStack => RocksDbTable Int Int -> IO ()
 assertEmptyTable t = do
     assertIO (tableLookup t 1) Nothing
     assertEntries t []
 
-assertEntries :: RocksDbTable Int Int -> [(Int, Int)] -> IO ()
+assertEntries :: HasCallStack => RocksDbTable Int Int -> [(Int, Int)] -> IO ()
 assertEntries t l_ = do
     forM_ l $ \(k,v) ->
         assertIO (tableLookup t k) (Just v)
@@ -97,27 +107,31 @@ assertEntries t l_ = do
   where
     l = sort l_
 
-tableTests :: RocksDb -> B8.ByteString -> IO ()
+tableTests :: HasCallStack => RocksDb -> B8.ByteString -> IO ()
 tableTests db tableName = do
-
     assertEmptyTable t
+
     tableInsert t 1 8
     assertEntries t [(1,8)]
+
     tableInsert t 2 9
     assertEntries t [(1,8), (2,9)]
+
     tableDelete t 1
     assertEntries t [(2,9)]
+
     tableInsert t 2 8
     assertEntries t [(2,8)]
+
     tableDelete t 2
     assertEmptyTable t
   where
     t = intTable db tableName
 
-tableBatchTests :: RocksDb -> B8.ByteString -> IO ()
+tableBatchTests :: HasCallStack => RocksDb -> B8.ByteString -> IO ()
 tableBatchTests db tableName = do
-
     assertEmptyTable t
+
     updateBatch []
     assertEmptyTable t
 
@@ -126,32 +140,140 @@ tableBatchTests db tableName = do
 
     updateBatch [RocksDbInsert t 2 9]
     assertEntries t [(1,8), (2,9)]
+
     updateBatch [RocksDbDelete t 2]
     assertEntries t [(1,8)]
+
     updateBatch [RocksDbInsert t 2 9, RocksDbDelete t 2]
     assertEntries t [(1,8)]
+
     updateBatch [RocksDbInsert t 2 9, RocksDbDelete t 2, RocksDbInsert t 2 9]
     assertEntries t [(1,8), (2,9)]
+
     updateBatch [RocksDbInsert t 1 8, RocksDbDelete t 1]
     assertEntries t [(2,9)]
+
     updateBatch [RocksDbInsert t 1 7, RocksDbInsert t 1 8, RocksDbInsert t 1 8]
     assertEntries t [(1,8), (2,9)]
+
     updateBatch [RocksDbDelete t 1, RocksDbInsert t 3 7]
     assertEntries t [(2,9), (3,7)]
+
     updateBatch [RocksDbInsert t 4 6, RocksDbInsert t 5 5]
     assertEntries t [(2,9), (3,7), (4,6), (5,5)]
+
     updateBatch [RocksDbDelete t 2, RocksDbDelete t 3, RocksDbDelete t 4, RocksDbDelete t 5]
     assertEmptyTable t
   where
     t = intTable db tableName
 
+
+-- Orphan instance
+--
+instance IsCasValue Int where
+    type CasKeyType Int = Int
+    casKey = (+10)
+    {-# INLINE casKey #-}
+
+casBatchTests :: HasCallStack => RocksDb -> B8.ByteString -> IO ()
+casBatchTests db tableName = do
+    assertEmptyTable t
+
+    casInsertBatch t mempty
+    assertEmptyTable t
+
+    casInsertBatch t [1]
+    assertCasEntries t [1]
+
+    casInsertBatch t [2]
+    assertCasEntries t [1, 2]
+
+    casDeleteBatch t [casKey @Int 2]
+    assertCasEntries t [1]
+
+    casInsertBatch t [1]
+    assertCasEntries t [1]
+
+    casInsertBatch t [2, 2, 2]
+    assertCasEntries t [1, 2]
+
+    casInsertBatch t [1,2,3,4]
+    assertCasEntries t [1, 2, 3, 4]
+
+    casDeleteBatch t [casKey @Int 5]
+    assertCasEntries t [1, 2, 3, 4]
+
+    casDeleteBatch t $ casKey @Int <$> [1, 3, 1]
+    assertCasEntries t [2, 4]
+
+    casDeleteBatch t []
+    assertCasEntries t [2, 4]
+
+    casDeleteBatch t $ casKey @Int <$> [2, 4]
+    assertEmptyTable t
+  where
+    t = intTable db tableName
+
+casTests :: HasCallStack => RocksDb -> B8.ByteString -> IO ()
+casTests db tableName = do
+    assertEmptyTable t
+    assertIO (casMember t 1) False
+    assertIO (casLookup t 1) Nothing
+
+    casInsertBatch t mempty
+    assertEmptyTable t
+
+    casInsert t 1
+    assertCasEntries t [1]
+    assertIO (casMember t $ casKey @Int 1) True
+    assertIO (casLookup t $ casKey @Int 1) (Just 1)
+
+    casInsert t 2
+    assertCasEntries t [1, 2]
+    assertIO (casMember t $ casKey @Int 1) True
+    assertIO (casMember t $ casKey @Int 2) True
+    assertIO (casLookup t $ casKey @Int 1) (Just 1)
+    assertIO (casLookup t $ casKey @Int 2) (Just 2)
+
+    casDelete t $ casKey @Int 2
+    assertCasEntries t [1]
+    assertIO (casMember t $ casKey @Int 1) True
+    assertIO (casMember t $ casKey @Int 2) False
+    assertIO (casLookup t $ casKey @Int 1) (Just 1)
+    assertIO (casLookup t $ casKey @Int 2) Nothing
+
+    casInsert t 1
+    assertCasEntries t [1]
+
+    traverse_ @[] (casInsert t) [2, 2, 2]
+    assertCasEntries t [1, 2]
+
+    traverse_ @[] (casInsert t) [1,2,3,4]
+    assertCasEntries t [1, 2, 3, 4]
+
+    casDelete t $ casKey @Int 5
+    assertCasEntries t [1, 2, 3, 4]
+
+    traverse_ @[] (casDelete t) $ casKey @Int <$> [1, 3, 1]
+    assertCasEntries t [2, 4]
+
+    traverse_ @[] (casDelete t) $ casKey @Int <$> [2, 4]
+    assertEmptyTable t
+  where
+    t = intTable db tableName
+
+assertCasEntries :: HasCallStack => RocksDbTable Int Int -> [Int] -> IO ()
+assertCasEntries t = assertEntries t . fmap (\x -> (casKey x, x))
+
 -- -------------------------------------------------------------------------- --
 -- Main
 
 main :: IO ()
 main = withTempRocksDb "testDb" $ \db -> do
     mapConcurrently_
         (\i -> tableTests db $ "testTable" <> B8.pack (show i))
-        [0..100 :: Int]
+        ([0..100] :: [Int])
     tableBatchTests db "testTable"
+    casTests db "testTable"
+    casBatchTests db "testTable"