diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,5 @@ +# Changelog for `memory-pool` + +## 0.1.0.0 + +* Initial release diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -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 diff --git a/README.md b/README.md new file mode 100644 --- /dev/null +++ b/README.md @@ -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 diff --git a/Setup.hs b/Setup.hs new file mode 100644 --- /dev/null +++ b/Setup.hs @@ -0,0 +1,4 @@ +import Distribution.Simple + +main :: IO () +main = defaultMain diff --git a/bench/Bench.hs b/bench/Bench.hs new file mode 100644 --- /dev/null +++ b/bench/Bench.hs @@ -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)) + ] + ] diff --git a/memory-pool.cabal b/memory-pool.cabal new file mode 100644 --- /dev/null +++ b/memory-pool.cabal @@ -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 diff --git a/src/System/Memory/Pool.hs b/src/System/Memory/Pool.hs new file mode 100644 --- /dev/null +++ b/src/System/Memory/Pool.hs @@ -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 #-} diff --git a/tests/Common.hs b/tests/Common.hs new file mode 100644 --- /dev/null +++ b/tests/Common.hs @@ -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 diff --git a/tests/Main.hs b/tests/Main.hs new file mode 100644 --- /dev/null +++ b/tests/Main.hs @@ -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 diff --git a/tests/Test/System/Memory/PoolTests.hs b/tests/Test/System/Memory/PoolTests.hs new file mode 100644 --- /dev/null +++ b/tests/Test/System/Memory/PoolTests.hs @@ -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