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

    -- * 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
  { Resource a -> a
resource           :: a
  , Resource a -> Int
stripeNumber       :: !Int
  , Resource a -> Int
availableResources :: !Int
  , Resource a -> Acquisition
acquisition        :: !Acquisition
  , Resource a -> Double
acquisitionTime    :: !Double
  , Resource a -> Maybe Double
creationTime       :: !(Maybe Double)
  } deriving (Resource a -> Resource a -> Bool
(Resource a -> Resource a -> Bool)
-> (Resource a -> Resource a -> Bool) -> Eq (Resource a)
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
[Resource a] -> ShowS
Resource a -> String
(Int -> Resource a -> ShowS)
-> (Resource a -> String)
-> ([Resource a] -> ShowS)
-> Show (Resource a)
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 x. Resource a -> Rep (Resource a) x)
-> (forall x. Rep (Resource a) x -> Resource a)
-> Generic (Resource a)
forall x. Rep (Resource a) x -> Resource a
forall x. Resource a -> Rep (Resource a) x
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
(Acquisition -> Acquisition -> Bool)
-> (Acquisition -> Acquisition -> Bool) -> Eq Acquisition
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
(Int -> Acquisition -> ShowS)
-> (Acquisition -> String)
-> ([Acquisition] -> ShowS)
-> Show Acquisition
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. Acquisition -> Rep Acquisition x)
-> (forall x. Rep Acquisition x -> Acquisition)
-> Generic Acquisition
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 :: Pool a -> (Resource a -> IO r) -> IO r
withResource Pool a
pool Resource a -> IO r
act = ((forall a. IO a -> IO a) -> IO r) -> IO r
forall b. ((forall a. IO a -> IO a) -> IO b) -> IO b
mask (((forall a. IO a -> IO a) -> IO r) -> IO r)
-> ((forall a. IO a -> IO a) -> IO r) -> IO r
forall a b. (a -> b) -> a -> b
$ \forall a. IO a -> IO a
unmask -> do
  (Resource a
res, LocalPool a
localPool) <- Pool a -> IO (Resource a, LocalPool a)
forall a. Pool a -> IO (Resource a, LocalPool a)
takeResource Pool a
pool
  r
r <- IO r -> IO r
forall a. IO a -> IO a
unmask (Resource a -> IO r
act Resource a
res) IO r -> IO () -> IO r
forall a b. IO a -> IO b -> IO a
`onException` Pool a -> LocalPool a -> a -> IO ()
forall a. Pool a -> LocalPool a -> a -> IO ()
destroyResource Pool a
pool LocalPool a
localPool (Resource a -> a
forall a. Resource a -> a
resource Resource a
res)
  LocalPool a -> a -> IO ()
forall a. LocalPool a -> a -> IO ()
putResource LocalPool a
localPool (Resource a -> a
forall a. Resource a -> a
resource Resource a
res)
  r -> IO r
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 :: Pool a -> IO (Resource a, LocalPool a)
takeResource Pool a
pool = IO (Resource a, LocalPool a) -> IO (Resource a, LocalPool a)
forall a. IO a -> IO a
mask_ (IO (Resource a, LocalPool a) -> IO (Resource a, LocalPool a))
-> IO (Resource a, LocalPool a) -> IO (Resource a, LocalPool a)
forall a b. (a -> b) -> a -> b
$ do
  Double
t1 <- IO Double
getMonotonicTime
  LocalPool a
lp <- SmallArray (LocalPool a) -> IO (LocalPool a)
forall a. SmallArray (LocalPool a) -> IO (LocalPool a)
getLocalPool (Pool a -> SmallArray (LocalPool a)
forall a. Pool a -> SmallArray (LocalPool a)
localPools Pool a
pool)
  Stripe a
stripe <- MVar (Stripe a) -> IO (Stripe a)
forall a. MVar a -> IO a
takeMVar (LocalPool a -> MVar (Stripe a)
forall a. LocalPool a -> MVar (Stripe a)
stripeVar LocalPool a
lp)
  if Stripe a -> Int
forall a. Stripe a -> Int
available Stripe a
stripe Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0
    then do
      MVar (Maybe a)
q <- IO (MVar (Maybe a))
forall a. IO (MVar a)
newEmptyMVar
      MVar (Stripe a) -> Stripe a -> IO ()
forall a. MVar a -> a -> IO ()
putMVar (LocalPool a -> MVar (Stripe a)
forall a. LocalPool a -> MVar (Stripe a)
stripeVar LocalPool a
lp) (Stripe a -> IO ()) -> Stripe a -> IO ()
forall a b. (a -> b) -> a -> b
$! Stripe a
stripe { queueR :: Queue a
queueR = MVar (Maybe a) -> Queue a -> Queue a
forall a. MVar (Maybe a) -> Queue a -> Queue a
Queue MVar (Maybe a)
q (Stripe a -> Queue a
forall a. Stripe a -> Queue a
queueR Stripe a
stripe) }
      MVar (Stripe a) -> MVar (Maybe a) -> IO (Maybe a)
forall a. MVar (Stripe a) -> MVar (Maybe a) -> IO (Maybe a)
waitForResource (LocalPool a -> MVar (Stripe a)
forall a. LocalPool a -> MVar (Stripe a)
stripeVar LocalPool a
lp) MVar (Maybe a)
q IO (Maybe a)
-> (Maybe a -> IO (Resource a, LocalPool a))
-> IO (Resource a, LocalPool a)
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 :: forall a.
a
-> Int
-> Int
-> Acquisition
-> Double
-> Maybe Double
-> Resource a
Resource
                { resource :: a
resource           = a
a
                , stripeNumber :: Int
stripeNumber       = LocalPool a -> Int
forall a. LocalPool a -> Int
stripeId LocalPool a
lp
                , availableResources :: Int
availableResources = Int
0
                , acquisition :: Acquisition
acquisition        = Acquisition
Delayed
                , acquisitionTime :: Double
acquisitionTime    = Double
t2 Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
t1
                , creationTime :: Maybe Double
creationTime       = Maybe Double
forall a. Maybe a
Nothing
                }
          (Resource a, LocalPool a) -> IO (Resource a, LocalPool a)
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  <- PoolConfig a -> IO a
forall a. PoolConfig a -> IO a
createResource (Pool a -> PoolConfig a
forall a. Pool a -> PoolConfig a
poolConfig Pool a
pool) IO a -> IO () -> IO a
forall a b. IO a -> IO b -> IO a
`onException` MVar (Stripe a) -> IO ()
forall a. MVar (Stripe a) -> IO ()
restoreSize (LocalPool a -> MVar (Stripe a)
forall a. LocalPool a -> MVar (Stripe a)
stripeVar LocalPool a
lp)
          Double
t3 <- IO Double
getMonotonicTime
          let res :: Resource a
res = Resource :: forall a.
a
-> Int
-> Int
-> Acquisition
-> Double
-> Maybe Double
-> Resource a
Resource
                { resource :: a
resource           = a
a
                , stripeNumber :: Int
stripeNumber       = LocalPool a -> Int
forall a. LocalPool a -> Int
stripeId LocalPool a
lp
                , availableResources :: Int
availableResources = Int
0
                , acquisition :: Acquisition
acquisition        = Acquisition
Delayed
                , acquisitionTime :: Double
acquisitionTime    = Double
t2 Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
t1
                , creationTime :: Maybe Double
creationTime       = Double -> Maybe Double
forall a. a -> Maybe a
Just (Double -> Maybe Double) -> Double -> Maybe Double
forall a b. (a -> b) -> a -> b
$! Double
t3 Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
t2
                }
          (Resource a, LocalPool a) -> IO (Resource a, LocalPool a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Resource a
res, LocalPool a
lp)
    else Pool a
-> Double
-> LocalPool a
-> Stripe a
-> IO (Resource a, LocalPool a)
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 capability-local pool is exhausted.
tryWithResource :: Pool a -> (Resource a -> IO r) -> IO (Maybe r)
tryWithResource :: Pool a -> (Resource a -> IO r) -> IO (Maybe r)
tryWithResource Pool a
pool Resource a -> IO r
act = ((forall a. IO a -> IO a) -> IO (Maybe r)) -> IO (Maybe r)
forall b. ((forall a. IO a -> IO a) -> IO b) -> IO b
mask (((forall a. IO a -> IO a) -> IO (Maybe r)) -> IO (Maybe r))
-> ((forall a. IO a -> IO a) -> IO (Maybe r)) -> IO (Maybe r)
forall a b. (a -> b) -> a -> b
$ \forall a. IO a -> IO a
unmask -> Pool a -> IO (Maybe (Resource a, LocalPool a))
forall a. Pool a -> IO (Maybe (Resource a, LocalPool a))
tryTakeResource Pool a
pool IO (Maybe (Resource a, LocalPool a))
-> (Maybe (Resource a, LocalPool a) -> IO (Maybe r))
-> IO (Maybe r)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
  Just (Resource a
res, LocalPool a
localPool) -> do
    r
r <- IO r -> IO r
forall a. IO a -> IO a
unmask (Resource a -> IO r
act Resource a
res) IO r -> IO () -> IO r
forall a b. IO a -> IO b -> IO a
`onException` Pool a -> LocalPool a -> a -> IO ()
forall a. Pool a -> LocalPool a -> a -> IO ()
destroyResource Pool a
pool LocalPool a
localPool (Resource a -> a
forall a. Resource a -> a
resource Resource a
res)
    LocalPool a -> a -> IO ()
forall a. LocalPool a -> a -> IO ()
putResource LocalPool a
localPool (Resource a -> a
forall a. Resource a -> a
resource Resource a
res)
    Maybe r -> IO (Maybe r)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (r -> Maybe r
forall a. a -> Maybe a
Just r
r)
  Maybe (Resource a, LocalPool a)
Nothing -> Maybe r -> IO (Maybe r)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe r
forall a. Maybe a
Nothing

-- | A variant of 'takeResource' that returns 'Nothing' instead of blocking if
-- the capability-local pool is exhausted.
tryTakeResource :: Pool a -> IO (Maybe (Resource a, LocalPool a))
tryTakeResource :: Pool a -> IO (Maybe (Resource a, LocalPool a))
tryTakeResource Pool a
pool = IO (Maybe (Resource a, LocalPool a))
-> IO (Maybe (Resource a, LocalPool a))
forall a. IO a -> IO a
mask_ (IO (Maybe (Resource a, LocalPool a))
 -> IO (Maybe (Resource a, LocalPool a)))
-> IO (Maybe (Resource a, LocalPool a))
-> IO (Maybe (Resource a, LocalPool a))
forall a b. (a -> b) -> a -> b
$ do
  Double
t1 <- IO Double
getMonotonicTime
  LocalPool a
lp <- SmallArray (LocalPool a) -> IO (LocalPool a)
forall a. SmallArray (LocalPool a) -> IO (LocalPool a)
getLocalPool (Pool a -> SmallArray (LocalPool a)
forall a. Pool a -> SmallArray (LocalPool a)
localPools Pool a
pool)
  Stripe a
stripe <- MVar (Stripe a) -> IO (Stripe a)
forall a. MVar a -> IO a
takeMVar (LocalPool a -> MVar (Stripe a)
forall a. LocalPool a -> MVar (Stripe a)
stripeVar LocalPool a
lp)
  if Stripe a -> Int
forall a. Stripe a -> Int
available Stripe a
stripe Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0
    then do
      MVar (Stripe a) -> Stripe a -> IO ()
forall a. MVar a -> a -> IO ()
putMVar (LocalPool a -> MVar (Stripe a)
forall a. LocalPool a -> MVar (Stripe a)
stripeVar LocalPool a
lp) Stripe a
stripe
      Maybe (Resource a, LocalPool a)
-> IO (Maybe (Resource a, LocalPool a))
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe (Resource a, LocalPool a)
forall a. Maybe a
Nothing
    else (Resource a, LocalPool a) -> Maybe (Resource a, LocalPool a)
forall a. a -> Maybe a
Just ((Resource a, LocalPool a) -> Maybe (Resource a, LocalPool a))
-> IO (Resource a, LocalPool a)
-> IO (Maybe (Resource a, LocalPool a))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Pool a
-> Double
-> LocalPool a
-> Stripe a
-> IO (Resource a, LocalPool a)
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 :: 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 Stripe a -> [Entry a]
forall a. Stripe a -> [Entry a]
cache Stripe a
stripe of
  [] -> do
    let newAvailable :: Int
newAvailable = Stripe a -> Int
forall a. Stripe a -> Int
available Stripe a
stripe Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1
    MVar (Stripe a) -> Stripe a -> IO ()
forall a. MVar a -> a -> IO ()
putMVar (LocalPool a -> MVar (Stripe a)
forall a. LocalPool a -> MVar (Stripe a)
stripeVar LocalPool a
lp) (Stripe a -> IO ()) -> Stripe a -> IO ()
forall a b. (a -> b) -> a -> b
$! Stripe a
stripe { available :: Int
available = Int
newAvailable }
    Double
t2 <- IO Double
getMonotonicTime
    a
a  <- PoolConfig a -> IO a
forall a. PoolConfig a -> IO a
createResource (Pool a -> PoolConfig a
forall a. Pool a -> PoolConfig a
poolConfig Pool a
pool) IO a -> IO () -> IO a
forall a b. IO a -> IO b -> IO a
`onException` MVar (Stripe a) -> IO ()
forall a. MVar (Stripe a) -> IO ()
restoreSize (LocalPool a -> MVar (Stripe a)
forall a. LocalPool a -> MVar (Stripe a)
stripeVar LocalPool a
lp)
    Double
t3 <- IO Double
getMonotonicTime
    let res :: Resource a
res = Resource :: forall a.
a
-> Int
-> Int
-> Acquisition
-> Double
-> Maybe Double
-> Resource a
Resource
          { resource :: a
resource           = a
a
          , stripeNumber :: Int
stripeNumber       = LocalPool a -> Int
forall a. LocalPool a -> Int
stripeId LocalPool a
lp
          , availableResources :: Int
availableResources = Int
newAvailable
          , acquisition :: Acquisition
acquisition        = Acquisition
Immediate
          , acquisitionTime :: Double
acquisitionTime    = Double
t2 Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
t1
          , creationTime :: Maybe Double
creationTime       = Double -> Maybe Double
forall a. a -> Maybe a
Just (Double -> Maybe Double) -> Double -> Maybe Double
forall a b. (a -> b) -> a -> b
$! Double
t3 Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
t2
          }
    (Resource a, LocalPool a) -> IO (Resource a, LocalPool a)
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 = Stripe a -> Int
forall a. Stripe a -> Int
available Stripe a
stripe Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1
    MVar (Stripe a) -> Stripe a -> IO ()
forall a. MVar a -> a -> IO ()
putMVar (LocalPool a -> MVar (Stripe a)
forall a. LocalPool a -> MVar (Stripe a)
stripeVar LocalPool a
lp) (Stripe a -> IO ()) -> Stripe a -> IO ()
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 :: forall a.
a
-> Int
-> Int
-> Acquisition
-> Double
-> Maybe Double
-> Resource a
Resource
          { resource :: a
resource           = a
a
          , stripeNumber :: Int
stripeNumber       = LocalPool a -> Int
forall a. LocalPool a -> Int
stripeId LocalPool a
lp
          , availableResources :: Int
availableResources = Int
newAvailable
          , acquisition :: Acquisition
acquisition        = Acquisition
Immediate
          , acquisitionTime :: Double
acquisitionTime    = Double
t2 Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
t1
          , creationTime :: Maybe Double
creationTime       = Maybe Double
forall a. Maybe a
Nothing
          }
    (Resource a, LocalPool a) -> IO (Resource a, LocalPool a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Resource a
res, LocalPool a
lp)