concurrent-insert-bench.hs

module Main where

import System.TimeIt
import Data.ByteString.Lazy qualified as LBS
import Data.ByteString qualified as BS
import Data.ByteString.Lazy (ByteString)
import Crypto.Random (getRandomBytes)
import Control.Monad
import Data.Functor
import Data.Function
import UnliftIO
import Data.List (unfoldr)
import Data.Foldable
import Data.Traversable
import System.FilePath
import Data.Map qualified as Map
import Data.Map (Map)
import Data.Hashable
import Data.Vector qualified as V
import Data.Vector ((!))
import Safe
import Control.Concurrent.STM (flushTQueue)
-- import Data.IORef (atomicallyModifyIORef
import Prettyprinter
import Data.HashMap.Strict (HashMap)
import Data.HashMap.Strict qualified as HashMap
import System.Random.Stateful
-- .MWC qualified as MWC
import System.Random.MWC qualified as MWC

import StmContainers.Map qualified as StmMap
import Control.Concurrent.STM.TSkipList as TS
import Data.Trie as Trie

import DBPipe.SQLite as SQL

import Data.CritBit.Map.Lazy qualified as C

runUpdates :: Int -> Int -> Int -> (ByteString -> IO ()) -> IO ()
runUpdates total k n fn = do
  tot <- newTVarIO total

  threads <- replicateM k $ async do
               fix \next -> do
                 rest <- atomically do
                           n <- readTVar tot
                           modifyTVar tot pred
                           pure n

                 when (rest > 0) do
                   bs <- getRandomBytes n
                   fn (LBS.fromStrict bs)
                   next

  mapM_ wait threads

randomRead :: (StatefulGen g IO) => Int -> g -> Double -> (ByteString -> IO (Maybe ByteString)) -> IO (Maybe ByteString)
randomRead n g p rfn = do
  k <- getRandomBytes n <&> LBS.fromStrict
  pr <- MWC.uniformRM (0, 1.0 :: Double)  g

  if pr > (1 - p) then do
    rfn k
  else
    pure Nothing

main :: IO ()
main = do

  let threads = 8
  let total = 500_000
  let keylen = 32

  g <- MWC.createSystemRandom


  timeItNamed "just-count" do
    tm <- newIORef 0
    runUpdates total threads keylen $ \_ -> do
      atomicModifyIORef' tm (\x -> (succ x, ()))
    readIORef tm >>= print

  timeItNamed "just-tvar-map" do
    tvm <- newTVarIO mempty

    runUpdates total threads keylen $ \bs -> do
      atomically $ modifyTVar' tvm (Map.insert bs bs)

  timeItNamed "just-ioref-map" do
    tvm <- newIORef mempty
    runUpdates total threads keylen $ \bs -> do
      atomicModifyIORef' tvm (\x -> (Map.insert bs bs x, ()))

  timeItNamed "just-tvar-hashmap" do
    tvm <- newTVarIO mempty
    runUpdates total threads keylen $ \bs -> do
      atomically $ modifyTVar' tvm (HashMap.insert bs bs)

  timeItNamed "just-tvar-buck-map" do
    let buc = threads * 2
    tvm <- replicateM buc (newTVarIO mempty) <&> V.fromList
    runUpdates total threads keylen $ \bs -> do
      let i = maybe 0 fst (LBS.uncons bs) `mod` fromIntegral buc
      let t = tvm ! fromIntegral i
      atomically $ modifyTVar' t (HashMap.insert bs bs)

  timeItNamed "just-stm-map" do
    tm <- StmMap.newIO @ByteString @ByteString
    runUpdates total 8 keylen $ \bs -> do
      atomically $ StmMap.insert bs bs tm

  timeItNamed "bs-trie" do
    tm <- newTVarIO Trie.empty
    runUpdates total threads keylen $ \bs -> do
      let k = LBS.toStrict bs
      atomically $ modifyTVar tm (Trie.insert k bs)

  timeItNamed "critbit" do
    tm <- newTVarIO C.empty
    runUpdates total threads keylen $ \bs -> do
      let k = LBS.toStrict bs
      atomically $ modifyTVar tm (C.insert k bs)

  env <- newDBPipeEnv dbPipeOptsDef ":memory:"

  withDB env do
    ddl "create table kv (k text not null, v text not null, primary key (k))"
    commitAll

  timeItNamed "sqlite" do
    runUpdates total threads keylen $ \bs -> do
      withDB env $ do
        SQL.insert "insert into kv (k,v) values (?,?) on conflict (k) do nothing" (bs,bs)
    withDB env commitAll

  pure ()

  -- timeItNamed "just-tskiplist 8" do
  --   tm <- TS.newIO @ByteString @ByteString
  --   runUpdates total 8 keylen $ \bs -> do
  --     atomically $ TS.insert bs bs tm