-- | A variant of "Data.Pool" with introspection capabilities.
module Data.Pool.Introspection
  ( -- * Pool
    Pool
  , LocalPool
  , newPool

  -- ** Configuration
  , PoolConfig
  , defaultPoolConfig
  , setNumStripes

    -- * Resource management
  , Resource(..)
  , Acquisition(..)
  , withResource
  , takeResource
  , tryWithResource
  , tryTakeResource
  , putResource
  , destroyResource
  , destroyAllResources
  ) where

import Control.Concurrent
import Control.Exception
import GHC.Clock
import GHC.Generics (Generic)

import Data.Pool.Internal

-- | A resource taken from the pool along with additional information.
data Resource a = Resource
  { forall a. Resource a -> a
resource           :: a
  , forall a. Resource a -> Int
stripeNumber       :: !Int
  , forall a. Resource a -> Int
availableResources :: !Int
  , forall a. Resource a -> Acquisition
acquisition        :: !Acquisition
  , forall a. Resource a -> Double
acquisitionTime    :: !Double
  , forall a. Resource a -> Maybe Double
creationTime       :: !(Maybe Double)
  } deriving (Resource a -> Resource a -> Bool
forall a. Eq a => Resource a -> Resource a -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Resource a -> Resource a -> Bool
$c/= :: forall a. Eq a => Resource a -> Resource a -> Bool
== :: Resource a -> Resource a -> Bool
$c== :: forall a. Eq a => Resource a -> Resource a -> Bool
Eq, Int -> Resource a -> ShowS
forall a. Show a => Int -> Resource a -> ShowS
forall a. Show a => [Resource a] -> ShowS
forall a. Show a => Resource a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Resource a] -> ShowS
$cshowList :: forall a. Show a => [Resource a] -> ShowS
show :: Resource a -> String
$cshow :: forall a. Show a => Resource a -> String
showsPrec :: Int -> Resource a -> ShowS
$cshowsPrec :: forall a. Show a => Int -> Resource a -> ShowS
Show, forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
forall a x. Rep (Resource a) x -> Resource a
forall a x. Resource a -> Rep (Resource a) x
$cto :: forall a x. Rep (Resource a) x -> Resource a
$cfrom :: forall a x. Resource a -> Rep (Resource a) x
Generic)

-- | Describes how a resource was acquired from the pool.
data Acquisition
  = Immediate
  -- ^ A resource was taken from the pool immediately.
  | Delayed
  -- ^ The thread had to wait until a resource was released.
  deriving (Acquisition -> Acquisition -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Acquisition -> Acquisition -> Bool
$c/= :: Acquisition -> Acquisition -> Bool
== :: Acquisition -> Acquisition -> Bool
$c== :: Acquisition -> Acquisition -> Bool
Eq, Int -> Acquisition -> ShowS
[Acquisition] -> ShowS
Acquisition -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Acquisition] -> ShowS
$cshowList :: [Acquisition] -> ShowS
show :: Acquisition -> String
$cshow :: Acquisition -> String
showsPrec :: Int -> Acquisition -> ShowS
$cshowsPrec :: Int -> Acquisition -> ShowS
Show, forall x. Rep Acquisition x -> Acquisition
forall x. Acquisition -> Rep Acquisition x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cto :: forall x. Rep Acquisition x -> Acquisition
$cfrom :: forall x. Acquisition -> Rep Acquisition x
Generic)

-- | 'Data.Pool.withResource' with introspection capabilities.
withResource :: Pool a -> (Resource a -> IO r) -> IO r
withResource :: forall a r. Pool a -> (Resource a -> IO r) -> IO r
withResource Pool a
pool Resource a -> IO r
act = 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
  (Resource a
res, LocalPool a
localPool) <- forall a. Pool a -> IO (Resource a, LocalPool a)
takeResource Pool a
pool
  r
r <- forall a. IO a -> IO a
unmask (Resource a -> IO r
act Resource a
res) forall a b. IO a -> IO b -> IO a
`onException` forall a. Pool a -> LocalPool a -> a -> IO ()
destroyResource Pool a
pool LocalPool a
localPool (forall a. Resource a -> a
resource Resource a
res)
  forall a. LocalPool a -> a -> IO ()
putResource LocalPool a
localPool (forall a. Resource a -> a
resource Resource a
res)
  forall (f :: * -> *) a. Applicative f => a -> f a
pure r
r

-- | 'Data.Pool.takeResource' with introspection capabilities.
takeResource :: Pool a -> IO (Resource a, LocalPool a)
takeResource :: forall a. Pool a -> IO (Resource a, LocalPool a)
takeResource Pool a
pool = forall a. IO a -> IO a
mask_ forall a b. (a -> b) -> a -> b
$ do
  Double
t1 <- IO Double
getMonotonicTime
  LocalPool a
lp <- forall a. SmallArray (LocalPool a) -> IO (LocalPool a)
getLocalPool (forall a. Pool a -> SmallArray (LocalPool a)
localPools Pool a
pool)
  Stripe a
stripe <- forall a. MVar a -> IO a
takeMVar (forall a. LocalPool a -> MVar (Stripe a)
stripeVar LocalPool a
lp)
  if forall a. Stripe a -> Int
available Stripe a
stripe forall a. Eq a => a -> a -> Bool
== Int
0
    then do
      MVar (Maybe a)
q <- forall a. IO (MVar a)
newEmptyMVar
      forall a. MVar a -> a -> IO ()
putMVar (forall a. LocalPool a -> MVar (Stripe a)
stripeVar LocalPool a
lp) forall a b. (a -> b) -> a -> b
$! Stripe a
stripe { queueR :: Queue a
queueR = forall a. MVar (Maybe a) -> Queue a -> Queue a
Queue MVar (Maybe a)
q (forall a. Stripe a -> Queue a
queueR Stripe a
stripe) }
      forall a. MVar (Stripe a) -> MVar (Maybe a) -> IO (Maybe a)
waitForResource (forall a. LocalPool a -> MVar (Stripe a)
stripeVar LocalPool a
lp) MVar (Maybe a)
q forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
        Just a
a -> do
          Double
t2 <- IO Double
getMonotonicTime
          let res :: Resource a
res = Resource
                { resource :: a
resource           = a
a
                , stripeNumber :: Int
stripeNumber       = forall a. LocalPool a -> Int
stripeId LocalPool a
lp
                , availableResources :: Int
availableResources = Int
0
                , acquisition :: Acquisition
acquisition        = Acquisition
Delayed
                , acquisitionTime :: Double
acquisitionTime    = Double
t2 forall a. Num a => a -> a -> a
- Double
t1
                , creationTime :: Maybe Double
creationTime       = forall a. Maybe a
Nothing
                }
          forall (f :: * -> *) a. Applicative f => a -> f a
pure (Resource a
res, LocalPool a
lp)
        Maybe a
Nothing -> do
          Double
t2 <- IO Double
getMonotonicTime
          a
a  <- forall a. PoolConfig a -> IO a
createResource (forall a. Pool a -> PoolConfig a
poolConfig Pool a
pool) forall a b. IO a -> IO b -> IO a
`onException` forall a. MVar (Stripe a) -> IO ()
restoreSize (forall a. LocalPool a -> MVar (Stripe a)
stripeVar LocalPool a
lp)
          Double
t3 <- IO Double
getMonotonicTime
          let res :: Resource a
res = Resource
                { resource :: a
resource           = a
a
                , stripeNumber :: Int
stripeNumber       = forall a. LocalPool a -> Int
stripeId LocalPool a
lp
                , availableResources :: Int
availableResources = Int
0
                , acquisition :: Acquisition
acquisition        = Acquisition
Delayed
                , acquisitionTime :: Double
acquisitionTime    = Double
t2 forall a. Num a => a -> a -> a
- Double
t1
                , creationTime :: Maybe Double
creationTime       = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$! Double
t3 forall a. Num a => a -> a -> a
- Double
t2
                }
          forall (f :: * -> *) a. Applicative f => a -> f a
pure (Resource a
res, LocalPool a
lp)
    else forall a.
Pool a
-> Double
-> LocalPool a
-> Stripe a
-> IO (Resource a, LocalPool a)
takeAvailableResource Pool a
pool Double
t1 LocalPool a
lp Stripe a
stripe

-- | A variant of 'withResource' that doesn't execute the action and returns
-- 'Nothing' instead of blocking if the local pool is exhausted.
tryWithResource :: Pool a -> (Resource a -> IO r) -> IO (Maybe r)
tryWithResource :: forall a r. Pool a -> (Resource a -> IO r) -> IO (Maybe r)
tryWithResource Pool a
pool Resource a -> IO r
act = 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 -> forall a. Pool a -> IO (Maybe (Resource a, LocalPool a))
tryTakeResource Pool a
pool forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
  Just (Resource a
res, LocalPool a
localPool) -> do
    r
r <- forall a. IO a -> IO a
unmask (Resource a -> IO r
act Resource a
res) forall a b. IO a -> IO b -> IO a
`onException` forall a. Pool a -> LocalPool a -> a -> IO ()
destroyResource Pool a
pool LocalPool a
localPool (forall a. Resource a -> a
resource Resource a
res)
    forall a. LocalPool a -> a -> IO ()
putResource LocalPool a
localPool (forall a. Resource a -> a
resource Resource a
res)
    forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall a. a -> Maybe a
Just r
r)
  Maybe (Resource a, LocalPool a)
Nothing -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a. Maybe a
Nothing

-- | A variant of 'takeResource' that returns 'Nothing' instead of blocking if
-- the local pool is exhausted.
tryTakeResource :: Pool a -> IO (Maybe (Resource a, LocalPool a))
tryTakeResource :: forall a. Pool a -> IO (Maybe (Resource a, LocalPool a))
tryTakeResource Pool a
pool = forall a. IO a -> IO a
mask_ forall a b. (a -> b) -> a -> b
$ do
  Double
t1 <- IO Double
getMonotonicTime
  LocalPool a
lp <- forall a. SmallArray (LocalPool a) -> IO (LocalPool a)
getLocalPool (forall a. Pool a -> SmallArray (LocalPool a)
localPools Pool a
pool)
  Stripe a
stripe <- forall a. MVar a -> IO a
takeMVar (forall a. LocalPool a -> MVar (Stripe a)
stripeVar LocalPool a
lp)
  if forall a. Stripe a -> Int
available Stripe a
stripe forall a. Eq a => a -> a -> Bool
== Int
0
    then do
      forall a. MVar a -> a -> IO ()
putMVar (forall a. LocalPool a -> MVar (Stripe a)
stripeVar LocalPool a
lp) Stripe a
stripe
      forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a. Maybe a
Nothing
    else forall a. a -> Maybe a
Just forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a.
Pool a
-> Double
-> LocalPool a
-> Stripe a
-> IO (Resource a, LocalPool a)
takeAvailableResource Pool a
pool Double
t1 LocalPool a
lp Stripe a
stripe

----------------------------------------
-- Helpers

takeAvailableResource
  :: Pool a
  -> Double
  -> LocalPool a
  -> Stripe a
  -> IO (Resource a, LocalPool a)
takeAvailableResource :: forall a.
Pool a
-> Double
-> LocalPool a
-> Stripe a
-> IO (Resource a, LocalPool a)
takeAvailableResource Pool a
pool Double
t1 LocalPool a
lp Stripe a
stripe = case forall a. Stripe a -> [Entry a]
cache Stripe a
stripe of
  [] -> do
    let newAvailable :: Int
newAvailable = forall a. Stripe a -> Int
available Stripe a
stripe forall a. Num a => a -> a -> a
- Int
1
    forall a. MVar a -> a -> IO ()
putMVar (forall a. LocalPool a -> MVar (Stripe a)
stripeVar LocalPool a
lp) forall a b. (a -> b) -> a -> b
$! Stripe a
stripe { available :: Int
available = Int
newAvailable }
    Double
t2 <- IO Double
getMonotonicTime
    a
a  <- forall a. PoolConfig a -> IO a
createResource (forall a. Pool a -> PoolConfig a
poolConfig Pool a
pool) forall a b. IO a -> IO b -> IO a
`onException` forall a. MVar (Stripe a) -> IO ()
restoreSize (forall a. LocalPool a -> MVar (Stripe a)
stripeVar LocalPool a
lp)
    Double
t3 <- IO Double
getMonotonicTime
    let res :: Resource a
res = Resource
          { resource :: a
resource           = a
a
          , stripeNumber :: Int
stripeNumber       = forall a. LocalPool a -> Int
stripeId LocalPool a
lp
          , availableResources :: Int
availableResources = Int
newAvailable
          , acquisition :: Acquisition
acquisition        = Acquisition
Immediate
          , acquisitionTime :: Double
acquisitionTime    = Double
t2 forall a. Num a => a -> a -> a
- Double
t1
          , creationTime :: Maybe Double
creationTime       = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$! Double
t3 forall a. Num a => a -> a -> a
- Double
t2
          }
    forall (f :: * -> *) a. Applicative f => a -> f a
pure (Resource a
res, LocalPool a
lp)
  Entry a
a Double
_ : [Entry a]
as -> do
    let newAvailable :: Int
newAvailable = forall a. Stripe a -> Int
available Stripe a
stripe forall a. Num a => a -> a -> a
- Int
1
    forall a. MVar a -> a -> IO ()
putMVar (forall a. LocalPool a -> MVar (Stripe a)
stripeVar LocalPool a
lp) forall a b. (a -> b) -> a -> b
$! Stripe a
stripe { available :: Int
available = Int
newAvailable, cache :: [Entry a]
cache = [Entry a]
as }
    Double
t2 <- IO Double
getMonotonicTime
    let res :: Resource a
res = Resource
          { resource :: a
resource           = a
a
          , stripeNumber :: Int
stripeNumber       = forall a. LocalPool a -> Int
stripeId LocalPool a
lp
          , availableResources :: Int
availableResources = Int
newAvailable
          , acquisition :: Acquisition
acquisition        = Acquisition
Immediate
          , acquisitionTime :: Double
acquisitionTime    = Double
t2 forall a. Num a => a -> a -> a
- Double
t1
          , creationTime :: Maybe Double
creationTime       = forall a. Maybe a
Nothing
          }
    forall (f :: * -> *) a. Applicative f => a -> f a
pure (Resource a
res, LocalPool a
lp)