Variant of withResource that allows explicit resource destruction

[chris-martin]

The resource-pool library assumes that a corrupted resource will cause an action that uses to it to throw an exception. I do not think this is universally true. I have an action involving a query on a database connection where, if the query throws an exception, the action can catch the exception and still return a meaningful value. However, I still want the connection resource to be destroyed when this happens.

I have written the following function for this, and I would like to propose the adding it to the library.

import Control.Exception (mask, onException)
import Control.Monad (unless)
import Control.Monad.IO.Class (MonadIO(..))
import Control.Monad.Trans.Control (MonadBaseControl, control)
import Data.IORef (atomicModifyIORef', newIORef, readIORef)
import Data.Pool (Pool, destroyResource, putResource, takeResource)

{- | Like 'withResource', but allows the action to explicitly destroy the resource

If the action throws an exception, the resource will be destroyed. If the action returns
normally and does not run the destroy action, the resource will be returned to the pool.
-}
withResource' :: (MonadBaseControl IO m, MonadIO m) =>
    Pool a
    -> (a -> m () -> m b) -- ^ The first argument is the resource; the second
                          --   argument is an action to destroy the resource
    -> m b
withResource' pool action = control $ \runInIO -> mask $ \restore ->
  do
    -- Acquire the resource
    (resource, local) <- takeResource pool

    -- Keep track of whether the resource has been destroyed, to avoid destroying it repeatedly
    destroyedRef <- newIORef False

    -- This action destroys the resource, if it has not already been destroyed
    let destroy = do alreadyDestroyed <- atomicModifyIORef' destroyedRef (\x -> (True, x))
                     unless alreadyDestroyed $ destroyResource pool local resource

    -- Run the user's action; if it throws an exception, destroy the resource
    result <- restore (runInIO (action resource (liftIO destroy))) `onException` destroy

    -- Return the resource to the pool, if it has not been destroyed
    destroyed <- readIORef destroyedRef
    unless destroyed $ putResource local resource

    -- Return the result from the user's action
    return result