memory-pool (empty) → 0.1.0.0
raw patch · 10 files changed
+911/−0 lines, 10 filesdep +QuickCheckdep +asyncdep +basesetup-changed
Dependencies added: QuickCheck, async, base, criterion, deepseq, memory-pool, primitive, pvar, random, reflection, tasty, tasty-hunit, tasty-quickcheck, unliftio
Files
- CHANGELOG.md +5/−0
- LICENSE +177/−0
- README.md +31/−0
- Setup.hs +4/−0
- bench/Bench.hs +57/−0
- memory-pool.cabal +87/−0
- src/System/Memory/Pool.hs +299/−0
- tests/Common.hs +7/−0
- tests/Main.hs +11/−0
- tests/Test/System/Memory/PoolTests.hs +233/−0
+ CHANGELOG.md view
@@ -0,0 +1,5 @@+# Changelog for `memory-pool`++## 0.1.0.0++* Initial release
+ LICENSE view
@@ -0,0 +1,177 @@++ Apache License+ Version 2.0, January 2004+ http://www.apache.org/licenses/++ TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION++ 1. Definitions.++ "License" shall mean the terms and conditions for use, reproduction,+ and distribution as defined by Sections 1 through 9 of this document.++ "Licensor" shall mean the copyright owner or entity authorized by+ the copyright owner that is granting the License.++ "Legal Entity" shall mean the union of the acting entity and all+ other entities that control, are controlled by, or are under common+ control with that entity. For the purposes of this definition,+ "control" means (i) the power, direct or indirect, to cause the+ direction or management of such entity, whether by contract or+ otherwise, or (ii) ownership of fifty percent (50%) or more of the+ outstanding shares, or (iii) beneficial ownership of such entity.++ "You" (or "Your") shall mean an individual or Legal Entity+ exercising permissions granted by this License.++ "Source" form shall mean the preferred form for making modifications,+ including but not limited to software source code, documentation+ source, and configuration files.++ "Object" form shall mean any form resulting from mechanical+ transformation or translation of a Source form, including but+ not limited to compiled object code, generated documentation,+ and conversions to other media types.++ "Work" shall mean the work of authorship, whether in Source or+ Object form, made available under the License, as indicated by a+ copyright notice that is included in or attached to the work+ (an example is provided in the Appendix below).++ "Derivative Works" shall mean any work, whether in Source or Object+ form, that is based on (or derived from) the Work and for which the+ editorial revisions, annotations, elaborations, or other modifications+ represent, as a whole, an original work of authorship. For the purposes+ of this License, Derivative Works shall not include works that remain+ separable from, or merely link (or bind by name) to the interfaces of,+ the Work and Derivative Works thereof.++ "Contribution" shall mean any work of authorship, including+ the original version of the Work and any modifications or additions+ to that Work or Derivative Works thereof, that is intentionally+ submitted to Licensor for inclusion in the Work by the copyright owner+ or by an individual or Legal Entity authorized to submit on behalf of+ the copyright owner. For the purposes of this definition, "submitted"+ means any form of electronic, verbal, or written communication sent+ to the Licensor or its representatives, including but not limited to+ communication on electronic mailing lists, source code control systems,+ and issue tracking systems that are managed by, or on behalf of, the+ Licensor for the purpose of discussing and improving the Work, but+ excluding communication that is conspicuously marked or otherwise+ designated in writing by the copyright owner as "Not a Contribution."++ "Contributor" shall mean Licensor and any individual or Legal Entity+ on behalf of whom a Contribution has been received by Licensor and+ subsequently incorporated within the Work.++ 2. Grant of Copyright License. Subject to the terms and conditions of+ this License, each Contributor hereby grants to You a perpetual,+ worldwide, non-exclusive, no-charge, royalty-free, irrevocable+ copyright license to reproduce, prepare Derivative Works of,+ publicly display, publicly perform, sublicense, and distribute the+ Work and such Derivative Works in Source or Object form.++ 3. Grant of Patent License. Subject to the terms and conditions of+ this License, each Contributor hereby grants to You a perpetual,+ worldwide, non-exclusive, no-charge, royalty-free, irrevocable+ (except as stated in this section) patent license to make, have made,+ use, offer to sell, sell, import, and otherwise transfer the Work,+ where such license applies only to those patent claims licensable+ by such Contributor that are necessarily infringed by their+ Contribution(s) alone or by combination of their Contribution(s)+ with the Work to which such Contribution(s) was submitted. If You+ institute patent litigation against any entity (including a+ cross-claim or counterclaim in a lawsuit) alleging that the Work+ or a Contribution incorporated within the Work constitutes direct+ or contributory patent infringement, then any patent licenses+ granted to You under this License for that Work shall terminate+ as of the date such litigation is filed.++ 4. Redistribution. You may reproduce and distribute copies of the+ Work or Derivative Works thereof in any medium, with or without+ modifications, and in Source or Object form, provided that You+ meet the following conditions:++ (a) You must give any other recipients of the Work or+ Derivative Works a copy of this License; and++ (b) You must cause any modified files to carry prominent notices+ stating that You changed the files; and++ (c) You must retain, in the Source form of any Derivative Works+ that You distribute, all copyright, patent, trademark, and+ attribution notices from the Source form of the Work,+ excluding those notices that do not pertain to any part of+ the Derivative Works; and++ (d) If the Work includes a "NOTICE" text file as part of its+ distribution, then any Derivative Works that You distribute must+ include a readable copy of the attribution notices contained+ within such NOTICE file, excluding those notices that do not+ pertain to any part of the Derivative Works, in at least one+ of the following places: within a NOTICE text file distributed+ as part of the Derivative Works; within the Source form or+ documentation, if provided along with the Derivative Works; or,+ within a display generated by the Derivative Works, if and+ wherever such third-party notices normally appear. The contents+ of the NOTICE file are for informational purposes only and+ do not modify the License. You may add Your own attribution+ notices within Derivative Works that You distribute, alongside+ or as an addendum to the NOTICE text from the Work, provided+ that such additional attribution notices cannot be construed+ as modifying the License.++ You may add Your own copyright statement to Your modifications and+ may provide additional or different license terms and conditions+ for use, reproduction, or distribution of Your modifications, or+ for any such Derivative Works as a whole, provided Your use,+ reproduction, and distribution of the Work otherwise complies with+ the conditions stated in this License.++ 5. Submission of Contributions. Unless You explicitly state otherwise,+ any Contribution intentionally submitted for inclusion in the Work+ by You to the Licensor shall be under the terms and conditions of+ this License, without any additional terms or conditions.+ Notwithstanding the above, nothing herein shall supersede or modify+ the terms of any separate license agreement you may have executed+ with Licensor regarding such Contributions.++ 6. Trademarks. This License does not grant permission to use the trade+ names, trademarks, service marks, or product names of the Licensor,+ except as required for reasonable and customary use in describing the+ origin of the Work and reproducing the content of the NOTICE file.++ 7. Disclaimer of Warranty. Unless required by applicable law or+ agreed to in writing, Licensor provides the Work (and each+ Contributor provides its Contributions) on an "AS IS" BASIS,+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or+ implied, including, without limitation, any warranties or conditions+ of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A+ PARTICULAR PURPOSE. You are solely responsible for determining the+ appropriateness of using or redistributing the Work and assume any+ risks associated with Your exercise of permissions under this License.++ 8. Limitation of Liability. In no event and under no legal theory,+ whether in tort (including negligence), contract, or otherwise,+ unless required by applicable law (such as deliberate and grossly+ negligent acts) or agreed to in writing, shall any Contributor be+ liable to You for damages, including any direct, indirect, special,+ incidental, or consequential damages of any character arising as a+ result of this License or out of the use or inability to use the+ Work (including but not limited to damages for loss of goodwill,+ work stoppage, computer failure or malfunction, or any and all+ other commercial damages or losses), even if such Contributor+ has been advised of the possibility of such damages.++ 9. Accepting Warranty or Additional Liability. While redistributing+ the Work or Derivative Works thereof, You may choose to offer,+ and charge a fee for, acceptance of support, warranty, indemnity,+ or other liability obligations and/or rights consistent with this+ License. However, in accepting such obligations, You may act only+ on Your own behalf and on Your sole responsibility, not on behalf+ of any other Contributor, and only if You agree to indemnify,+ defend, and hold each Contributor harmless for any liability+ incurred by, or claims asserted against, such Contributor by reason+ of your accepting any such warranty or additional liability.++ END OF TERMS AND CONDITIONS
+ README.md view
@@ -0,0 +1,31 @@+# `memory-pool`++Thread-safe lock-free memory pool manager.++## Overview++Memory management utility that allows allocating large chunks of memory at once while+using it at a fine grain smaller block level. There is a restriction on the blocks of+memory to having the same size in a single pool.++This package was designed for Cardano blockchain, but it is a general purpose Haskell+library. The particular use cases in Cardano is to allocate one large page of memlocked+memory using libsodium and treat it as many smaller regions for secure storage and cleanup+of private keys.++## Status++| Github Actions | Coveralls | Hackage | Nightly | LTS |+|:--------------:|:---------:|:-------:|:-------:|:---:|+| [![Build Status][GA-badge]][GA-link] | [![Coverage Status][Coveralls-badge]][Coveralls-link] | [![Hackage][Hackage-badge]][Hackage-link] | [![Nightly][Nightly-badge]][Nightly-link] | [![LTS][LTS-badge]][LTS-link]++[GA-badge]: https://github.com/input-output-hk/memory-pool/workflows/CI/badge.svg+[GA-link]: https://github.com/input-output-hk/memory-pool/actions+[Coveralls-badge]: https://coveralls.io/repos/github/input-output-hk/memory-pool/badge.svg?branch=master+[Coveralls-link]: https://coveralls.io/github/input-output-hk/memory-pool?branch=master+[Hackage-badge]: https://img.shields.io/hackage/v/memory-pool.svg+[Hackage-link]: https://hackage.haskell.org/package/memory-pool+[Nightly-badge]: https://www.stackage.org/package/memory-pool/badge/nightly+[Nightly-link]: https://www.stackage.org/nightly/package/memory-pool+[LTS-badge]: https://www.stackage.org/package/memory-pool/badge/lts+[LTS-link]: https://www.stackage.org/lts/package/memory-pool
+ Setup.hs view
@@ -0,0 +1,4 @@+import Distribution.Simple++main :: IO ()+main = defaultMain
+ bench/Bench.hs view
@@ -0,0 +1,57 @@+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE TypeApplications #-}+{-# OPTIONS_GHC -Wno-orphans #-}++module Main where++import Control.DeepSeq+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 System.Memory.Pool+import UnliftIO.Async (pooledReplicateConcurrently)++instance NFData (Pool n s) where+ rnf !_ = ()++instance NFData (ForeignPtr a) where+ rnf !_ = ()++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+ ptr <- mallocBytes n+ newForeignPtr finalizerFree ptr++main :: IO ()+main = do+ let n = 10240+ blockSize = 32+ defaultMain+ [ 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))+ ]+ ]
+ memory-pool.cabal view
@@ -0,0 +1,87 @@+name: memory-pool+version: 0.1.0.0+synopsis: Short description+description: Lock-free threadsafe pinned memory pool manager+homepage: https://github.com/input-output-hk/memory-pool+license: Apache-2.0+license-file: LICENSE+author: IOHK+maintainer: operations@iohk.io+category: Memory+build-type: Simple+extra-source-files: README.md+extra-doc-files: CHANGELOG.md+cabal-version: 1.18+tested-with: GHC == 8.0.2+ , GHC == 8.2.2+ , GHC == 8.4.4+ , GHC == 8.6.5+ , GHC == 8.8.4+ , GHC == 8.10.7+ , GHC == 9.0.2+ , GHC == 9.2.8+ , GHC == 9.4.8+ , GHC == 9.6.6+ , GHC == 9.8.2+ , GHC == 9.10.1++library+ hs-source-dirs: src+ exposed-modules: System.Memory.Pool++ other-modules:+ build-depends: base >= 4.9 && < 5+ , primitive+ , pvar >= 1.0++ default-language: Haskell2010+ ghc-options: -Wall+ -Wincomplete-record-updates+ -Wincomplete-uni-patterns+ -Wredundant-constraints++test-suite tests+ type: exitcode-stdio-1.0+ hs-source-dirs: tests+ main-is: Main.hs+ other-modules: Common+ , Test.System.Memory.PoolTests+ build-depends: base+ , async+ , memory-pool+ , primitive+ , pvar+ , random+ , reflection+ , tasty+ , tasty-quickcheck+ , tasty-hunit+ , QuickCheck++ default-language: Haskell2010+ ghc-options: -Wall+ -Wincomplete-record-updates+ -Wincomplete-uni-patterns+ -Wredundant-constraints+ -fno-warn-orphans+ -threaded+ -with-rtsopts=-N++benchmark bench+ type: exitcode-stdio-1.0+ hs-source-dirs: bench+ main-is: Bench.hs+ ghc-options: -Wall+ -threaded+ -O2+ -with-rtsopts=-N+ build-depends: base+ , criterion+ , deepseq+ , memory-pool+ , unliftio+ default-language: Haskell2010++source-repository head+ type: git+ location: https://github.com/input-output-hk/memory-pool
+ src/System/Memory/Pool.hs view
@@ -0,0 +1,299 @@+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE KindSignatures #-}+{-# LANGUAGE LambdaCase #-}+{-# LANGUAGE MagicHash #-}+{-# LANGUAGE NamedFieldPuns #-}+{-# LANGUAGE RankNTypes #-}+{-# LANGUAGE RecordWildCards #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE UnboxedTuples #-}++-- | The goal of this Memory Pool is to provide the ability to allocate big chunks of+-- memory that can fit many `Block`s. Some memory allocators out there have a fairly large+-- minimal size requirement, which would be wasteful if many chunks of small size (eg. 32+-- bytes) are needed at once. Memory pool will allocate one page at a time as more blocks+-- is needed.+--+-- Currently there is no functionality for releasing unused pages. So, once a page is+-- allocated, it will be re-used when more `Block`s is needed, but it will not be GCed+-- until the whole `Pool` is GCed.+module System.Memory.Pool (+ -- * Pool+ Pool,+ initPool,++ -- * Block+ Block (..),+ blockByteCount,+ grabNextBlock,++ -- * Helpers++ -- Exported for testing+ countPages,+ findNextZeroIndex,+) where++import Control.Applicative+import Control.Monad+import Data.Bits+import Data.Primitive.MutVar+import Data.Primitive.PVar+import Data.Primitive.PVar.Unsafe (atomicModifyIntArray#)+import Data.Primitive.PrimArray+import Foreign.ForeignPtr+import Foreign.Ptr+import GHC.Exts (fetchAndIntArray#)+import GHC.ForeignPtr (addForeignPtrConcFinalizer)+import GHC.IO+import GHC.Int+import GHC.ST+import GHC.TypeLits++-- | This is just a proxy type that carries information at the type level about the size+-- of the block in bytes supported by a particular instance of a `Pool`. Use+-- `blockByteCount` to get the byte size at the value level.+data Block (n :: Nat) = Block++-- | Number of bytes in a `Block`+blockByteCount :: KnownNat n => Block n -> Int+blockByteCount = fromIntegral . natVal++-- | Internal helper type that manages each individual page. This is essentially a mutable+-- linked list, which contains a memory buffer, a bit array that tracks which blocks in+-- the buffer are free and which ones are taken.+data Page n s = Page+ { pageMemory :: !(ForeignPtr (Block n))+ -- ^ Contiguous memory buffer that holds all the blocks in the page.+ , pageBitArray :: !(MutablePrimArray s Int)+ -- ^ We use an Int array, because there are no built-in atomic primops for Word.+ , pageFull :: !(PVar Int s)+ -- ^ This is a boolean flag which indicates when a page is full. It here as+ -- optimization only, because it allows us to skip iteration of the above bit+ -- array. It is an `Int` instead of a `Bool`, because GHC provides atomic primops for+ -- ByteArray, whcih is what `PVar` is based on.+ , pageNextPage :: !(MutVar s (Maybe (Page n s)))+ -- ^ Link to the next page. Last page when this IORef contains `Nothing`+ }++-- | Thread-safe lock-free memory pool for managing large memory pages that are made up of+-- many small `Block`s.+data Pool n s = Pool+ { poolFirstPage :: !(Page n s)+ -- ^ Initial page, which itself contains references to subsequent pages+ , poolPageInitializer :: !(ST s (Page n s))+ -- ^ Page initializing action+ , poolBlockFinalizer :: !(Ptr (Block n) -> IO ())+ -- ^ Finilizer that will be attached to each individual `ForeignPtr` of a reserved+ -- `Block`.+ }++-- | Useful function for testing. Check how many pages have been allocated thus far.+countPages :: Pool n s -> ST s Int+countPages pool = go 1 (poolFirstPage pool)+ where+ go n Page {pageNextPage} = do+ readMutVar pageNextPage >>= \case+ Nothing -> pure n+ Just nextPage -> go (n + 1) nextPage++ixBitSize :: Int+ixBitSize = finiteBitSize (0 :: Word)++-- | Initilizes the `Pool` that can be used for further allocation of @`ForeignPtr`+-- `Block` n@ with `grabNextBlock`.+initPool ::+ forall n s.+ KnownNat n =>+ -- | Number of groups per page. Must be a posititve number, otherwise error. One group+ -- contains as many blocks as the operating system has bits. A 64bit architecture will+ -- have 64 blocks per group. For example, if program is compiled on a 64 bit OS and you+ -- know ahead of time the maximum number of blocks that will be allocated through out+ -- the program, then the optimal value for this argument will @maxBlockNum/64@+ Int ->+ -- | Mempool page allocator. Some allocated pages might be immediately discarded,+ -- therefore number of pages utilized will not necessesarely match the number of times+ -- this action will be called.+ (forall a. Int -> ST s (ForeignPtr a)) ->+ -- | Finalizer to use for each block. It is an IO action because it will be executed by+ -- the Garbage Collector in a separate thread once the `Block` is no longer referenced.+ (Ptr (Block n) -> IO ()) ->+ ST s (Pool n s)+initPool groupsPerPage memAlloc blockFinalizer = do+ unless (groupsPerPage > 0) $+ error $+ "Groups per page should be a positive number, but got: "+ ++ show groupsPerPage+ let pageInit = do+ pageMemory <-+ memAlloc $ groupsPerPage * ixBitSize * blockByteCount (Block :: Block n)+ pageBitArray <- newPrimArray groupsPerPage+ setPrimArray pageBitArray 0 groupsPerPage 0+ pageFull <- newPVar 0+ pageNextPage <- newMutVar Nothing+ pure Page {..}+ firstPage <- pageInit+ pure+ Pool+ { poolFirstPage = firstPage+ , poolPageInitializer = pageInit+ , poolBlockFinalizer = blockFinalizer+ }++-- | Reserve a `ForeignPtr` of the `blockByteCount` size in the `Pool`. There is a default+-- finalizer attached to the `ForeignPtr` that will run `Block` pointer finalizer and+-- release that memory for re-use by other blocks allocated in the future. It is safe to+-- add more Haskell finalizers with `addForeignPtrConcFinalizer` if necessary.+grabNextBlock :: KnownNat n => Pool n s -> ST s (ForeignPtr (Block n))+grabNextBlock = grabNextPoolBlockWith grabNextPageForeignPtr+{-# INLINE grabNextBlock #-}++-- | This is a helper function that will allocate a `Page` if the current `Page` in the+-- `Pool` is full. Whenever there are still block slots are available then supplied+-- @grabNext@ function will be used to reserve the slot in that `Page`.+grabNextPoolBlockWith ::+ (Page n s -> (Ptr (Block n) -> IO ()) -> ST s (Maybe (ForeignPtr (Block n)))) ->+ Pool n s ->+ ST s (ForeignPtr (Block n))+grabNextPoolBlockWith grabNext pool = go (poolFirstPage pool)+ where+ go !page = do+ isPageFull <- atomicReadIntPVar (pageFull page)+ if intToBool isPageFull+ then+ readMutVar (pageNextPage page) >>= \case+ Nothing -> do+ newPage <- poolPageInitializer pool+ -- There is a slight chance of a race condition in that the next page could+ -- have been allocated and assigned to 'pageNextPage' by another thread+ -- since we last checked for it. This is not a problem since we can safely+ -- discard the page created in this thread and switch to the one that was+ -- assigned to 'pageNextPage'.+ mNextPage <-+ atomicModifyMutVar' (pageNextPage page) $ \mNextPage ->+ (mNextPage <|> Just newPage, mNextPage)+ case mNextPage of+ Nothing -> go newPage+ Just existingPage -> do+ -- Here we cleanup the newly allocated page in favor of the one that+ -- was potentially created by another thread. It is important to+ -- eagerly free up scarce resources.+ --+ -- This operation is idempotent and thread safe+ unsafeIOToST $ finalizeForeignPtr (pageMemory newPage)+ go existingPage+ Just nextPage -> go nextPage+ else+ grabNext page (poolBlockFinalizer pool) >>= \case+ Nothing -> go page+ Just ma -> pure ma+{-# INLINE grabNextPoolBlockWith #-}++intToBool :: Int -> Bool+intToBool 0 = False+intToBool _ = True++-- | This is a helper function that will attempt to find the next available slot for the+-- `Block` and create a `ForeignPtr` with the size of `Block` in the `Page`. In case when+-- `Page` is full it will return `Nothing`.+grabNextPageForeignPtr ::+ forall n s.+ KnownNat n =>+ -- | Page to grab the block from+ Page n s ->+ -- | Finalizer to run, once the `ForeignPtr` holding on to `Ptr` `Block` is no longer used+ (Ptr (Block n) -> IO ()) ->+ ST s (Maybe (ForeignPtr (Block n)))+grabNextPageForeignPtr page finalizer =+ grabNextPageWithAllocator page $ \blockPtr resetIndex -> do+ fp <- newForeignPtr_ blockPtr+ addForeignPtrConcFinalizer fp $ finalizer blockPtr >> resetIndex+ pure fp+{-# INLINE grabNextPageForeignPtr #-}++grabNextPageWithAllocator ::+ forall n s.+ KnownNat n =>+ Page n s ->+ (Ptr (Block n) -> IO () -> IO (ForeignPtr (Block n))) ->+ ST s (Maybe (ForeignPtr (Block n)))+grabNextPageWithAllocator Page {..} allocator = do+ setNextZero pageBitArray >>= \case+ -- There is a slight chance that some Blocks will be cleared before the pageFull is+ -- set to True. This is not a problem because that memory will be recovered as soon as+ -- any other Block in the Page is finalized+ --+ -- TODO: Potentially verify that first Int in pageBitArray has all bits set, in+ -- order to prevent the degenerate case of all Blocks beeing finalized right before+ -- the page is marked as full.+ Nothing -> Nothing <$ atomicWriteIntPVar pageFull 1+ Just ix ->+ fmap Just $+ unsafeIOToST $+ withForeignPtr pageMemory $ \pagePtr ->+ let !blockPtr =+ plusPtr pagePtr $ ix * blockByteCount (Block :: Block n)+ in allocator blockPtr $ do+ let !(!q, !r) = ix `quotRem` ixBitSize+ !pageBitMask = clearBit (complement 0) r+ touchForeignPtr pageMemory+ unsafeSTToIO $ atomicAndIntMutablePrimArray pageBitArray q pageBitMask+ unsafeSTToIO $ atomicWriteIntPVar pageFull 0+{-# INLINE grabNextPageWithAllocator #-}++-- | Atomically AND an element of the array+atomicAndIntMutablePrimArray :: MutablePrimArray s Int -> Int -> Int -> ST s ()+atomicAndIntMutablePrimArray (MutablePrimArray mba#) (I# i#) (I# m#) =+ ST $ \s# ->+ case fetchAndIntArray# mba# i# m# s# of+ (# s'#, _ #) -> (# s'#, () #)+{-# INLINE atomicAndIntMutablePrimArray #-}++-- | Atomically modify an element of the array+atomicModifyMutablePrimArray :: MutablePrimArray s Int -> Int -> (Int -> (Int, a)) -> ST s a+atomicModifyMutablePrimArray (MutablePrimArray mba#) (I# i#) f =+ ST $ atomicModifyIntArray# mba# i# (\x# -> case f (I# x#) of (I# y#, a) -> (# y#, a #))+{-# INLINE atomicModifyMutablePrimArray #-}++-- | Helper function that finds an index of the left-most bit that is not set.+findNextZeroIndex :: forall b. FiniteBits b => b -> Maybe Int+findNextZeroIndex b =+ let !i0 = countTrailingZeros b+ i1 = countTrailingZeros (complement b)+ maxBits = finiteBitSize (undefined :: b)+ in if i0 == 0+ then+ if i1 == maxBits+ then Nothing+ else Just i1+ else Just (i0 - 1)+{-# INLINE findNextZeroIndex #-}++-- | Finds an index of the next bit that is not set in the bit array and flips it+-- atomically. In case when all bits are set, then `Nothing` is returned. It is possible+-- that while search is ongoing bits that where checked get cleared. This is totally fine+-- for our implementation of mempool.+setNextZero :: MutablePrimArray s Int -> ST s (Maybe Int)+setNextZero ma = ifindAtomicMutablePrimArray ma f+ where+ f i !w =+ case findNextZeroIndex w of+ Nothing -> (w, Nothing)+ Just !bitIx -> (setBit w bitIx, Just (ixBitSize * i + bitIx))+{-# INLINE setNextZero #-}++ifindAtomicMutablePrimArray ::+ MutablePrimArray s Int ->+ (Int -> Int -> (Int, Maybe a)) ->+ ST s (Maybe a)+ifindAtomicMutablePrimArray ma f = do+ n <- getSizeofMutablePrimArray ma+ let go i+ | i >= n = pure Nothing+ | otherwise =+ atomicModifyMutablePrimArray ma i (f i) >>= \case+ Nothing -> go (i + 1)+ Just a -> pure $ Just a+ go 0+{-# INLINE ifindAtomicMutablePrimArray #-}
+ tests/Common.hs view
@@ -0,0 +1,7 @@+module Common (+ module X,+) where++import Test.Tasty as X+import Test.Tasty.HUnit as X+import Test.Tasty.QuickCheck as X
+ tests/Main.hs view
@@ -0,0 +1,11 @@+module Main where++import System.IO (BufferMode (LineBuffering), hSetBuffering, hSetEncoding, stdout, utf8)+import Test.System.Memory.PoolTests (poolTests)+import Test.Tasty++main :: IO ()+main = do+ hSetBuffering stdout LineBuffering+ hSetEncoding stdout utf8+ defaultMain poolTests
+ tests/Test/System/Memory/PoolTests.hs view
@@ -0,0 +1,233 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE ScopedTypeVariables #-}++module Test.System.Memory.PoolTests (poolTests) where++import Common+import Control.Concurrent (threadDelay)+import Control.Concurrent.Async+import Control.Concurrent.Chan+import Control.Monad+import Control.Monad.Primitive+import Data.Bits+import Data.Function+import Data.Primitive.PVar+import Data.Primitive.Ptr (setPtr)+import Data.Proxy+import Data.Reflection+import Data.Word+import Foreign.ForeignPtr+import Foreign.Ptr+import Foreign.Storable+import GHC.TypeLits+import System.Mem (performGC)+import System.Memory.Pool+import System.Random+import Test.QuickCheck.Monadic++poolTests :: TestTree+poolTests = do+ testGroup+ "Pool"+ [ testProperty "findNextZeroIndex" $ propFindNextZeroIndex+ , poolProps (Block :: Block 32)+ , poolProps (Block :: Block 64)+ , poolPropsArbSizeBlock+ ]++poolProps :: KnownNat n => Block n -> TestTree+poolProps block =+ testGroup+ ("Block " ++ show (blockByteCount block))+ [ testProperty "PoolGarbageCollected" $ propPoolGarbageCollected block+ , testProperty "PoolAllocateAndFinalize" $ propPoolAllocateAndFinalize block+ ]++poolPropsArbSizeBlock :: TestTree+poolPropsArbSizeBlock =+ testGroup+ "Arbitrary sized block"+ [ testProperty "PoolGarbageCollected" $ \(Positive n) ->+ reifyNat n (propPoolGarbageCollected . proxyToBlock)+ , testProperty "PoolAllocateAndFinalize" $ \(Positive n) ->+ reifyNat n (propPoolAllocateAndFinalize . proxyToBlock)+ ]+ where+ proxyToBlock :: Proxy n -> Block n+ proxyToBlock Proxy = Block++propFindNextZeroIndex :: Int -> Property+propFindNextZeroIndex w = monadicIO . run $+ case findNextZeroIndex w of+ Nothing -> w @?= complement 0+ Just ix -> do+ testBit w ix @?= False+ case findNextZeroIndex (setBit w ix) of+ Nothing -> setBit w ix @?= complement 0+ Just ix' -> do+ assertBool+ ("Expected found index to be different, but got same: " ++ show ix)+ (ix' /= ix)+ assertBool+ ("Expected the bit under index: " ++ show ix' ++ " to not be set")+ $ not (testBit w ix')++-- We allow one extra page be allocated due to concurrency false positives in block+-- reservations+checkNumPages :: Pool n RealWorld -> Int -> Int -> Assertion+checkNumPages pool n numBlocks = do+ let estimatedUpperBoundOfPages = 1 + max 1 (numBlocks `div` n `div` 64)+ numPages <- stToPrim $ countPages pool+ assertBool+ ( concat+ [ "Number of pages should not exceed the expected amount: "+ , show estimatedUpperBoundOfPages+ , " but allocated: "+ , show numPages+ ]+ )+ (numPages <= estimatedUpperBoundOfPages)++checkBlockBytes ::+ (KnownNat n, Storable a, Eq a, Show a) =>+ Block n ->+ a ->+ Ptr b ->+ Assertion+checkBlockBytes block byte ptr =+ let checkFillByte i =+ when (i >= 0) $ do+ byte' <- peekByteOff ptr i+ byte' @?= byte+ checkFillByte (i - 1)+ in checkFillByte (blockByteCount block - 1)++mallocPreFilled :: Word8 -> Int -> ST s (ForeignPtr b)+mallocPreFilled preFillByte bc = unsafeIOToPrim $ do+ mfp <- mallocForeignPtrBytes bc+ withForeignPtr mfp $ \ptr -> setPtr (castPtr ptr) bc preFillByte+ pure mfp++-- | @ensureAllGCedWith iterations delay expectedCount registerCounter@ waits+-- for all items to be GCed by triggering garbage collection @iterations@+-- times, once every @delay@ milliseconds. After @iterations@ attempts, if the+-- counter hook was not called exactly @expectedCount@ times, a test failure is+-- raised via 'assertFailure'. Garbage collection is tracked via finalizers+-- (see below).+ensureAllGCedWith ::+ -- | Number of GC attempts to make before failing+ Int ->+ -- | Delay between attempts, in milliseconds+ Int ->+ -- | Expected number of counter hook firings (in practice: individual+ -- garbage collections on 'ForeignPtr's, as per their finalizers).+ Int ->+ -- | Function for registering the counter hook. The argument to this+ -- function should be attached to each 'ForeignPtr' we're interested in+ -- as a finalizer.+ (IO () -> IO a) ->+ IO a+ensureAllGCedWith iterations delay expectedCount registerCounter = do+ countRef <- newPVar (0 :: Int)+ res <- registerCounter (void $ atomicAddIntPVar countRef 1)+ let go i = do+ performGC+ threadDelay (delay * 1000)+ n <- atomicReadIntPVar countRef+ unless (n == expectedCount) $ do+ if i <= 1+ then+ assertFailure $+ "Expected all "+ ++ show expectedCount+ ++ " pointers to be GCed in "+ ++ show (delay * iterations)+ ++ "ms, but "+ ++ show n+ ++ " where GCed instead"+ else go (i - 1)+ res <$ go iterations++-- | 'ensureAllGCedWith' with default values: 100 iterations, 10ms delay.+ensureAllGCed :: Int -> (IO () -> IO a) -> IO a+ensureAllGCed = ensureAllGCedWith 100 10++propPoolGarbageCollected ::+ forall n.+ KnownNat n =>+ Block n ->+ Positive Int ->+ Word16 ->+ Word8 ->+ Word8 ->+ Property+propPoolGarbageCollected block (Positive n) numBlocks16 preFillByte fillByte =+ monadicIO . run $ do+ let numBlocks = 1 + (fromIntegral numBlocks16 `div` 20) -- make it not too big+ (pool, ptrs) <-+ ensureAllGCed numBlocks $ \countOneBlockGCed -> do+ pool <-+ stToPrim $ initPool n (mallocPreFilled preFillByte) $ \ptr -> do+ setPtr (castPtr ptr) (blockByteCount block) fillByte+ countOneBlockGCed+ fmps :: [ForeignPtr (Block n)] <-+ replicateConcurrently numBlocks (stToPrim $ grabNextBlock pool)+ touch fmps+ -- Here we return just the pointers and let the GC collect the ForeignPtrs+ ptrs <-+ forM fmps $ \fma ->+ withForeignPtr fma $ \ptr -> do+ let bytePtr = castPtr ptr+ checkBlockBytes block preFillByte bytePtr+ setPtr bytePtr (blockByteCount block) fillByte+ pure bytePtr+ pure (pool, ptrs)+ forM_ ptrs (checkBlockBytes block fillByte)+ checkNumPages pool n numBlocks+ -- Ensure that memory to that the pointers are referencing to is still alive+ touch pool++propPoolAllocateAndFinalize ::+ forall n.+ KnownNat n =>+ Block n ->+ Positive Int ->+ Word16 ->+ Word8 ->+ Word8 ->+ Property+propPoolAllocateAndFinalize block (Positive n) numBlocks16 emptyByte fullByte =+ monadicIO . run $ do+ let numBlocks = 1 + (fromIntegral numBlocks16 `div` 20)+ pool <-+ ensureAllGCed numBlocks $ \countOneBlockGCed -> do+ chan <- newChan+ pool <-+ stToPrim $ initPool n (mallocPreFilled emptyByte) $ \(ptr :: Ptr (Block n)) -> do+ setPtr (castPtr ptr) (blockByteCount block) emptyByte+ countOneBlockGCed+ -- allocate and finalize blocks concurrently+ pool+ <$ concurrently_+ ( do+ replicateConcurrently_ numBlocks $ do+ fp <- stToPrim $ grabNextBlock pool+ withForeignPtr fp (checkBlockBytes block emptyByte)+ writeChan chan (Just fp)+ -- place Nothing to indicate that we are done allocating blocks+ writeChan chan Nothing+ )+ ( fix $ \loop -> do+ mfp <- readChan chan+ forM_ mfp $ \fp -> do+ withForeignPtr fp $ \ptr ->+ -- fill the newly allocated block+ setPtr (castPtr ptr) (blockByteCount block) fullByte+ -- manually finalize every other block and let the GC to pick the rest+ shouldFinalize <- randomIO+ when shouldFinalize $ finalizeForeignPtr fp+ loop+ )+ -- verify number of pages+ checkNumPages pool n numBlocks