ppad-lmdb-0.2.0: bench/Weight.hs

{-# LANGUAGE BangPatterns      #-}
{-# LANGUAGE OverloadedStrings #-}

module Main where

import Control.DeepSeq (force)
import Control.Exception (evaluate)
import Control.Monad (forM_)
import qualified Data.ByteString as BS
import qualified Data.ByteString.Builder as BB
import qualified Data.ByteString.Lazy as BSL
import qualified Database.LMDB as L
import System.Directory
  ( createDirectory
  , doesDirectoryExist
  , getTemporaryDirectory
  , removeDirectoryRecursive
  )
import System.FilePath ((</>))
import Weigh

main :: IO ()
main = do
  tmpRoot <- getTemporaryDirectory
  let dir = tmpRoot </> "ppad-lmdb-weigh"
  exists <- doesDirectoryExist dir
  if exists then removeDirectoryRecursive dir else pure ()
  createDirectory dir
  let envFlags = L.defaultEnvFlags
                   { L.envMapSize = 64 * 1024 * 1024
                   , L.envMaxDbs  = 4
                   }
  -- build keys and values outside the weighed region
  kv1 <- evaluate (force (kvs 1000))
  ks1 <- evaluate (force (map fst kv1))
  L.withEnv dir envFlags $ \env -> do
    dbi <- L.openDbi env Nothing True
    -- pre-populate; the put case then overwrites these keys
    insert env dbi kv1
    mainWith $ do
      io "1k overwrites (one txn)" (insert env dbi) kv1
      io "1k gets (one txn)"       (readN env dbi)  ks1
      io "scan 1k entries"         (scanAll env)    dbi
  existsAfter <- doesDirectoryExist dir
  if existsAfter then removeDirectoryRecursive dir else pure ()

kvs :: Int -> [(BS.ByteString, BS.ByteString)]
kvs n = [ (mkKey i, mkVal i) | i <- [0 .. n - 1] ]

mkKey :: Int -> BS.ByteString
mkKey i = BSL.toStrict (BB.toLazyByteString (BB.string7 "k" <> BB.intDec i))

mkVal :: Int -> BS.ByteString
mkVal i = BSL.toStrict (BB.toLazyByteString (BB.string7 "v" <> BB.intDec i))

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

readN :: L.Env -> L.Dbi -> [BS.ByteString] -> IO ()
readN env dbi ks = L.withReadTxn env $ \txn ->
  forM_ ks $ \k -> do
    _ <- L.get txn dbi k
    pure ()

scanAll :: L.Env -> L.Dbi -> IO Int
scanAll env dbi = L.withReadTxn env $ \txn ->
  L.withCursor txn dbi $ \cur -> do
    mfirst <- L.cursorFirst cur
    let go !acc Nothing  = pure acc
        go !acc (Just _) = L.cursorNext cur >>= go (acc + 1)
    go (0 :: Int) mfirst