ppad-lmdb-0.1.0: bench/Weight.hs

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

module Main where

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
                   }
  L.withEnv dir envFlags $ \env -> do
    L.withWriteTxn env $ \txn -> do
      dbi <- L.openDbi txn Nothing True
      forM_ (kvs 1000) $ \(k, v) -> L.put txn dbi k v
    mainWith $ do
      io "1k put (one txn)"  (insert env) 1000
      io "1k get (one txn)"  (readN env)  1000
      io "scan 1k entries"   scanAll      env
  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 -> Int -> IO ()
insert env n = L.withWriteTxn env $ \txn -> do
  dbi <- L.openDbi txn Nothing False
  forM_ (kvs n) $ \(k, v) -> L.put txn dbi k v

readN :: L.Env -> Int -> IO ()
readN env n = L.withReadTxn env $ \txn -> do
  dbi <- L.openDbi txn Nothing False
  forM_ [0 .. n - 1] $ \i -> do
    _ <- L.get txn dbi (mkKey i)
    pure ()

scanAll :: L.Env -> IO Int
scanAll env = L.withReadTxn env $ \txn -> do
  dbi <- L.openDbi txn Nothing False
  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