{-# LANGUAGE BinaryLiterals #-}
{-# LANGUAGE ImportQualifiedPost #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE TupleSections #-}
import Control.Concurrent (forkIO, killThread, threadDelay)
import Control.Monad
import Data.ByteString.Char8 qualified as C
import Data.Default (def)
import Data.Maybe
import Database.RocksDB
import Test.Hspec
import Text.Printf
import UnliftIO
conf :: Config
conf =
def
{ createIfMissing = True,
errorIfExists = True,
bloomFilter = True,
prefixLength = Just 3
}
withTestDBCF :: (MonadUnliftIO m) => [String] -> (DB -> m a) -> m a
withTestDBCF cfs go =
withSystemTempDirectory "rocksdb-tests-cf" $ \path ->
withDBCF path conf (map (,conf) cfs) 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"
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"
describe "Snapshots" $ do
it "getCF and iterator see same data during concurrent writes" $ \db -> do
let cf = head $ columnFamilies db
-- Initial data
putCF db cf "key1" "val1"
putCF db cf "key2" "val2"
putCF db cf "key3" "val3"
-- Create snapshot
(snapDB, snap) <- createSnapshot db
-- Start writer thread modifying the same keys
started <- newEmptyMVar
writerThread <- forkIO $ do
putMVar started ()
let loop n
| n > 100 = pure ()
| otherwise = do
let suffix = C.pack $ "-v" <> show n
putCF db cf "key1" ("val1" <> suffix)
putCF db cf "key2" ("val2" <> suffix)
putCF db cf "key3" ("val3" <> suffix)
threadDelay 500
loop (n + 1 :: Int)
loop (0 :: Int)
takeMVar started
threadDelay 1000 -- Let writer get going
-- Read via getCF on snapshot
v1 <- getCF snapDB cf "key1"
threadDelay 10000
v2 <- getCF snapDB cf "key2"
threadDelay 10000
v3 <- getCF snapDB cf "key3"
-- Read via iterator on same snapshot (using the same column family)
iterEntries <- withIterCF snapDB cf $ \itr -> do
iterFirst itr
let collect acc = do
valid <- iterValid itr
if valid
then do
mentry <- iterEntry itr
iterNext itr
collect (mentry : acc)
else pure (reverse acc)
catMaybes <$> collect []
-- Clean up
killThread writerThread
releaseSnapshot (snapDB, snap)
-- Verify: getCF and iterator see the same original values
v1 `shouldBe` Just "val1"
v2 `shouldBe` Just "val2"
v3 `shouldBe` Just "val3"
iterEntries `shouldBe` [("key1", "val1"), ("key2", "val2"), ("key3", "val3")]
it "withIterSnap sees snapshot data" $ \db -> do
put db "a" "1"
put db "b" "2"
(_, snap) <- createSnapshot db
-- Modify after snapshot
put db "a" "modified"
put db "c" "3"
-- Iterator with explicit snapshot should see original data
entries <- withIterSnap db (Just snap) $ \itr -> do
iterFirst itr
let collect acc = do
valid <- iterValid itr
if valid
then do
mentry <- iterEntry itr
iterNext itr
collect (mentry : acc)
else pure (reverse acc)
catMaybes <$> collect []
entries `shouldBe` [("a", "1"), ("b", "2")]
it "withIterSnapCF sees snapshot data on column family" $ \db -> do
let cf = head $ columnFamilies db
putCF db cf "x" "100"
putCF db cf "y" "200"
(_, snap) <- createSnapshot db
-- Modify after snapshot
putCF db cf "x" "modified"
putCF db cf "z" "300"
-- Iterator with explicit snapshot on CF
entries <- withIterSnapCF db (Just snap) cf $ \itr -> do
iterFirst itr
let collect acc = do
valid <- iterValid itr
if valid
then do
mentry <- iterEntry itr
iterNext itr
collect (mentry : acc)
else pure (reverse acc)
catMaybes <$> collect []
entries `shouldBe` [("x", "100"), ("y", "200")]
it "createIteratorSnap with manual management" $ \db -> do
put db "m" "10"
put db "n" "20"
(_, snap) <- createSnapshot db
put db "m" "changed"
-- Manual iterator creation with snapshot
(itr, mReadOpts) <- createIteratorSnap db (Just snap) Nothing
iterFirst itr
k1 <- iterKey itr
v1 <- iterValue itr
iterNext itr
k2 <- iterKey itr
v2 <- iterValue itr
destroyIterator itr
forM_ mReadOpts destroyReadOpts
(k1, v1) `shouldBe` (Just "m", Just "10")
(k2, v2) `shouldBe` (Just "n", Just "20")