{-# OPTIONS_HADDOCK not-home #-}
module Data.Pool.Internal where
import Control.Concurrent
import Control.Exception
import Control.Monad
import Data.Hashable (hash)
import Data.IORef
import Data.Primitive.SmallArray
import GHC.Clock
import qualified Data.List as L
data Pool a = Pool
{ forall a. Pool a -> PoolConfig a
poolConfig :: !(PoolConfig a)
, forall a. Pool a -> SmallArray (LocalPool a)
localPools :: !(SmallArray (LocalPool a))
, forall a. Pool a -> IORef ()
reaperRef :: !(IORef ())
}
data LocalPool a = LocalPool
{ forall a. LocalPool a -> Int
stripeId :: !Int
, forall a. LocalPool a -> MVar (Stripe a)
stripeVar :: !(MVar (Stripe a))
, forall a. LocalPool a -> IORef ()
cleanerRef :: !(IORef ())
}
data Stripe a = Stripe
{ forall a. Stripe a -> Int
available :: !Int
, forall a. Stripe a -> [Entry a]
cache :: ![Entry a]
, forall a. Stripe a -> Queue a
queue :: !(Queue a)
, forall a. Stripe a -> Queue a
queueR :: !(Queue a)
}
data Entry a = Entry
{ forall a. Entry a -> a
entry :: a
, forall a. Entry a -> Double
lastUsed :: !Double
}
data Queue a = Queue !(MVar (Maybe a)) (Queue a) | Empty
data PoolConfig a = PoolConfig
{ forall a. PoolConfig a -> IO a
createResource :: !(IO a)
, forall a. PoolConfig a -> a -> IO ()
freeResource :: !(a -> IO ())
, forall a. PoolConfig a -> Double
poolCacheTTL :: !Double
, forall a. PoolConfig a -> Int
poolMaxResources :: !Int
, forall a. PoolConfig a -> Maybe Int
poolNumStripes :: !(Maybe Int)
}
defaultPoolConfig
:: IO a
-> (a -> IO ())
-> Double
-> Int
-> PoolConfig a
defaultPoolConfig :: forall a. IO a -> (a -> IO ()) -> Double -> Int -> PoolConfig a
defaultPoolConfig IO a
create a -> IO ()
free Double
cacheTTL Int
maxResources = PoolConfig
{ createResource :: IO a
createResource = IO a
create
, freeResource :: a -> IO ()
freeResource = a -> IO ()
free
, poolCacheTTL :: Double
poolCacheTTL = Double
cacheTTL
, poolMaxResources :: Int
poolMaxResources = Int
maxResources
, poolNumStripes :: Maybe Int
poolNumStripes = forall a. Maybe a
Nothing
}
setNumStripes :: Maybe Int -> PoolConfig a -> PoolConfig a
setNumStripes :: forall a. Maybe Int -> PoolConfig a -> PoolConfig a
setNumStripes Maybe Int
numStripes PoolConfig a
pc = PoolConfig a
pc { poolNumStripes :: Maybe Int
poolNumStripes = Maybe Int
numStripes }
newPool :: PoolConfig a -> IO (Pool a)
newPool :: forall a. PoolConfig a -> IO (Pool a)
newPool PoolConfig a
pc = do
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (forall a. PoolConfig a -> Double
poolCacheTTL PoolConfig a
pc forall a. Ord a => a -> a -> Bool
< Double
0.5) forall a b. (a -> b) -> a -> b
$ do
forall a. HasCallStack => [Char] -> a
error [Char]
"poolCacheTTL must be at least 0.5"
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (forall a. PoolConfig a -> Int
poolMaxResources PoolConfig a
pc forall a. Ord a => a -> a -> Bool
< Int
1) forall a b. (a -> b) -> a -> b
$ do
forall a. HasCallStack => [Char] -> a
error [Char]
"poolMaxResources must be at least 1"
Int
numStripes <- forall b a. b -> (a -> b) -> Maybe a -> b
maybe IO Int
getNumCapabilities forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall a. PoolConfig a -> Maybe Int
poolNumStripes PoolConfig a
pc)
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
numStripes forall a. Ord a => a -> a -> Bool
< Int
1) forall a b. (a -> b) -> a -> b
$ do
forall a. HasCallStack => [Char] -> a
error [Char]
"numStripes must be at least 1"
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (forall a. PoolConfig a -> Int
poolMaxResources PoolConfig a
pc forall a. Ord a => a -> a -> Bool
< Int
numStripes) forall a b. (a -> b) -> a -> b
$ do
forall a. HasCallStack => [Char] -> a
error [Char]
"poolMaxResources must not be smaller than numStripes"
SmallArray (LocalPool a)
pools <- forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall a. Int -> [a] -> SmallArray a
smallArrayFromListN Int
numStripes) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [Int
1..Int
numStripes] forall a b. (a -> b) -> a -> b
$ \Int
n -> do
IORef ()
ref <- forall a. a -> IO (IORef a)
newIORef ()
MVar (Stripe a)
stripe <- forall a. a -> IO (MVar a)
newMVar Stripe
{ available :: Int
available = forall a. PoolConfig a -> Int
poolMaxResources PoolConfig a
pc Int -> Int -> Int
`quotCeil` Int
numStripes
, cache :: [Entry a]
cache = []
, queue :: Queue a
queue = forall a. Queue a
Empty
, queueR :: Queue a
queueR = forall a. Queue a
Empty
}
forall (f :: * -> *) a. Functor f => f a -> f ()
void forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. IORef a -> IO () -> IO (Weak (IORef a))
mkWeakIORef IORef ()
ref forall a b. (a -> b) -> a -> b
$ forall a.
(Entry a -> Bool) -> (a -> IO ()) -> MVar (Stripe a) -> IO ()
cleanStripe (forall a b. a -> b -> a
const Bool
True) (forall a. PoolConfig a -> a -> IO ()
freeResource PoolConfig a
pc) MVar (Stripe a)
stripe
forall (f :: * -> *) a. Applicative f => a -> f a
pure LocalPool { stripeId :: Int
stripeId = Int
n
, stripeVar :: MVar (Stripe a)
stripeVar = MVar (Stripe a)
stripe
, cleanerRef :: IORef ()
cleanerRef = IORef ()
ref
}
forall a. IO a -> IO a
mask_ forall a b. (a -> b) -> a -> b
$ do
IORef ()
ref <- forall a. a -> IO (IORef a)
newIORef ()
ThreadId
collectorA <- ((forall a. IO a -> IO a) -> IO ()) -> IO ThreadId
forkIOWithUnmask forall a b. (a -> b) -> a -> b
$ \forall a. IO a -> IO a
unmask -> forall a. IO a -> IO a
unmask forall a b. (a -> b) -> a -> b
$ forall {t :: * -> *} {b}. Foldable t => t (LocalPool a) -> IO b
collector SmallArray (LocalPool a)
pools
forall (f :: * -> *) a. Functor f => f a -> f ()
void forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. IORef a -> IO () -> IO (Weak (IORef a))
mkWeakIORef IORef ()
ref forall a b. (a -> b) -> a -> b
$ do
ThreadId -> IO ()
killThread ThreadId
collectorA
forall (f :: * -> *) a. Applicative f => a -> f a
pure Pool { poolConfig :: PoolConfig a
poolConfig = PoolConfig a
pc
, localPools :: SmallArray (LocalPool a)
localPools = SmallArray (LocalPool a)
pools
, reaperRef :: IORef ()
reaperRef = IORef ()
ref
}
where
quotCeil :: Int -> Int -> Int
quotCeil :: Int -> Int -> Int
quotCeil Int
x Int
y =
let (Int
z, Int
r) = Int
x forall a. Integral a => a -> a -> (a, a)
`quotRem` Int
y in if Int
r forall a. Eq a => a -> a -> Bool
== Int
0 then Int
z else Int
z forall a. Num a => a -> a -> a
+ Int
1
collector :: t (LocalPool a) -> IO b
collector t (LocalPool a)
pools = forall (f :: * -> *) a b. Applicative f => f a -> f b
forever forall a b. (a -> b) -> a -> b
$ do
Int -> IO ()
threadDelay Int
1000000
Double
now <- IO Double
getMonotonicTime
let isStale :: Entry a -> Bool
isStale Entry a
e = Double
now forall a. Num a => a -> a -> a
- forall a. Entry a -> Double
lastUsed Entry a
e forall a. Ord a => a -> a -> Bool
> forall a. PoolConfig a -> Double
poolCacheTTL PoolConfig a
pc
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (forall a.
(Entry a -> Bool) -> (a -> IO ()) -> MVar (Stripe a) -> IO ()
cleanStripe forall {a}. Entry a -> Bool
isStale (forall a. PoolConfig a -> a -> IO ()
freeResource PoolConfig a
pc) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. LocalPool a -> MVar (Stripe a)
stripeVar) t (LocalPool a)
pools
destroyResource :: Pool a -> LocalPool a -> a -> IO ()
destroyResource :: forall a. Pool a -> LocalPool a -> a -> IO ()
destroyResource Pool a
pool LocalPool a
lp a
a = do
forall a. IO a -> IO a
uninterruptibleMask_ forall a b. (a -> b) -> a -> b
$ do
Stripe a
stripe <- forall a. MVar a -> IO a
takeMVar (forall a. LocalPool a -> MVar (Stripe a)
stripeVar LocalPool a
lp)
Stripe a
newStripe <- forall a. Stripe a -> Maybe a -> IO (Stripe a)
signal Stripe a
stripe forall a. Maybe a
Nothing
forall a. MVar a -> a -> IO ()
putMVar (forall a. LocalPool a -> MVar (Stripe a)
stripeVar LocalPool a
lp) Stripe a
newStripe
forall (f :: * -> *) a. Functor f => f a -> f ()
void forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall e a. Exception e => IO a -> IO (Either e a)
try @SomeException forall a b. (a -> b) -> a -> b
$ forall a. PoolConfig a -> a -> IO ()
freeResource (forall a. Pool a -> PoolConfig a
poolConfig Pool a
pool) a
a
putResource :: LocalPool a -> a -> IO ()
putResource :: forall a. LocalPool a -> a -> IO ()
putResource LocalPool a
lp a
a = do
forall a. IO a -> IO a
uninterruptibleMask_ forall a b. (a -> b) -> a -> b
$ do
Stripe a
stripe <- forall a. MVar a -> IO a
takeMVar (forall a. LocalPool a -> MVar (Stripe a)
stripeVar LocalPool a
lp)
Stripe a
newStripe <- forall a. Stripe a -> Maybe a -> IO (Stripe a)
signal Stripe a
stripe (forall a. a -> Maybe a
Just a
a)
forall a. MVar a -> a -> IO ()
putMVar (forall a. LocalPool a -> MVar (Stripe a)
stripeVar LocalPool a
lp) Stripe a
newStripe
destroyAllResources :: Pool a -> IO ()
destroyAllResources :: forall a. Pool a -> IO ()
destroyAllResources Pool a
pool = forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ (forall a. Pool a -> SmallArray (LocalPool a)
localPools Pool a
pool) forall a b. (a -> b) -> a -> b
$ \LocalPool a
lp -> do
forall a.
(Entry a -> Bool) -> (a -> IO ()) -> MVar (Stripe a) -> IO ()
cleanStripe (forall a b. a -> b -> a
const Bool
True) (forall a. PoolConfig a -> a -> IO ()
freeResource (forall a. Pool a -> PoolConfig a
poolConfig Pool a
pool)) (forall a. LocalPool a -> MVar (Stripe a)
stripeVar LocalPool a
lp)
getLocalPool :: SmallArray (LocalPool a) -> IO (LocalPool a)
getLocalPool :: forall a. SmallArray (LocalPool a) -> IO (LocalPool a)
getLocalPool SmallArray (LocalPool a)
pools = do
Int
sid <- if Int
stripes forall a. Eq a => a -> a -> Bool
== Int
1
then forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
0
else do
Int
capabilities <- IO Int
getNumCapabilities
if Int
stripes forall a. Ord a => a -> a -> Bool
< Int
capabilities Bool -> Bool -> Bool
&& Int
capabilities forall a. Integral a => a -> a -> a
`rem` Int
stripes forall a. Eq a => a -> a -> Bool
/= Int
0
then forall a. Hashable a => a -> Int
hash forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IO ThreadId
myThreadId
else forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a b. (a, b) -> a
fst forall b c a. (b -> c) -> (a -> b) -> a -> c
. ThreadId -> IO (Int, Bool)
threadCapability forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< IO ThreadId
myThreadId
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ SmallArray (LocalPool a)
pools forall a. SmallArray a -> Int -> a
`indexSmallArray` (Int
sid forall a. Integral a => a -> a -> a
`rem` Int
stripes)
where
stripes :: Int
stripes = forall a. SmallArray a -> Int
sizeofSmallArray SmallArray (LocalPool a)
pools
waitForResource :: MVar (Stripe a) -> MVar (Maybe a) -> IO (Maybe a)
waitForResource :: forall a. MVar (Stripe a) -> MVar (Maybe a) -> IO (Maybe a)
waitForResource MVar (Stripe a)
mstripe MVar (Maybe a)
q = forall a. MVar a -> IO a
takeMVar MVar (Maybe a)
q forall a b. IO a -> IO b -> IO a
`onException` IO ()
cleanup
where
cleanup :: IO ()
cleanup = forall a. IO a -> IO a
uninterruptibleMask_ forall a b. (a -> b) -> a -> b
$ do
Stripe a
stripe <- forall a. MVar a -> IO a
takeMVar MVar (Stripe a)
mstripe
Stripe a
newStripe <- forall a. MVar a -> IO (Maybe a)
tryTakeMVar MVar (Maybe a)
q forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
Just Maybe a
ma -> do
forall a. Stripe a -> Maybe a -> IO (Stripe a)
signal Stripe a
stripe Maybe a
ma
Maybe (Maybe a)
Nothing -> do
forall a. MVar a -> a -> IO ()
putMVar MVar (Maybe a)
q forall a b. (a -> b) -> a -> b
$ forall a. HasCallStack => [Char] -> a
error [Char]
"unreachable"
forall (f :: * -> *) a. Applicative f => a -> f a
pure Stripe a
stripe
forall a. MVar a -> a -> IO ()
putMVar MVar (Stripe a)
mstripe Stripe a
newStripe
restoreSize :: MVar (Stripe a) -> IO ()
restoreSize :: forall a. MVar (Stripe a) -> IO ()
restoreSize MVar (Stripe a)
mstripe = forall a. IO a -> IO a
uninterruptibleMask_ forall a b. (a -> b) -> a -> b
$ do
Stripe a
stripe <- forall a. MVar a -> IO a
takeMVar MVar (Stripe a)
mstripe
forall a. MVar a -> a -> IO ()
putMVar MVar (Stripe a)
mstripe forall a b. (a -> b) -> a -> b
$! Stripe a
stripe { available :: Int
available = forall a. Stripe a -> Int
available Stripe a
stripe forall a. Num a => a -> a -> a
+ Int
1 }
cleanStripe
:: (Entry a -> Bool)
-> (a -> IO ())
-> MVar (Stripe a)
-> IO ()
cleanStripe :: forall a.
(Entry a -> Bool) -> (a -> IO ()) -> MVar (Stripe a) -> IO ()
cleanStripe Entry a -> Bool
isStale a -> IO ()
free MVar (Stripe a)
mstripe = forall b. ((forall a. IO a -> IO a) -> IO b) -> IO b
mask forall a b. (a -> b) -> a -> b
$ \forall a. IO a -> IO a
unmask -> do
[a]
stale <- forall a b. MVar a -> (a -> IO (a, b)) -> IO b
modifyMVar MVar (Stripe a)
mstripe forall a b. (a -> b) -> a -> b
$ \Stripe a
stripe -> forall a. IO a -> IO a
unmask forall a b. (a -> b) -> a -> b
$ do
let ([Entry a]
stale, [Entry a]
fresh) = forall a. (a -> Bool) -> [a] -> ([a], [a])
L.partition Entry a -> Bool
isStale (forall a. Stripe a -> [Entry a]
cache Stripe a
stripe)
newStripe :: Stripe a
newStripe = Stripe a
stripe { cache :: [Entry a]
cache = [Entry a]
fresh }
Stripe a
newStripe seq :: forall a b. a -> b -> b
`seq` forall (f :: * -> *) a. Applicative f => a -> f a
pure (Stripe a
newStripe, forall a b. (a -> b) -> [a] -> [b]
map forall a. Entry a -> a
entry [Entry a]
stale)
forall a. IO a -> IO a
uninterruptibleMask_ forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [a]
stale forall a b. (a -> b) -> a -> b
$ forall e a. Exception e => IO a -> IO (Either e a)
try @SomeException forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> IO ()
free
signal :: Stripe a -> Maybe a -> IO (Stripe a)
signal :: forall a. Stripe a -> Maybe a -> IO (Stripe a)
signal Stripe a
stripe Maybe a
ma = if forall a. Stripe a -> Int
available Stripe a
stripe forall a. Eq a => a -> a -> Bool
== Int
0
then Queue a -> Queue a -> IO (Stripe a)
loop (forall a. Stripe a -> Queue a
queue Stripe a
stripe) (forall a. Stripe a -> Queue a
queueR Stripe a
stripe)
else do
[Entry a]
newCache <- case Maybe a
ma of
Just a
a -> do
Double
now <- IO Double
getMonotonicTime
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall a. a -> Double -> Entry a
Entry a
a Double
now forall a. a -> [a] -> [a]
: forall a. Stripe a -> [Entry a]
cache Stripe a
stripe
Maybe a
Nothing -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall a. Stripe a -> [Entry a]
cache Stripe a
stripe
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$! Stripe a
stripe { available :: Int
available = forall a. Stripe a -> Int
available Stripe a
stripe forall a. Num a => a -> a -> a
+ Int
1
, cache :: [Entry a]
cache = [Entry a]
newCache
}
where
loop :: Queue a -> Queue a -> IO (Stripe a)
loop Queue a
Empty Queue a
Empty = do
[Entry a]
newCache <- case Maybe a
ma of
Just a
a -> do
Double
now <- IO Double
getMonotonicTime
forall (f :: * -> *) a. Applicative f => a -> f a
pure [forall a. a -> Double -> Entry a
Entry a
a Double
now]
Maybe a
Nothing -> forall (f :: * -> *) a. Applicative f => a -> f a
pure []
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$! Stripe { available :: Int
available = Int
1
, cache :: [Entry a]
cache = [Entry a]
newCache
, queue :: Queue a
queue = forall a. Queue a
Empty
, queueR :: Queue a
queueR = forall a. Queue a
Empty
}
loop Queue a
Empty Queue a
qR = Queue a -> Queue a -> IO (Stripe a)
loop (forall a. Queue a -> Queue a
reverseQueue Queue a
qR) forall a. Queue a
Empty
loop (Queue MVar (Maybe a)
q Queue a
qs) Queue a
qR = forall a. MVar a -> a -> IO Bool
tryPutMVar MVar (Maybe a)
q Maybe a
ma forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
Bool
False -> Queue a -> Queue a -> IO (Stripe a)
loop Queue a
qs Queue a
qR
Bool
True -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$! Stripe a
stripe { available :: Int
available = Int
0
, queue :: Queue a
queue = Queue a
qs
, queueR :: Queue a
queueR = Queue a
qR
}
reverseQueue :: Queue a -> Queue a
reverseQueue :: forall a. Queue a -> Queue a
reverseQueue = forall {a}. Queue a -> Queue a -> Queue a
go forall a. Queue a
Empty
where
go :: Queue a -> Queue a -> Queue a
go Queue a
acc = \case
Queue a
Empty -> Queue a
acc
Queue MVar (Maybe a)
x Queue a
xs -> Queue a -> Queue a -> Queue a
go (forall a. MVar (Maybe a) -> Queue a -> Queue a
Queue MVar (Maybe a)
x Queue a
acc) Queue a
xs