{-# LANGUAGE DeriveFoldable #-}
{-# LANGUAGE ScopedTypeVariables #-}
-- | Similar to Data.Pool from resource-pool, but resources are
-- identified by some key. To clarify semantics of this module:
--
-- * The pool holds onto and tracks idle resources. Active resources
-- (those checked out via 'takeKeyedPool') are not tracked at all by
-- 'KeyedPool' itself.
--
-- * The pool limits the number of idle resources per key and the
-- total number of idle resources.
--
-- * There is no limit placed on /active/ resources. As such: there
-- will be no delay when calling 'takeKeyedPool': it will either use
-- an idle resource already present, or create a new one
-- immediately.
--
-- * Once the garbage collector cleans up the 'kpAlive' value, the
-- pool will be shut down, by placing a 'PoolClosed' into the
-- 'kpVar' and destroying all existing idle connection.
--
-- * A reaper thread will destroy unused idle resources regularly. It
-- will stop running once 'kpVar' contains a 'PoolClosed' value.
--
-- * 'takeKeyedPool' is async exception safe, but relies on the
-- /caller/ to ensure prompt cleanup. See its comment for more
-- information.
module Data.KeyedPool
    ( KeyedPool
    , createKeyedPool
    , takeKeyedPool
    , Managed
    , managedResource
    , managedReused
    , managedRelease
    , keepAlive
    , Reuse (..)
    , dummyManaged
    ) where

import Control.Concurrent (forkIOWithUnmask, threadDelay)
import Control.Concurrent.STM
import Control.Exception (mask_, catch, SomeException)
import Control.Monad (join, unless, void)
import Data.Map (Map)
import Data.Maybe (isJust)
import qualified Data.Map.Strict as Map
import Data.Time (UTCTime, getCurrentTime, addUTCTime)
import Data.IORef (IORef, newIORef, mkWeakIORef, readIORef)
import qualified Data.Foldable as F
import GHC.Conc (unsafeIOToSTM)
import System.IO.Unsafe (unsafePerformIO)

data KeyedPool key resource = KeyedPool
    { forall key resource. KeyedPool key resource -> key -> IO resource
kpCreate :: !(key -> IO resource)
    , forall key resource. KeyedPool key resource -> resource -> IO ()
kpDestroy :: !(resource -> IO ())
    , forall key resource. KeyedPool key resource -> Int
kpMaxPerKey :: !Int
    , forall key resource. KeyedPool key resource -> Int
kpMaxTotal :: !Int
    , forall key resource.
KeyedPool key resource -> TVar (PoolMap key resource)
kpVar :: !(TVar (PoolMap key resource))
    , forall key resource. KeyedPool key resource -> IORef ()
kpAlive :: !(IORef ())
    }

data PoolMap key resource
    = PoolClosed
    | PoolOpen
        -- Total number of resources in the pool
        {-# UNPACK #-} !Int
        !(Map key (PoolList resource))
    deriving forall a. PoolMap key a -> Bool
forall key a. Eq a => a -> PoolMap key a -> Bool
forall key a. Num a => PoolMap key a -> a
forall key a. Ord a => PoolMap key a -> a
forall m a. Monoid m => (a -> m) -> PoolMap key a -> m
forall key m. Monoid m => PoolMap key m -> m
forall key a. PoolMap key a -> Bool
forall key a. PoolMap key a -> Int
forall key a. PoolMap key a -> [a]
forall a b. (a -> b -> b) -> b -> PoolMap key a -> b
forall key a. (a -> a -> a) -> PoolMap key a -> a
forall key m a. Monoid m => (a -> m) -> PoolMap key a -> m
forall key b a. (b -> a -> b) -> b -> PoolMap key a -> b
forall key a b. (a -> b -> b) -> b -> PoolMap key a -> b
forall (t :: * -> *).
(forall m. Monoid m => t m -> m)
-> (forall m a. Monoid m => (a -> m) -> t a -> m)
-> (forall m a. Monoid m => (a -> m) -> t a -> m)
-> (forall a b. (a -> b -> b) -> b -> t a -> b)
-> (forall a b. (a -> b -> b) -> b -> t a -> b)
-> (forall b a. (b -> a -> b) -> b -> t a -> b)
-> (forall b a. (b -> a -> b) -> b -> t a -> b)
-> (forall a. (a -> a -> a) -> t a -> a)
-> (forall a. (a -> a -> a) -> t a -> a)
-> (forall a. t a -> [a])
-> (forall a. t a -> Bool)
-> (forall a. t a -> Int)
-> (forall a. Eq a => a -> t a -> Bool)
-> (forall a. Ord a => t a -> a)
-> (forall a. Ord a => t a -> a)
-> (forall a. Num a => t a -> a)
-> (forall a. Num a => t a -> a)
-> Foldable t
product :: forall a. Num a => PoolMap key a -> a
$cproduct :: forall key a. Num a => PoolMap key a -> a
sum :: forall a. Num a => PoolMap key a -> a
$csum :: forall key a. Num a => PoolMap key a -> a
minimum :: forall a. Ord a => PoolMap key a -> a
$cminimum :: forall key a. Ord a => PoolMap key a -> a
maximum :: forall a. Ord a => PoolMap key a -> a
$cmaximum :: forall key a. Ord a => PoolMap key a -> a
elem :: forall a. Eq a => a -> PoolMap key a -> Bool
$celem :: forall key a. Eq a => a -> PoolMap key a -> Bool
length :: forall a. PoolMap key a -> Int
$clength :: forall key a. PoolMap key a -> Int
null :: forall a. PoolMap key a -> Bool
$cnull :: forall key a. PoolMap key a -> Bool
toList :: forall a. PoolMap key a -> [a]
$ctoList :: forall key a. PoolMap key a -> [a]
foldl1 :: forall a. (a -> a -> a) -> PoolMap key a -> a
$cfoldl1 :: forall key a. (a -> a -> a) -> PoolMap key a -> a
foldr1 :: forall a. (a -> a -> a) -> PoolMap key a -> a
$cfoldr1 :: forall key a. (a -> a -> a) -> PoolMap key a -> a
foldl' :: forall b a. (b -> a -> b) -> b -> PoolMap key a -> b
$cfoldl' :: forall key b a. (b -> a -> b) -> b -> PoolMap key a -> b
foldl :: forall b a. (b -> a -> b) -> b -> PoolMap key a -> b
$cfoldl :: forall key b a. (b -> a -> b) -> b -> PoolMap key a -> b
foldr' :: forall a b. (a -> b -> b) -> b -> PoolMap key a -> b
$cfoldr' :: forall key a b. (a -> b -> b) -> b -> PoolMap key a -> b
foldr :: forall a b. (a -> b -> b) -> b -> PoolMap key a -> b
$cfoldr :: forall key a b. (a -> b -> b) -> b -> PoolMap key a -> b
foldMap' :: forall m a. Monoid m => (a -> m) -> PoolMap key a -> m
$cfoldMap' :: forall key m a. Monoid m => (a -> m) -> PoolMap key a -> m
foldMap :: forall m a. Monoid m => (a -> m) -> PoolMap key a -> m
$cfoldMap :: forall key m a. Monoid m => (a -> m) -> PoolMap key a -> m
fold :: forall m. Monoid m => PoolMap key m -> m
$cfold :: forall key m. Monoid m => PoolMap key m -> m
F.Foldable

-- | A non-empty list which keeps track of its own length and when
-- each resource was created.
data PoolList a
    = One a {-# UNPACK #-} !UTCTime
    | Cons
        a

        -- size of the list from this point and on
        {-# UNPACK #-} !Int

        {-# UNPACK #-} !UTCTime
        !(PoolList a)
    deriving forall a. Eq a => a -> PoolList a -> Bool
forall a. Num a => PoolList a -> a
forall a. Ord a => PoolList a -> a
forall m. Monoid m => PoolList m -> m
forall a. PoolList a -> Bool
forall a. PoolList a -> Int
forall a. PoolList a -> [a]
forall a. (a -> a -> a) -> PoolList a -> a
forall m a. Monoid m => (a -> m) -> PoolList a -> m
forall b a. (b -> a -> b) -> b -> PoolList a -> b
forall a b. (a -> b -> b) -> b -> PoolList a -> b
forall (t :: * -> *).
(forall m. Monoid m => t m -> m)
-> (forall m a. Monoid m => (a -> m) -> t a -> m)
-> (forall m a. Monoid m => (a -> m) -> t a -> m)
-> (forall a b. (a -> b -> b) -> b -> t a -> b)
-> (forall a b. (a -> b -> b) -> b -> t a -> b)
-> (forall b a. (b -> a -> b) -> b -> t a -> b)
-> (forall b a. (b -> a -> b) -> b -> t a -> b)
-> (forall a. (a -> a -> a) -> t a -> a)
-> (forall a. (a -> a -> a) -> t a -> a)
-> (forall a. t a -> [a])
-> (forall a. t a -> Bool)
-> (forall a. t a -> Int)
-> (forall a. Eq a => a -> t a -> Bool)
-> (forall a. Ord a => t a -> a)
-> (forall a. Ord a => t a -> a)
-> (forall a. Num a => t a -> a)
-> (forall a. Num a => t a -> a)
-> Foldable t
product :: forall a. Num a => PoolList a -> a
$cproduct :: forall a. Num a => PoolList a -> a
sum :: forall a. Num a => PoolList a -> a
$csum :: forall a. Num a => PoolList a -> a
minimum :: forall a. Ord a => PoolList a -> a
$cminimum :: forall a. Ord a => PoolList a -> a
maximum :: forall a. Ord a => PoolList a -> a
$cmaximum :: forall a. Ord a => PoolList a -> a
elem :: forall a. Eq a => a -> PoolList a -> Bool
$celem :: forall a. Eq a => a -> PoolList a -> Bool
length :: forall a. PoolList a -> Int
$clength :: forall a. PoolList a -> Int
null :: forall a. PoolList a -> Bool
$cnull :: forall a. PoolList a -> Bool
toList :: forall a. PoolList a -> [a]
$ctoList :: forall a. PoolList a -> [a]
foldl1 :: forall a. (a -> a -> a) -> PoolList a -> a
$cfoldl1 :: forall a. (a -> a -> a) -> PoolList a -> a
foldr1 :: forall a. (a -> a -> a) -> PoolList a -> a
$cfoldr1 :: forall a. (a -> a -> a) -> PoolList a -> a
foldl' :: forall b a. (b -> a -> b) -> b -> PoolList a -> b
$cfoldl' :: forall b a. (b -> a -> b) -> b -> PoolList a -> b
foldl :: forall b a. (b -> a -> b) -> b -> PoolList a -> b
$cfoldl :: forall b a. (b -> a -> b) -> b -> PoolList a -> b
foldr' :: forall a b. (a -> b -> b) -> b -> PoolList a -> b
$cfoldr' :: forall a b. (a -> b -> b) -> b -> PoolList a -> b
foldr :: forall a b. (a -> b -> b) -> b -> PoolList a -> b
$cfoldr :: forall a b. (a -> b -> b) -> b -> PoolList a -> b
foldMap' :: forall m a. Monoid m => (a -> m) -> PoolList a -> m
$cfoldMap' :: forall m a. Monoid m => (a -> m) -> PoolList a -> m
foldMap :: forall m a. Monoid m => (a -> m) -> PoolList a -> m
$cfoldMap :: forall m a. Monoid m => (a -> m) -> PoolList a -> m
fold :: forall m. Monoid m => PoolList m -> m
$cfold :: forall m. Monoid m => PoolList m -> m
F.Foldable

plistToList :: PoolList a -> [(UTCTime, a)]
plistToList :: forall a. PoolList a -> [(UTCTime, a)]
plistToList (One a
a UTCTime
t) = [(UTCTime
t, a
a)]
plistToList (Cons a
a Int
_ UTCTime
t PoolList a
plist) = (UTCTime
t, a
a) forall a. a -> [a] -> [a]
: forall a. PoolList a -> [(UTCTime, a)]
plistToList PoolList a
plist

plistFromList :: [(UTCTime, a)] -> Maybe (PoolList a)
plistFromList :: forall a. [(UTCTime, a)] -> Maybe (PoolList a)
plistFromList [] = forall a. Maybe a
Nothing
plistFromList [(UTCTime
t, a
a)] = forall a. a -> Maybe a
Just (forall a. a -> UTCTime -> PoolList a
One a
a UTCTime
t)
plistFromList [(UTCTime, a)]
xs =
    forall a. a -> Maybe a
Just forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> b
snd forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall {a}. [(UTCTime, a)] -> (Int, PoolList a)
go forall a b. (a -> b) -> a -> b
$ [(UTCTime, a)]
xs
  where
    go :: [(UTCTime, a)] -> (Int, PoolList a)
go [] = forall a. HasCallStack => [Char] -> a
error [Char]
"plistFromList.go []"
    go [(UTCTime
t, a
a)] = (Int
2, forall a. a -> UTCTime -> PoolList a
One a
a UTCTime
t)
    go ((UTCTime
t, a
a):[(UTCTime, a)]
rest) =
        let (Int
i, PoolList a
rest') = [(UTCTime, a)] -> (Int, PoolList a)
go [(UTCTime, a)]
rest
            i' :: Int
i' = Int
i forall a. Num a => a -> a -> a
+ Int
1
         in Int
i' seq :: forall a b. a -> b -> b
`seq` (Int
i', forall a. a -> Int -> UTCTime -> PoolList a -> PoolList a
Cons a
a Int
i UTCTime
t PoolList a
rest')

-- | Create a new 'KeyedPool' which will automatically clean up after
-- itself when all referenced to the 'KeyedPool' are gone. It will
-- also fork a reaper thread to regularly kill off unused resource.
createKeyedPool
    :: Ord key
    => (key -> IO resource) -- ^ create a new resource
    -> (resource -> IO ())
       -- ^ Destroy a resource. Note that exceptions thrown by this will be
       -- silently discarded. If you want reporting, please install an
       -- exception handler yourself.
    -> Int -- ^ number of resources per key to allow in the pool
    -> Int -- ^ number of resources to allow in the pool across all keys
    -> (SomeException -> IO ()) -- ^ what to do if the reaper throws an exception
    -> IO (KeyedPool key resource)
createKeyedPool :: forall key resource.
Ord key =>
(key -> IO resource)
-> (resource -> IO ())
-> Int
-> Int
-> (SomeException -> IO ())
-> IO (KeyedPool key resource)
createKeyedPool key -> IO resource
create resource -> IO ()
destroy Int
maxPerKey Int
maxTotal SomeException -> IO ()
onReaperException = do
    TVar (PoolMap key resource)
var <- forall a. a -> IO (TVar a)
newTVarIO forall a b. (a -> b) -> a -> b
$ forall key resource.
Int -> Map key (PoolList resource) -> PoolMap key resource
PoolOpen Int
0 forall k a. Map k a
Map.empty

    -- We use a different IORef for the weak ref instead of the var
    -- above since the reaper thread will always be holding onto a
    -- reference.
    IORef ()
alive <- forall a. a -> IO (IORef a)
newIORef ()
    forall (f :: * -> *) a. Functor f => f a -> f ()
void forall a b. (a -> b) -> a -> b
$ forall a. IORef a -> IO () -> IO (Weak (IORef a))
mkWeakIORef IORef ()
alive forall a b. (a -> b) -> a -> b
$ forall resource key.
(resource -> IO ()) -> TVar (PoolMap key resource) -> IO ()
destroyKeyedPool' resource -> IO ()
destroy TVar (PoolMap key resource)
var

    -- Make sure to fork _after_ we've established the mkWeakIORef. If
    -- we did it the other way around, it would be possible for an
    -- async exception to happen before our destroyKeyedPool' handler
    -- was installed, and then reap would have to rely on detecting an
    -- STM deadlock before it could ever exit. This way, the reap
    -- function will only start running when we're guaranteed that
    -- cleanup will be triggered.

    -- Ensure that we have a normal masking state in the new thread.
    ThreadId
_ <- ((forall a. IO a -> IO a) -> IO ()) -> IO ThreadId
forkIOWithUnmask forall a b. (a -> b) -> a -> b
$ \forall a. IO a -> IO a
restore -> forall a. IO a -> IO a
keepRunning forall a b. (a -> b) -> a -> b
$ forall a. IO a -> IO a
restore forall a b. (a -> b) -> a -> b
$ forall key resource.
Ord key =>
(resource -> IO ()) -> TVar (PoolMap key resource) -> IO ()
reap resource -> IO ()
destroy TVar (PoolMap key resource)
var
    forall (m :: * -> *) a. Monad m => a -> m a
return KeyedPool
        { kpCreate :: key -> IO resource
kpCreate = key -> IO resource
create
        , kpDestroy :: resource -> IO ()
kpDestroy = resource -> IO ()
destroy
        , kpMaxPerKey :: Int
kpMaxPerKey = Int
maxPerKey
        , kpMaxTotal :: Int
kpMaxTotal = Int
maxTotal
        , kpVar :: TVar (PoolMap key resource)
kpVar = TVar (PoolMap key resource)
var
        , kpAlive :: IORef ()
kpAlive = IORef ()
alive
        }
  where
    keepRunning :: IO a -> IO a
keepRunning IO a
action =
        IO a
loop
      where
        loop :: IO a
loop = IO a
action forall e a. Exception e => IO a -> (e -> IO a) -> IO a
`catch` \SomeException
e -> SomeException -> IO ()
onReaperException SomeException
e forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> IO a
loop

-- | Make a 'KeyedPool' inactive and destroy all idle resources.
destroyKeyedPool' :: (resource -> IO ())
                  -> TVar (PoolMap key resource)
                  -> IO ()
destroyKeyedPool' :: forall resource key.
(resource -> IO ()) -> TVar (PoolMap key resource) -> IO ()
destroyKeyedPool' resource -> IO ()
destroy TVar (PoolMap key resource)
var = do
    PoolMap key resource
m <- forall a. STM a -> IO a
atomically forall a b. (a -> b) -> a -> b
$ forall a. TVar a -> a -> STM a
swapTVar TVar (PoolMap key resource)
var forall key resource. PoolMap key resource
PoolClosed
    forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
F.mapM_ (IO () -> IO ()
ignoreExceptions forall b c a. (b -> c) -> (a -> b) -> a -> c
. resource -> IO ()
destroy) PoolMap key resource
m

-- | Run a reaper thread, which will destroy old resources. It will
-- stop running once our pool switches to PoolClosed, which is handled
-- via the mkWeakIORef in the creation of the pool.
reap :: forall key resource.
        Ord key
     => (resource -> IO ())
     -> TVar (PoolMap key resource)
     -> IO ()
reap :: forall key resource.
Ord key =>
(resource -> IO ()) -> TVar (PoolMap key resource) -> IO ()
reap resource -> IO ()
destroy TVar (PoolMap key resource)
var =
    IO ()
loop
  where
    loop :: IO ()
loop = do
        Int -> IO ()
threadDelay (Int
5 forall a. Num a => a -> a -> a
* Int
1000 forall a. Num a => a -> a -> a
* Int
1000)
        forall (m :: * -> *) a. Monad m => m (m a) -> m a
join forall a b. (a -> b) -> a -> b
$ forall a. STM a -> IO a
atomically forall a b. (a -> b) -> a -> b
$ do
            PoolMap key resource
m'' <- forall a. TVar a -> STM a
readTVar TVar (PoolMap key resource)
var
            case PoolMap key resource
m'' of
                PoolMap key resource
PoolClosed -> forall (m :: * -> *) a. Monad m => a -> m a
return (forall (m :: * -> *) a. Monad m => a -> m a
return ())
                PoolOpen Int
idleCount Map key (PoolList resource)
m
                    | forall k a. Map k a -> Bool
Map.null Map key (PoolList resource)
m -> forall a. STM a
retry
                    | Bool
otherwise -> do
                        (PoolMap key resource
m', [resource]
toDestroy) <- Int
-> Map key (PoolList resource)
-> STM (PoolMap key resource, [resource])
findStale Int
idleCount Map key (PoolList resource)
m
                        forall a. TVar a -> a -> STM ()
writeTVar TVar (PoolMap key resource)
var PoolMap key resource
m'
                        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ do
                            forall a. IO a -> IO a
mask_ (forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (IO () -> IO ()
ignoreExceptions forall b c a. (b -> c) -> (a -> b) -> a -> c
. resource -> IO ()
destroy) [resource]
toDestroy)
                            IO ()
loop

    findStale :: Int
              -> Map key (PoolList resource)
              -> STM (PoolMap key resource, [resource])
    findStale :: Int
-> Map key (PoolList resource)
-> STM (PoolMap key resource, [resource])
findStale Int
idleCount Map key (PoolList resource)
m = do
        -- We want to make sure to get the time _after_ any delays
        -- occur due to the retry call above. Since getCurrentTime has
        -- no side effects outside of the STM block, this is a safe
        -- usage.
        UTCTime
now <- forall a. IO a -> STM a
unsafeIOToSTM IO UTCTime
getCurrentTime
        let isNotStale :: UTCTime -> Bool
isNotStale UTCTime
time = NominalDiffTime
30 NominalDiffTime -> UTCTime -> UTCTime
`addUTCTime` UTCTime
time forall a. Ord a => a -> a -> Bool
>= UTCTime
now
        let findStale' :: ([(a, PoolList a)] -> [(k, a)])
-> ([a] -> b) -> [(a, PoolList a)] -> (Map k a, b)
findStale' [(a, PoolList a)] -> [(k, a)]
toKeep [a] -> b
toDestroy [] =
                (forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList ([(a, PoolList a)] -> [(k, a)]
toKeep []), [a] -> b
toDestroy [])
            findStale' [(a, PoolList a)] -> [(k, a)]
toKeep [a] -> b
toDestroy ((a
key, PoolList a
plist):[(a, PoolList a)]
rest) =
                ([(a, PoolList a)] -> [(k, a)])
-> ([a] -> b) -> [(a, PoolList a)] -> (Map k a, b)
findStale' [(a, PoolList a)] -> [(k, a)]
toKeep' [a] -> b
toDestroy' [(a, PoolList a)]
rest
              where
                -- Note: By definition, the timestamps must be in
                -- descending order, so we don't need to traverse the
                -- whole list.
                ([(UTCTime, a)]
notStale, [(UTCTime, a)]
stale) = forall a. (a -> Bool) -> [a] -> ([a], [a])
span (UTCTime -> Bool
isNotStale forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> a
fst) forall a b. (a -> b) -> a -> b
$ forall a. PoolList a -> [(UTCTime, a)]
plistToList PoolList a
plist
                toDestroy' :: [a] -> b
toDestroy' = [a] -> b
toDestroy forall b c a. (b -> c) -> (a -> b) -> a -> c
. (forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> b
snd [(UTCTime, a)]
staleforall a. [a] -> [a] -> [a]
++)
                toKeep' :: [(a, PoolList a)] -> [(k, a)]
toKeep' =
                    case forall a. [(UTCTime, a)] -> Maybe (PoolList a)
plistFromList [(UTCTime, a)]
notStale of
                        Maybe (PoolList a)
Nothing -> [(a, PoolList a)] -> [(k, a)]
toKeep
                        Just PoolList a
x -> [(a, PoolList a)] -> [(k, a)]
toKeep forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((a
key, PoolList a
x)forall a. a -> [a] -> [a]
:)
        let (Map key (PoolList resource)
toKeep, [resource]
toDestroy) = forall {k} {a} {a} {a} {b}.
Ord k =>
([(a, PoolList a)] -> [(k, a)])
-> ([a] -> b) -> [(a, PoolList a)] -> (Map k a, b)
findStale' forall a. a -> a
id forall a. a -> a
id (forall k a. Map k a -> [(k, a)]
Map.toList Map key (PoolList resource)
m)
        let idleCount' :: Int
idleCount' = Int
idleCount forall a. Num a => a -> a -> a
- forall (t :: * -> *) a. Foldable t => t a -> Int
length [resource]
toDestroy
        forall (m :: * -> *) a. Monad m => a -> m a
return (forall key resource.
Int -> Map key (PoolList resource) -> PoolMap key resource
PoolOpen Int
idleCount' Map key (PoolList resource)
toKeep, [resource]
toDestroy)

-- | Check out a value from the 'KeyedPool' with the given key.
--
-- This function will internally call 'mask_' to ensure async safety,
-- and will return a value which uses weak references to ensure that
-- the value is cleaned up. However, if you want to ensure timely
-- resource cleanup, you should bracket this operation together with
-- 'managedRelease'.
takeKeyedPool :: Ord key => KeyedPool key resource -> key -> IO (Managed resource)
takeKeyedPool :: forall key resource.
Ord key =>
KeyedPool key resource -> key -> IO (Managed resource)
takeKeyedPool KeyedPool key resource
kp key
key = forall a. IO a -> IO a
mask_ forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a. Monad m => m (m a) -> m a
join forall a b. (a -> b) -> a -> b
$ forall a. STM a -> IO a
atomically forall a b. (a -> b) -> a -> b
$ do
    (PoolMap key resource
m, Maybe resource
mresource) <- forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall {a}. PoolMap key a -> (PoolMap key a, Maybe a)
go forall a b. (a -> b) -> a -> b
$ forall a. TVar a -> STM a
readTVar (forall key resource.
KeyedPool key resource -> TVar (PoolMap key resource)
kpVar KeyedPool key resource
kp)
    forall a. TVar a -> a -> STM ()
writeTVar (forall key resource.
KeyedPool key resource -> TVar (PoolMap key resource)
kpVar KeyedPool key resource
kp) forall a b. (a -> b) -> a -> b
$! PoolMap key resource
m
    forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ do
        resource
resource <- forall b a. b -> (a -> b) -> Maybe a -> b
maybe (forall key resource. KeyedPool key resource -> key -> IO resource
kpCreate KeyedPool key resource
kp key
key) forall (m :: * -> *) a. Monad m => a -> m a
return Maybe resource
mresource
        IORef ()
alive <- forall a. a -> IO (IORef a)
newIORef ()
        TVar Bool
isReleasedVar <- forall a. a -> IO (TVar a)
newTVarIO Bool
False

        let release :: Reuse -> IO ()
release Reuse
action = forall a. IO a -> IO a
mask_ forall a b. (a -> b) -> a -> b
$ do
                Bool
isReleased <- forall a. STM a -> IO a
atomically forall a b. (a -> b) -> a -> b
$ forall a. TVar a -> a -> STM a
swapTVar TVar Bool
isReleasedVar Bool
True
                forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless Bool
isReleased forall a b. (a -> b) -> a -> b
$
                    case Reuse
action of
                        Reuse
Reuse -> forall key resource.
Ord key =>
KeyedPool key resource -> key -> resource -> IO ()
putResource KeyedPool key resource
kp key
key resource
resource
                        Reuse
DontReuse -> IO () -> IO ()
ignoreExceptions forall a b. (a -> b) -> a -> b
$ forall key resource. KeyedPool key resource -> resource -> IO ()
kpDestroy KeyedPool key resource
kp resource
resource

        Weak (IORef ())
_ <- forall a. IORef a -> IO () -> IO (Weak (IORef a))
mkWeakIORef IORef ()
alive forall a b. (a -> b) -> a -> b
$ Reuse -> IO ()
release Reuse
DontReuse
        forall (m :: * -> *) a. Monad m => a -> m a
return Managed
            { _managedResource :: resource
_managedResource = resource
resource
            , _managedReused :: Bool
_managedReused = forall a. Maybe a -> Bool
isJust Maybe resource
mresource
            , _managedRelease :: Reuse -> IO ()
_managedRelease = Reuse -> IO ()
release
            , _managedAlive :: IORef ()
_managedAlive = IORef ()
alive
            }
  where
    go :: PoolMap key a -> (PoolMap key a, Maybe a)
go PoolMap key a
PoolClosed = (forall key resource. PoolMap key resource
PoolClosed, forall a. Maybe a
Nothing)
    go pcOrig :: PoolMap key a
pcOrig@(PoolOpen Int
idleCount Map key (PoolList a)
m) =
        case forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup key
key Map key (PoolList a)
m of
            Maybe (PoolList a)
Nothing -> (PoolMap key a
pcOrig, forall a. Maybe a
Nothing)
            Just (One a
a UTCTime
_) ->
                (forall key resource.
Int -> Map key (PoolList resource) -> PoolMap key resource
PoolOpen (Int
idleCount forall a. Num a => a -> a -> a
- Int
1) (forall k a. Ord k => k -> Map k a -> Map k a
Map.delete key
key Map key (PoolList a)
m), forall a. a -> Maybe a
Just a
a)
            Just (Cons a
a Int
_ UTCTime
_ PoolList a
rest) ->
                (forall key resource.
Int -> Map key (PoolList resource) -> PoolMap key resource
PoolOpen (Int
idleCount forall a. Num a => a -> a -> a
- Int
1) (forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert key
key PoolList a
rest Map key (PoolList a)
m), forall a. a -> Maybe a
Just a
a)

-- | Try to return a resource to the pool. If too many resources
-- already exist, then just destroy it.
putResource :: Ord key => KeyedPool key resource -> key -> resource -> IO ()
putResource :: forall key resource.
Ord key =>
KeyedPool key resource -> key -> resource -> IO ()
putResource KeyedPool key resource
kp key
key resource
resource = do
    UTCTime
now <- IO UTCTime
getCurrentTime
    forall (m :: * -> *) a. Monad m => m (m a) -> m a
join forall a b. (a -> b) -> a -> b
$ forall a. STM a -> IO a
atomically forall a b. (a -> b) -> a -> b
$ do
        (PoolMap key resource
m, IO ()
action) <- forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (UTCTime -> PoolMap key resource -> (PoolMap key resource, IO ())
go UTCTime
now) (forall a. TVar a -> STM a
readTVar (forall key resource.
KeyedPool key resource -> TVar (PoolMap key resource)
kpVar KeyedPool key resource
kp))
        forall a. TVar a -> a -> STM ()
writeTVar (forall key resource.
KeyedPool key resource -> TVar (PoolMap key resource)
kpVar KeyedPool key resource
kp) forall a b. (a -> b) -> a -> b
$! PoolMap key resource
m
        forall (m :: * -> *) a. Monad m => a -> m a
return IO ()
action
  where
    go :: UTCTime -> PoolMap key resource -> (PoolMap key resource, IO ())
go UTCTime
_ PoolMap key resource
PoolClosed = (forall key resource. PoolMap key resource
PoolClosed, forall key resource. KeyedPool key resource -> resource -> IO ()
kpDestroy KeyedPool key resource
kp resource
resource)
    go UTCTime
now pc :: PoolMap key resource
pc@(PoolOpen Int
idleCount Map key (PoolList resource)
m)
        | Int
idleCount forall a. Ord a => a -> a -> Bool
>= forall key resource. KeyedPool key resource -> Int
kpMaxTotal KeyedPool key resource
kp = (PoolMap key resource
pc, forall key resource. KeyedPool key resource -> resource -> IO ()
kpDestroy KeyedPool key resource
kp resource
resource)
        | Bool
otherwise = case forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup key
key Map key (PoolList resource)
m of
            Maybe (PoolList resource)
Nothing ->
                let cnt' :: Int
cnt' = Int
idleCount forall a. Num a => a -> a -> a
+ Int
1
                    m' :: PoolMap key resource
m' = forall key resource.
Int -> Map key (PoolList resource) -> PoolMap key resource
PoolOpen Int
cnt' (forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert key
key (forall a. a -> UTCTime -> PoolList a
One resource
resource UTCTime
now) Map key (PoolList resource)
m)
                 in (PoolMap key resource
m', forall (m :: * -> *) a. Monad m => a -> m a
return ())
            Just PoolList resource
l ->
                let (PoolList resource
l', Maybe resource
mx) = forall a.
UTCTime -> Int -> a -> PoolList a -> (PoolList a, Maybe a)
addToList UTCTime
now (forall key resource. KeyedPool key resource -> Int
kpMaxPerKey KeyedPool key resource
kp) resource
resource PoolList resource
l
                    cnt' :: Int
cnt' = Int
idleCount forall a. Num a => a -> a -> a
+ forall b a. b -> (a -> b) -> Maybe a -> b
maybe Int
1 (forall a b. a -> b -> a
const Int
0) Maybe resource
mx
                    m' :: PoolMap key resource
m' = forall key resource.
Int -> Map key (PoolList resource) -> PoolMap key resource
PoolOpen Int
cnt' (forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert key
key PoolList resource
l' Map key (PoolList resource)
m)
                 in (PoolMap key resource
m', forall b a. b -> (a -> b) -> Maybe a -> b
maybe (forall (m :: * -> *) a. Monad m => a -> m a
return ()) (forall key resource. KeyedPool key resource -> resource -> IO ()
kpDestroy KeyedPool key resource
kp) Maybe resource
mx)

-- | Add a new element to the list, up to the given maximum number. If we're
-- already at the maximum, return the new value as leftover.
addToList :: UTCTime -> Int -> a -> PoolList a -> (PoolList a, Maybe a)
addToList :: forall a.
UTCTime -> Int -> a -> PoolList a -> (PoolList a, Maybe a)
addToList UTCTime
_ Int
i a
x PoolList a
l | Int
i forall a. Ord a => a -> a -> Bool
<= Int
1 = (PoolList a
l, forall a. a -> Maybe a
Just a
x)
addToList UTCTime
now Int
_ a
x l :: PoolList a
l@One{} = (forall a. a -> Int -> UTCTime -> PoolList a -> PoolList a
Cons a
x Int
2 UTCTime
now PoolList a
l, forall a. Maybe a
Nothing)
addToList UTCTime
now Int
maxCount a
x l :: PoolList a
l@(Cons a
_ Int
currCount UTCTime
_ PoolList a
_)
    | Int
maxCount forall a. Ord a => a -> a -> Bool
> Int
currCount = (forall a. a -> Int -> UTCTime -> PoolList a -> PoolList a
Cons a
x (Int
currCount forall a. Num a => a -> a -> a
+ Int
1) UTCTime
now PoolList a
l, forall a. Maybe a
Nothing)
    | Bool
otherwise = (PoolList a
l, forall a. a -> Maybe a
Just a
x)

-- | A managed resource, which can be returned to the 'KeyedPool' when
-- work with it is complete. Using garbage collection, it will default
-- to destroying the resource if the caller does not explicitly use
-- 'managedRelease'.
data Managed resource = Managed
    { forall resource. Managed resource -> resource
_managedResource :: !resource
    , forall resource. Managed resource -> Bool
_managedReused :: !Bool
    , forall resource. Managed resource -> Reuse -> IO ()
_managedRelease :: !(Reuse -> IO ())
    , forall resource. Managed resource -> IORef ()
_managedAlive :: !(IORef ())
    }

-- | Get the raw resource from the 'Managed' value.
managedResource :: Managed resource -> resource
managedResource :: forall resource. Managed resource -> resource
managedResource = forall resource. Managed resource -> resource
_managedResource

-- | Was this value taken from the pool?
managedReused :: Managed resource -> Bool
managedReused :: forall resource. Managed resource -> Bool
managedReused = forall resource. Managed resource -> Bool
_managedReused

-- | Release the resource, after which it is invalid to use the
-- 'managedResource' value. 'Reuse' returns the resource to the
-- pool; 'DontReuse' destroys it.
managedRelease :: Managed resource -> Reuse -> IO ()
managedRelease :: forall resource. Managed resource -> Reuse -> IO ()
managedRelease = forall resource. Managed resource -> Reuse -> IO ()
_managedRelease

data Reuse = Reuse | DontReuse

-- | For testing purposes only: create a dummy Managed wrapper
dummyManaged :: resource -> Managed resource
dummyManaged :: forall resource. resource -> Managed resource
dummyManaged resource
resource = Managed
    { _managedResource :: resource
_managedResource = resource
resource
    , _managedReused :: Bool
_managedReused = Bool
False
    , _managedRelease :: Reuse -> IO ()
_managedRelease = forall a b. a -> b -> a
const (forall (m :: * -> *) a. Monad m => a -> m a
return ())
    , _managedAlive :: IORef ()
_managedAlive = forall a. IO a -> a
unsafePerformIO (forall a. a -> IO (IORef a)
newIORef ())
    }

ignoreExceptions :: IO () -> IO ()
ignoreExceptions :: IO () -> IO ()
ignoreExceptions IO ()
f = IO ()
f forall e a. Exception e => IO a -> (e -> IO a) -> IO a
`catch` \(SomeException
_ :: SomeException) -> forall (m :: * -> *) a. Monad m => a -> m a
return ()

-- | Prevent the managed resource from getting released before you want to use.
keepAlive :: Managed resource -> IO ()
keepAlive :: forall resource. Managed resource -> IO ()
keepAlive = forall a. IORef a -> IO a
readIORef forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall resource. Managed resource -> IORef ()
_managedAlive