mempool-bench: move more pre-calculation into withResource

Author: nfrisby

ouroboros-consensus/bench/mempool-bench/Main.hs

@@ -21,6 +21,7 @@ import           Data.Set ()
 import qualified Data.Text as Text
 import qualified Data.Text.Read as Text.Read
 import           Main.Utf8 (withStdTerminalHandles)
+import           Ouroboros.Consensus.Ledger.SupportsMempool (ByteSize)
 import qualified Ouroboros.Consensus.Mempool.Capacity as Mempool
 import           System.Exit (die, exitFailure)
 import qualified Test.Consensus.Mempool.Mocked as Mocked
@@ -54,31 +55,33 @@ main = withStdTerminalHandles $ do
           where
             benchAddNTxs n =
                 withResource
-                  (pure $!! mkNTryAddTxs n)
+                  (pure $!!
+                     let cmds = mkNTryAddTxs n
+                     in (cmds, foldMap TestBlock.txSize $ getCmdsTxs cmds))
                   (\_ -> pure ())
-                  (\getTxs -> do
+                  (\getCmds -> do
                       bgroup (show n <> " transactions") [
                           bench "setup mempool" $ whnfIO $ do
-                            txs <- getTxs
-                            openMempoolWithCapacityFor txs
+                            (_cmds, capacity) <- getCmds
+                            openMempoolWithCapacity capacity
                         , bench "setup mempool + benchmark" $ whnfIO $ do
-                            txs <- getTxs
-                            mempool <- openMempoolWithCapacityFor txs
-                            run mempool txs
+                            (cmds, capacity) <- getCmds
+                            mempool <- openMempoolWithCapacity capacity
+                            run mempool cmds
                         , testCase "test" $ do
-                            txs <- getTxs
-                            mempool <- openMempoolWithCapacityFor txs
-                            testAddTxs mempool txs
-                        , testCase "txs length" $ do
-                            txs <- getTxs
-                            length txs @?= n
+                            (cmds, capacity) <- getCmds
+                            mempool <- openMempoolWithCapacity capacity
+                            testAddCmds mempool cmds
+                        , testCase "cmds length" $ do
+                            (cmds, _capacity) <- getCmds
+                            length cmds @?= n
                         ]
                   )
               where
-                testAddTxs mempool txs = do
-                    run mempool txs
+                testAddCmds mempool cmds = do
+                    run mempool cmds
                     mempoolTxs <- Mocked.getTxs mempool
-                    mempoolTxs @?= getCmdsTxs txs
+                    mempoolTxs @?= getCmdsTxs cmds
 
     parseBenchmarkResults csvFilePath = do
         csvData <- BL.readFile csvFilePath
@@ -129,17 +132,14 @@ main = withStdTerminalHandles $ do
   Adding TestBlock transactions to a mempool
 -------------------------------------------------------------------------------}
 
-openMempoolWithCapacityFor :: [MempoolCmd TestBlock] ->  IO (MockedMempool IO TestBlock)
-openMempoolWithCapacityFor cmds =
-    Mocked.openMockedMempool capacityRequiredByCmds
+openMempoolWithCapacity :: ByteSize ->  IO (MockedMempool IO TestBlock)
+openMempoolWithCapacity capacity =
+    Mocked.openMockedMempool (Mempool.mkCapacityBytesOverride  capacity)
                              Tracer.nullTracer
                              Mocked.MempoolAndModelParams {
                                  Mocked.immpInitialState = TestBlock.initialLedgerState
                                , Mocked.immpLedgerConfig = TestBlock.sampleLedgerConfig
                              }
-  where
-    capacityRequiredByCmds = Mempool.mkCapacityBytesOverride totalTxsSize
-      where totalTxsSize = foldMap TestBlock.txSize $ getCmdsTxs cmds
 
 mkNTryAddTxs :: Int -> [MempoolCmd TestBlock.TestBlock]
 mkNTryAddTxs 0 = []