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

    -- ** Configuration
  , PoolConfig
  , defaultPoolConfig
  , setNumStripes
  , setPoolLabel

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

import Control.Concurrent.STM
import Control.Exception
import Control.Monad
import Data.Text qualified as T
import GHC.Clock (getMonotonicTime)
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 -> Text
poolLabel :: !T.Text
  , 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 stock (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
$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
/= :: Resource a -> Resource a -> Bool
Eq, (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
$cfrom :: forall a x. Resource a -> Rep (Resource a) x
from :: forall x. Resource a -> Rep (Resource a) x
$cto :: forall a x. Rep (Resource a) x -> Resource a
to :: forall x. Rep (Resource a) x -> Resource a
Generic, 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
$cshowsPrec :: forall a. Show a => Int -> Resource a -> ShowS
showsPrec :: Int -> Resource a -> ShowS
$cshow :: forall a. Show a => Resource a -> String
show :: Resource a -> String
$cshowList :: forall a. Show a => [Resource a] -> ShowS
showList :: [Resource a] -> ShowS
Show)

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

-- | '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 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
  (res, localPool) <- Pool a -> IO (Resource a, LocalPool a)
forall a. Pool a -> IO (Resource a, LocalPool a)
takeResource Pool a
pool
  r <- unmask (act res) `onException` destroyResource pool localPool (resource res)
  putResource localPool (resource res)
  pure 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 = 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
  t1 <- IO Double
getMonotonicTime
  lp <- getLocalPool (localPools pool)
  join . atomically $ do
    stripe <- readTVar (stripeVar lp)
    if available stripe == 0
      then do
        q <- newEmptyTMVar
        writeTVar (stripeVar lp) $! stripe {queueR = Queue q (queueR stripe)}
        pure
          $ waitForResource (stripeVar lp) q >>= \case
            Just a
a -> do
              t2 <- IO Double
getMonotonicTime
              let res =
                    Resource
                      { resource :: a
resource = a
a
                      , poolLabel :: Text
poolLabel = PoolConfig a -> Text
forall a. PoolConfig a -> Text
pcLabel (PoolConfig a -> Text) -> PoolConfig a -> Text
forall a b. (a -> b) -> a -> b
$ Pool a -> PoolConfig a
forall a. Pool a -> PoolConfig a
poolConfig Pool a
pool
                      , 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
                      }
              pure (res, lp)
            Maybe a
Nothing -> do
              t2 <- IO Double
getMonotonicTime
              a <- createResource (poolConfig pool) `onException` restoreSize (stripeVar lp)
              t3 <- getMonotonicTime
              let res =
                    Resource
                      { resource :: a
resource = a
a
                      , poolLabel :: Text
poolLabel = PoolConfig a -> Text
forall a. PoolConfig a -> Text
pcLabel (PoolConfig a -> Text) -> PoolConfig a -> Text
forall a b. (a -> b) -> a -> b
$ Pool a -> PoolConfig a
forall a. Pool a -> PoolConfig a
poolConfig Pool a
pool
                      , 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
                      }
              pure (res, lp)
      else takeAvailableResource pool t1 lp 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 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 a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
    Just (Resource a
res, LocalPool a
localPool) -> do
      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)
      putResource localPool (resource res)
      pure (Just r)
    Maybe (Resource a, LocalPool a)
Nothing -> Maybe r -> IO (Maybe r)
forall a. a -> IO a
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 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 = 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
  t1 <- IO Double
getMonotonicTime
  lp <- getLocalPool (localPools pool)
  join . atomically $ do
    stripe <- readTVar (stripeVar lp)
    if available stripe == 0
      then do
        writeTVar (stripeVar lp) stripe
        pure $ pure Nothing
      else fmap Just <$> takeAvailableResource pool t1 lp stripe

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

takeAvailableResource
  :: Pool a
  -> Double
  -> LocalPool a
  -> Stripe a
  -> STM (IO (Resource a, LocalPool a))
takeAvailableResource :: forall a.
Pool a
-> Double
-> LocalPool a
-> Stripe a
-> STM (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
    TVar (Stripe a) -> Stripe a -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar (LocalPool a -> TVar (Stripe a)
forall a. LocalPool a -> TVar (Stripe a)
stripeVar LocalPool a
lp) (Stripe a -> STM ()) -> Stripe a -> STM ()
forall a b. (a -> b) -> a -> b
$! Stripe a
stripe {available = newAvailable}
    IO (Resource a, LocalPool a) -> STM (IO (Resource a, LocalPool a))
forall a. a -> STM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (IO (Resource a, LocalPool a)
 -> STM (IO (Resource a, LocalPool a)))
-> IO (Resource a, LocalPool a)
-> STM (IO (Resource a, LocalPool a))
forall a b. (a -> b) -> a -> b
$ do
      t2 <- IO Double
getMonotonicTime
      a <- createResource (poolConfig pool) `onException` restoreSize (stripeVar lp)
      t3 <- getMonotonicTime
      let res =
            Resource
              { resource :: a
resource = a
a
              , poolLabel :: Text
poolLabel = PoolConfig a -> Text
forall a. PoolConfig a -> Text
pcLabel (PoolConfig a -> Text) -> PoolConfig a -> Text
forall a b. (a -> b) -> a -> b
$ Pool a -> PoolConfig a
forall a. Pool a -> PoolConfig a
poolConfig Pool a
pool
              , 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
              }
      pure (res, 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
    TVar (Stripe a) -> Stripe a -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar (LocalPool a -> TVar (Stripe a)
forall a. LocalPool a -> TVar (Stripe a)
stripeVar LocalPool a
lp) (Stripe a -> STM ()) -> Stripe a -> STM ()
forall a b. (a -> b) -> a -> b
$! Stripe a
stripe {available = newAvailable, cache = as}
    IO (Resource a, LocalPool a) -> STM (IO (Resource a, LocalPool a))
forall a. a -> STM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (IO (Resource a, LocalPool a)
 -> STM (IO (Resource a, LocalPool a)))
-> IO (Resource a, LocalPool a)
-> STM (IO (Resource a, LocalPool a))
forall a b. (a -> b) -> a -> b
$ do
      t2 <- IO Double
getMonotonicTime
      let res =
            Resource
              { resource :: a
resource = a
a
              , poolLabel :: Text
poolLabel = PoolConfig a -> Text
forall a. PoolConfig a -> Text
pcLabel (PoolConfig a -> Text) -> PoolConfig a -> Text
forall a b. (a -> b) -> a -> b
$ Pool a -> PoolConfig a
forall a. Pool a -> PoolConfig a
poolConfig Pool a
pool
              , 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
              }
      pure (res, lp)