Skip to content

Instantly share code, notes, and snippets.

@voidlizard
Created February 27, 2024 13:49
Show Gist options
  • Save voidlizard/39c8efabade45992f399c22519caee08 to your computer and use it in GitHub Desktop.
Save voidlizard/39c8efabade45992f399c22519caee08 to your computer and use it in GitHub Desktop.
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 Database.RocksDB qualified as R
import Data.CritBit.Map.Lazy qualified as C
import Data.Word
import Data.Int
import Data.Default
import Codec.Compression.GZip as GZip
import System.Directory
-- data SkipListNode =
-- SkipListNode
-- {
-- }
runUpdates :: MonadUnliftIO m => Int -> Int -> Int -> (ByteString -> m ()) -> m ()
runUpdates total k n fn = do
tot <- newTVarIO total
let (a,b) = total `divMod` k
let tasks = a + b : replicate (k-1) a
g <- liftIO $ MWC.createSystemRandom
threads <- for tasks $ \t -> async do
replicateM_ t do
bs <- liftIO $ replicateM n (MWC.uniformM @Word8 g) <&> LBS.pack
fn bs
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 = 1_000_000
let total = 1_000_000
let keylen = 32
g <- MWC.createSystemRandom
liftIO $ removePathForcibly "lsm.wal"
liftIO $ removePathForcibly "lsm.data"
liftIO $ removePathForcibly "jopakita.db"
liftIO $ removePathForcibly "jopakita"
timeItNamed "just-none" do
runUpdates total threads keylen $ \_ -> do
pure ()
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
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 (Map.insert bs bs)
timeItNamed "just-tvar-buck-map-disk-wal" do
let buc = threads
tvm <- replicateM buc (newTVarIO mempty) <&> V.fromList
walq <- newTQueueIO
walp <- async do
tw <- newIORef ( 0 :: Int64 )
withBinaryFile "lsm.wal" AppendMode $ \ha -> do
fix \next -> do
it <- atomically $ readTQueue walq
case it of
Nothing -> pure ()
Just (k,v) -> do
let ss = k <> v
LBS.hPutStr ha ss
modifyIORef tw (+ LBS.length ss)
sz <- readIORef tw
-- sz <- atomicModifyIORef tw (\x -> (x + LBS.length ss, x + LBS.length ss))
when (sz `mod` (1024*1024) == 0) do
hFlush ha
next
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 (Map.insert bs bs)
atomically $ writeTQueue walq (Just (bs,bs))
withBinaryFile "lsm.data" AppendMode $ \ha -> do
joined <- mapM readTVarIO tvm <&> Map.toList . Map.unions
for_ joined $ \(k,v) -> do
LBS.hPutStr ha (k <> v)
atomically $ writeTQueue walq Nothing
wait walp
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)
timeItNamed "qwriter" do
-- tm <- newTVarIO mempty
q <- newTQueueIO
z <- newIORef mempty
runUpdates total threads keylen $ \bs -> do
let k = LBS.toStrict bs
atomicModifyIORef z (\x -> (k : x, ()))
-- atomically $ writeTQueue q k -- modifyTVar tm (C.insert k bs)
env <- newDBPipeEnv dbPipeOptsDef "jopakita.db"
withDB env do
ddl "create table kv (k text not null, v text not null, primary key (k))"
commitAll
timeItNamed "sqlite" do
withDB env $ transactional 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
R.withDB "jopakita" (def { R.createIfMissing = True }) $ \db -> do
timeItNamed "rocks" do
R.withSnapshot db $ \db1 -> do
runUpdates total threads keylen $ \bs -> do
R.put db1 (LBS.toStrict bs) (LBS.toStrict bs)
pure ()
pure ()
-- timeItNamed "just-tskiplist 8" do
-- tm <- TS.newIO @ByteString @ByteString
-- runUpdates total 8 keylen $ \bs -> do
-- atomically $ TS.insert bs bs tm
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment