Fix cardano-mempool benchmarks

IntersectMBO · Jul 26, 2023 · 472804f · 472804f
1 parent f3592f6
commit 472804f
Showing 1 changed file with 30 additions and 24 deletions.
diff --git a/cardano-mempool/bench/Bench.hs b/cardano-mempool/bench/Bench.hs
@@ -2,25 +2,28 @@
 {-# LANGUAGE DataKinds #-}
 {-# LANGUAGE TypeApplications #-}
 {-# OPTIONS_GHC -Wno-orphans #-}
+
 module Main where
 
-import Foreign.Marshal.Alloc
-import GHC.TypeLits
-import Criterion.Main
 import Cardano.Memory.Pool
-import Foreign.ForeignPtr
 import Control.DeepSeq
-import UnliftIO.Async (pooledReplicateConcurrently)
 import Control.Monad
+import Control.Monad.ST (RealWorld, stToIO)
+import Criterion.Main
+import Foreign.ForeignPtr
+import Foreign.Marshal.Alloc
+import GHC.IO (ioToST)
+import GHC.TypeLits
+import UnliftIO.Async (pooledReplicateConcurrently)
 
-instance NFData (Pool n) where
+instance NFData (Pool n s) where
   rnf !_ = ()
 
 instance NFData (ForeignPtr a) where
   rnf !_ = ()
 
-initHaskellPool :: KnownNat n => Int -> IO (Pool n)
-initHaskellPool n = initPool n mallocForeignPtrBytes (const (pure ()))
+initHaskellPool :: (KnownNat n) => Int -> IO (Pool n RealWorld)
+initHaskellPool n = stToIO $ initPool n (ioToST . mallocForeignPtrBytes) (const (pure ()))
 
 cmallocForeignPtr :: Int -> IO (ForeignPtr a)
 cmallocForeignPtr n = do
@@ -32,20 +35,23 @@ main = do
   let n = 10240
       blockSize = 32
   defaultMain
-    [ bgroup "Sequential"
-      [ env (initHaskellPool @32 (n `div` 64)) $ \pool ->
-          bench "ForeignPtr (Pool)" $ nfIO (replicateM n (grabNextBlock pool))
-      , bench "ForeignPtr (ByteArray)" $
-          nfIO (replicateM n (mallocForeignPtrBytes blockSize))
-      , bench "ForeignPtr (malloc)" $
-          nfIO (replicateM n (cmallocForeignPtr blockSize))
-      ]
-    , bgroup "Concurrent"
-      [ env (initHaskellPool @32 (n `div` 64)) $ \pool ->
-          bench "ForeignPtr (Pool)" $ nfIO (pooledReplicateConcurrently n (grabNextBlock pool))
-      , bench "ForeignPtr (ByteArray)" $
-          nfIO (pooledReplicateConcurrently n (mallocForeignPtrBytes blockSize))
-      , bench "ForeignPtr (malloc)" $
-          nfIO (pooledReplicateConcurrently n (cmallocForeignPtr blockSize))
-      ]
+    [ bgroup
+        "Sequential"
+        [ env (initHaskellPool @32 (n `div` 64)) $ \pool ->
+            bench "ForeignPtr (Pool)" $ nfIO $ replicateM n (stToIO (grabNextBlock pool))
+        , bench "ForeignPtr (ByteArray)" $
+            nfIO (replicateM n (mallocForeignPtrBytes blockSize))
+        , bench "ForeignPtr (malloc)" $
+            nfIO (replicateM n (cmallocForeignPtr blockSize))
+        ]
+    , bgroup
+        "Concurrent"
+        [ env (initHaskellPool @32 (n `div` 64)) $ \pool ->
+            bench "ForeignPtr (Pool)" $
+              nfIO (pooledReplicateConcurrently n (stToIO (grabNextBlock pool)))
+        , bench "ForeignPtr (ByteArray)" $
+            nfIO (pooledReplicateConcurrently n (mallocForeignPtrBytes blockSize))
+        , bench "ForeignPtr (malloc)" $
+            nfIO (pooledReplicateConcurrently n (cmallocForeignPtr blockSize))
+        ]
     ]