ppad-lmdb-0.2.0: test/Main.hs

{-# LANGUAGE OverloadedStrings #-}

module Main where

import Control.Concurrent
  ( forkIO
  , isCurrentThreadBound
  , newEmptyMVar
  , putMVar
  , rtsSupportsBoundThreads
  , takeMVar
  , yield
  )
import qualified Control.Exception as E
import Control.Monad (forM, forM_, replicateM_)
import qualified Data.ByteString as BS
import qualified Data.ByteString.Char8 as BSC
import qualified Data.Map.Strict as Map
import qualified Database.LMDB as L
import System.Directory
  ( createDirectory
  , getTemporaryDirectory
  , removeDirectoryRecursive
  )
import System.FilePath ((</>))
import System.IO.Error (isAlreadyExistsError)
import Test.Tasty
import Test.Tasty.HUnit
import Test.Tasty.QuickCheck as QC

main :: IO ()
main = defaultMain $ testGroup "ppad-lmdb"
  [ unitTests
  , propTests
  ]

-- helpers --------------------------------------------------------------------

-- | Run an action against a fresh temporary directory. The directory
--   is unique to the call, so tests may run concurrently, and is
--   removed afterwards even if the action throws.
withTmpDir :: (FilePath -> IO a) -> IO a
withTmpDir k = do
  tmpRoot <- getTemporaryDirectory
  E.bracket (fresh tmpRoot (0 :: Int)) removeDirectoryRecursive k
  where
    fresh root n = do
      let dir = root </> ("ppad-lmdb-test-" ++ show n)
      r <- E.try (createDirectory dir)
      case r of
        Right () -> pure dir
        Left e
          | isAlreadyExistsError e -> fresh root (n + 1)
          | otherwise              -> E.throwIO e

withTmpEnv :: L.EnvFlags -> (L.Env -> IO a) -> IO a
withTmpEnv envFlags k = withTmpDir $ \dir -> L.withEnv dir envFlags k

-- | 'withTmpEnv', with the default database opened.
withTmpDb :: L.EnvFlags -> (L.Env -> L.Dbi -> IO a) -> IO a
withTmpDb envFlags k = withTmpEnv envFlags $ \env ->
  L.openDbi env Nothing True >>= k env

-- shared bigger map for tests
flags :: L.EnvFlags
flags = L.defaultEnvFlags
  { L.envMapSize = 64 * 1024 * 1024
  , L.envMaxDbs  = 4
  }

writeKV
  :: L.Env -> L.Dbi -> [(BS.ByteString, BS.ByteString)] -> IO ()
writeKV env dbi kvs = L.withWriteTxn env $ \txn ->
  forM_ kvs $ \(k, v) -> L.put txn dbi k v

readKey :: L.Env -> L.Dbi -> BS.ByteString -> IO (Maybe BS.ByteString)
readKey env dbi k = L.withReadTxn env $ \txn -> L.get txn dbi k

-- | Run each action in its own forkIO thread and wait for all of
--   them, rethrowing the first failure.
forkAll :: [IO ()] -> IO ()
forkAll acts = do
  dones <- forM acts $ \act -> do
    done <- newEmptyMVar
    _ <- forkIO $
      (E.try act :: IO (Either E.SomeException ())) >>= putMVar done
    pure done
  rs <- mapM takeMVar dones
  either E.throwIO pure (sequence_ rs)

-- unit tests -----------------------------------------------------------------

unitTests :: TestTree
unitTests = testGroup "unit"
  [ testCase "put then get roundtrip" $
      withTmpDb flags $ \env dbi -> do
        writeKV env dbi [("hello", "world")]
        v <- L.withReadTxn env $ \txn ->
          L.get txn dbi "hello"
        v @?= Just "world"

  , testCase "get on missing key returns Nothing" $
      withTmpDb flags $ \env dbi -> do
        writeKV env dbi [("a", "1")]
        v <- L.withReadTxn env $ \txn ->
          L.get txn dbi "absent"
        v @?= Nothing

  , testCase "del returns True only when key existed" $
      withTmpDb flags $ \env dbi -> do
        writeKV env dbi [("a", "1"), ("b", "2")]
        (existed, repeated, missing) <- L.withWriteTxn env $ \txn -> do
          e   <- L.del txn dbi "a"
          e2  <- L.del txn dbi "a"
          m   <- L.del txn dbi "never-existed"
          pure (e, e2, m)
        existed  @?= True
        repeated @?= False
        missing  @?= False

  , testCase "write txn aborted on exception" $
      withTmpDb flags $ \env dbi -> do
        -- attempt that throws after a write
        let attempt = L.withWriteTxn env $ \txn -> do
              L.put txn dbi "k" "v"
              E.throwIO (userError "boom")
        r <- E.try attempt :: IO (Either IOError ())
        case r of
          Left _  -> pure ()
          Right _ -> assertFailure "expected the user error"
        -- key must not be visible to a fresh read txn
        v <- L.withReadTxn env $ \txn ->
          L.get txn dbi "k"
        v @?= Nothing

  , testCase "failed commit leaves env usable" $
      withTmpDb flags { L.envMapSize = 1024 * 1024 } $ \env dbi -> do
        -- a caught MapFull poisons the txn, so its commit must fail
        r <- E.try $ L.withWriteTxn env $ \txn -> do
          full <- E.try $ forM_ [0 :: Int ..] $ \i ->
            L.put txn dbi (BSC.pack (show i)) (BS.replicate 512 0)
          case full of
            Left L.LMDBMapFull -> pure ()
            _ -> assertFailure "expected LMDBMapFull"
        case r of
          Left (L.LMDBOther (-30782) _) -> pure () -- MDB_BAD_TXN
          _ -> assertFailure ("expected MDB_BAD_TXN, got " ++ show r)
        forM_ [1 :: Int .. 3] $ \_ -> writeKV env dbi [("k", "v")]
        v <- L.withReadTxn env $ \txn ->
          L.get txn dbi "k"
        v @?= Just "v"

  , testCase "dbi name is used exactly as given" $
      withTmpEnv flags $ \env -> do
        -- a slice: the bytes after it are not a NUL terminator
        _ <- L.openDbi env (Just (BS.take 3 "foobar")) True
        r <- E.try (L.openDbi env (Just "foobar") False)
        case r of
          Left L.LMDBNotFound -> pure ()
          _ -> assertFailure "opened a database that was never created"
        _ <- L.openDbi env (Just "foo") False
        pure ()

  , testCase "dbi name containing NUL is rejected" $
      withTmpEnv flags $ \env -> do
        r <- E.try (L.openDbi env (Just "foo\0bar") True)
        case r of
          Left (L.LMDBOther 22 _) -> pure () -- EINVAL
          _ -> assertFailure "expected EINVAL"

  , testCase "named databases are independent" $
      withTmpEnv flags $ \env -> do
        a <- L.openDbi env (Just "a") True
        b <- L.openDbi env (Just "b") True
        L.withWriteTxn env $ \txn -> do
          L.put txn a "k" "1"
          L.put txn b "k" "2"
        vs <- L.withReadTxn env $ \txn ->
          (,) <$> L.get txn a "k" <*> L.get txn b "k"
        vs @?= (Just "1", Just "2")

  , testCase "concurrent openDbi calls" $
      withTmpEnv flags $ \env ->
        forkAll
          [ do dbi <- L.openDbi env (Just name) True
               writeKV env dbi [("k", name)]
               v <- readKey env dbi "k"
               v @?= Just name
          | name <- ["w", "x", "y", "z"] ]

  , testCase "read-only env serves reads, refuses writes" $
      withTmpDir $ \dir -> do
        L.withEnv dir flags $ \env -> do
          dbi <- L.openDbi env Nothing True
          writeKV env dbi [("k", "v")]
        L.withEnv dir flags { L.envReadOnly = True } $ \env -> do
          dbi <- L.openDbi env Nothing False
          v <- readKey env dbi "k"
          v @?= Just "v"
          w <- E.try (L.withWriteTxn env (\_ -> pure ()))
          case w of
            Left (L.LMDBOther 13 _) -> pure () -- EACCES
            _ -> assertFailure "expected EACCES from withWriteTxn"
          c <- E.try (L.openDbi env (Just "new") True)
          case c of
            Left (L.LMDBOther 13 _) -> pure ()
            _ -> assertFailure "expected EACCES from openDbi"

  , testCase "read txns nest" $
      withTmpDb flags $ \env dbi -> do
        writeKV env dbi [("k", "v")]
        v <- L.withReadTxn env $ \_ -> readKey env dbi "k"
        v @?= Just "v"

  , testCase "concurrent readers see their own snapshots" $
      withTmpDb flags $ \env dbi -> do
        writeKV env dbi [("k", "old")]
        entered <- newEmptyMVar
        written <- newEmptyMVar
        seen    <- newEmptyMVar
        forkAll
          [ L.withReadTxn env $ \txn -> do
              putMVar entered ()
              takeMVar written
              L.get txn dbi "k" >>= putMVar seen
          , do takeMVar entered
               -- a second reader, while the first is live
               old <- readKey env dbi "k"
               old @?= Just "old"
               writeKV env dbi [("k", "new")]
               putMVar written ()
          ]
        stale <- takeMVar seen
        stale @?= Just "old"
        fresh <- readKey env dbi "k"
        fresh @?= Just "new"

  , testCase "concurrent writers serialize" $
      withTmpDb flags { L.envNoSync = True } $ \env dbi -> do
        let bump = L.withWriteTxn env $ \txn -> do
              n <- maybe 0 (maybe 0 fst . BSC.readInt)
                     <$> L.get txn dbi "n"
              yield -- invite another writer in mid-txn
              L.put txn dbi "n" (BSC.pack (show (n + 1)))
        forkAll (replicate 8 (replicateM_ 50 bump))
        n <- readKey env dbi "n"
        n @?= Just "400"

  , testCase "write txns run bound where the RTS allows" $
      withTmpEnv flags $ \env -> do
        bound <- newEmptyMVar
        forkAll
          [ L.withWriteTxn env (\_ -> isCurrentThreadBound)
              >>= putMVar bound ]
        b <- takeMVar bound
        b @?= rtsSupportsBoundThreads

  , testCase "open rejects out-of-range flags" $
      forM_ [ flags { L.envMapSize = -1 }
            , flags { L.envMaxDbs  = -1 }
            , flags { L.envMaxDbs  = 2 ^ (32 :: Int) }
            ] $ \bad -> do
        r <- E.try (withTmpEnv bad (\_ -> pure ()))
        case r of
          Left (L.LMDBOther 22 _) -> pure () -- EINVAL
          _ -> assertFailure ("expected EINVAL for " ++ show bad)

  , testCase "sync flushes a no-sync env" $
      withTmpDb flags { L.envNoSync = True } $ \env dbi -> do
        writeKV env dbi [("k", "v")]
        L.sync env

  , testCase "cursor walks keys in ascending order" $
      withTmpDb flags $ \env dbi -> do
        let kvs = [ (BSC.pack [c], BSC.pack [c])
                  | c <- "abcdefghij" ]
        writeKV env dbi kvs
        seen <- L.withReadTxn env $ \txn ->
          L.withCursor txn dbi (cursorWalk [])
        seen @?= kvs

  , testCase "cursorSeek lands on first >= key" $
      withTmpDb flags $ \env dbi -> do
        writeKV env dbi [ ("a","1"), ("c","3"), ("e","5") ]
        r <- L.withReadTxn env $ \txn ->
          L.withCursor txn dbi $ \cur ->
            L.cursorSeek cur "b"
        r @?= Just ("c", "3")

  , testCase "cursorSeek past the last key finds nothing" $
      withTmpDb flags $ \env dbi -> do
        writeKV env dbi [ ("a","1"), ("c","3") ]
        r <- L.withReadTxn env $ \txn ->
          L.withCursor txn dbi $ \cur ->
            L.cursorSeek cur "d"
        r @?= Nothing

  , testCase "cursor walks keys in descending order" $
      withTmpDb flags $ \env dbi -> do
        let kvs = [ (BSC.pack [c], BSC.pack [c])
                  | c <- "abcdefghij" ]
        writeKV env dbi kvs
        seen <- L.withReadTxn env $ \txn ->
          L.withCursor txn dbi $ \cur -> do
            let go acc Nothing   = pure (reverse acc)
                go acc (Just kv) = L.cursorPrev cur >>= go (kv : acc)
            L.cursorLast cur >>= go []
        seen @?= reverse kvs

  , testCase "cursors on an empty database find nothing" $
      withTmpDb flags $ \env dbi -> do
        r <- L.withReadTxn env $ \txn ->
          L.withCursor txn dbi $ \cur ->
            (,,) <$> L.cursorFirst cur
                 <*> L.cursorLast cur
                 <*> L.cursorSeek cur "a"
        r @?= (Nothing, Nothing, Nothing)
  ]

-- walk a cursor, accumulating pairs from current position to end
cursorWalk
  :: [(BS.ByteString, BS.ByteString)]
  -> L.Cursor s
  -> IO [(BS.ByteString, BS.ByteString)]
cursorWalk acc cur = do
  step <- L.cursorFirst cur
  go (reverse acc) step
  where
    go ys Nothing      = pure (reverse ys)
    go ys (Just kv)    = do
      next <- L.cursorNext cur
      go (kv : ys) next

-- properties -----------------------------------------------------------------

newtype Key = Key BS.ByteString deriving Show
newtype Val = Val BS.ByteString deriving Show

-- LMDB rejects empty keys; restrict to non-empty bounded strings.
instance Arbitrary Key where
  arbitrary = do
    n <- QC.chooseInt (1, 32)
    Key . BS.pack <$> QC.vectorOf n QC.arbitrary

instance Arbitrary Val where
  arbitrary = do
    n <- QC.chooseInt (0, 64)
    Val . BS.pack <$> QC.vectorOf n QC.arbitrary

propTests :: TestTree
propTests = testGroup "props"
  [ QC.testProperty "put / get roundtrip" $ \(Key k) (Val v) ->
      QC.ioProperty $ withTmpDb flags $ \env dbi -> do
        writeKV env dbi [(k, v)]
        got <- L.withReadTxn env $ \txn ->
          L.get txn dbi k
        pure (got == Just v)

  , QC.testProperty "cursor enumerates Map.toAscList" $ \kvs0 ->
      let kvs = uniqueKeys [(k, v) | (Key k, Val v) <- kvs0]
      in  QC.ioProperty $ withTmpDb flags $ \env dbi -> do
            writeKV env dbi kvs
            seen <- L.withReadTxn env $ \txn ->
              L.withCursor txn dbi (cursorWalk [])
            pure (seen == Map.toAscList (Map.fromList kvs))
  ]

-- keep only the last write for each key, matching LMDB's overwrite
-- semantics under MDB_NOOVERWRITE-not-set
uniqueKeys
  :: [(BS.ByteString, BS.ByteString)]
  -> [(BS.ByteString, BS.ByteString)]
uniqueKeys = Map.toList . Map.fromList