{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE TypeFamilies #-}
module Data.ConcurrentResourceMap
  ( ConcurrentResourceMap
  , ResourceMap(..)
  , newResourceMap
  , withInitialisedResource
  , withSharedResource
  ) where

import Control.Exception
import qualified Control.Concurrent.MVar as MVar
import Control.Concurrent.MVar (MVar)

data Resource r = Uninitialised | Initialised !r

-- | Some resource with a count of the users (threads) using it.
--
-- Internal invariant: if users = 0 then resource = Uninitialised
data CountedResource r = CountedResource
  { users :: !Int
  , resource :: !(Resource r)
  }

-- | Resource maps should implement this small set of operations that
-- we expect maps to have.
--
-- This allows you to use whatever fast underlying map type you'd
-- like, depending on your resources.
class ResourceMap m where
  type Key m :: *
  empty :: m v
  delete :: Key m -> m v -> m v
  insert :: Key m -> v -> m v -> m v
  lookup :: Key m -> m v -> Maybe v

-- | A map of shared resources @r@ keyed by @k@.
newtype ConcurrentResourceMap m v = C (MVar (m (MVar (CountedResource v))))

-- | Create an empty resource map.
newResourceMap :: ResourceMap m => IO (ConcurrentResourceMap m r)
newResourceMap = fmap C $ MVar.newMVar Data.ConcurrentResourceMap.empty

-- | Use a resource that can be accessed concurrently via multiple
-- threads but is only initialised and destroyed on as-needed basis.
-- If number of users falls to 0, the resource is destroyed. If a new
-- user joins and resource is not available, it's created.
--
-- Calls to 'withSharedResource' can even be nested if you need access
-- to resources with different keys in the same map. Calling
-- 'withSharedResource' in a nested matter on same resource key should
-- have no real adverse effects either.
withSharedResource
  :: ResourceMap m
  => ConcurrentResourceMap m r
  -- ^ Resource map. Create with 'newResourceMap'.
  -> Key m
   -- ^ Key for the resource. This allows you to have many of the same
   -- type of resource but separated: for example, one group of
   -- threads could be holding onto a logging handle to syslog while
   -- another could be holding a handle to a file.
  -> IO r
  -- ^ Initialise resource. Only ran if the resource is not yet
  -- initialised. Does not run in masked context so if you need to
  -- stop async exceptions, you should use 'mask' yourself. If the
  -- action fails (throws an exception), the user fails and we enter
  -- cleanup.
  -> (r -> IO ())
  -- ^ Destroy the resource if it was initialised. Ran by last alive
  -- user when it's exiting. Unlike initialisation, this _is_ ran in
  -- masked context. If this action fails (by throwing an exception
  -- itself), the resource will be assumed to be uninitialised and the
  -- exception will be re-thrown.
  --
  -- Therefore, if your cleanup can fail in a way that you have to
  -- know about/recover from, you should catch exceptions coming out
  -- out 'withSharedResource'. As you get reference to the resource
  -- in the @act@, you're able to store it/monitor it yourself and
  -- decide to take any appropriate actions in the future such as
  -- blocking other threads from running initialisation again until
  -- you've cleaned up the resource yourself.
  -> (r -> IO a)
  -- ^ Run an action with the initialised resource. Note that the
  -- availability of this resource only ensures that the user-given
  -- initialisers/destructors have been ran appropriate number of
  -- times: it of course makes no guarantees as to what the resource
  -- represents. For example, if it's a 'System.Process.ProcessHandle'
  -- or a database connection, there's no guarantee that the process
  -- is alive or that the database connection is still available. For
  -- resources that can dynamically fail, you should implement some
  -- sort of monitoring yourself.
  -> IO a
withSharedResource vm k initResource destroyResource act = bracket
  (addUser vm k)
  (removeUser vm k destroyResource)
  -- Don't leak the internal MVar to the user! This ensures that we
  -- can safely remove it from the resource map when we exit through
  -- 'removeUser'.
  actWithResource
  where
    actWithResource rvar = do
      r <- MVar.modifyMVar rvar $ \cr -> case cr of
        CountedResource { resource = Uninitialised } -> do
          r <- initResource
          return (cr { resource = Initialised r }, r)
        CountedResource { resource = Initialised r } ->
          return (cr, r)
      act r


-- | This is like 'withSharedResource' but will only execute the user
-- action if the resource already exists. This is useful if you create
-- your resources in one place but would like to use them
-- conditionally in another place if they are still alive.
--
-- Action is given Nothing if the resource does not exist or is not
-- initialised.
withInitialisedResource
  :: ResourceMap m
  => ConcurrentResourceMap m r
  -- ^ Resource map. Create with 'newResourceMap'.
  -> Key m
   -- ^ Key for the resource. This allows you to have many of the same
   -- type of resource but separated: for example, one group of
   -- threads could be holding onto a logging handle to syslog while
   -- another could be holding a handle to a file.
  -> (r -> IO ())
  -- ^ Destroy the resource if it was initialised. Ran by last alive
  -- user when it's exiting. Unlike initialisation, this _is_ ran in
  -- masked context. If this action fails (by throwing an exception
  -- itself), the resource will be assumed to be uninitialised and the
  -- exception will be re-thrown.
  --
  -- Therefore, if your cleanup can fail in a way that you have to
  -- know about/recover from, you should catch exceptions coming out
  -- out 'withSharedResource'. As you get reference to the resource
  -- in the @act@, you're able to store it/monitor it yourself and
  -- decide to take any appropriate actions in the future such as
  -- blocking other threads from running initialisation again until
  -- you've cleaned up the resource yourself.
  -> (Maybe r -> IO a)
  -- ^ Run an action with the resource. Note that the availability of
  -- this resource only ensures that the user-given
  -- initialisers/destructors have been ran appropriate number of
  -- times: it of course makes no guarantees as to what the resource
  -- represents. For example, if it's a 'System.Process.ProcessHandle'
  -- or a database connection, there's no guarantee that the process
  -- is alive or that the database connection is still available. For
  -- resources that can dynamically fail, you should implement some
  -- sort of monitoring yourself.
  -> IO a
withInitialisedResource vm k destroyResource act = bracket
  (addUserIfPresent vm k)
  removeUserIfPresent
  actWithResource
  where
    removeUserIfPresent Nothing = pure ()
    removeUserIfPresent (Just rvar) = removeUser vm k destroyResource rvar

    actWithResource Nothing = act Nothing
    actWithResource (Just rvar) = MVar.readMVar rvar >>= \cr -> case cr of
        CountedResource { resource = Initialised r } -> act (Just r)
        _ -> act Nothing

-- | Adds a user at given key. If it's the first user, creates the
-- underlying map.
--
-- Should be used as initialising action in 'bracket' along with
-- 'removeUser'.
addUser
    :: ResourceMap m
    => ConcurrentResourceMap m r -> Key m -> IO (MVar (CountedResource r))
addUser (C vm) k = MVar.modifyMVar vm $ \m -> case Data.ConcurrentResourceMap.lookup k m of
  -- We're the first user of this resource, make the counted
  -- resource.
  Nothing -> do
    v <- MVar.newMVar CountedResource { users = 1, resource = Uninitialised }
    return (Data.ConcurrentResourceMap.insert k v m, v)
  -- Other users already exist, increase the count only.
  Just vc -> do
    MVar.modifyMVar_ vc $ \cr ->
      return cr { users = users cr + 1 }
    return (m, vc)


-- | Adds a user at given key but only if the given key already exists..
--
-- Should be used as initialising action in 'bracket' along with
-- 'removeUser'.
addUserIfPresent
    :: ResourceMap m
    => ConcurrentResourceMap m r
    -> Key m
    -> IO (Maybe (MVar (CountedResource r)))
addUserIfPresent (C vm) k = MVar.modifyMVar vm $ \m -> case Data.ConcurrentResourceMap.lookup k m of
  -- We're the first user of this resource, make the counted
  -- resource.
  Nothing -> return (m, Nothing)
  -- Other users already exist, increase the count only.
  Just vc -> do
    MVar.modifyMVar_ vc $ \cr ->
      return cr { users = users cr + 1 }
    return (m, Just vc)

-- | Remove user for the given key. If it's the last user, removes the
-- counted resource from the map completely.
--
-- Should be used as cleanup action in 'bracket' along with 'addUser'.
removeUser
  :: ResourceMap m
  => ConcurrentResourceMap m r
  -> Key m
  -- ^ The resource from inside the map.
  -> (r -> IO ())
  -- ^ Destroy resource.
  -> MVar (CountedResource r)
  -- ^ Internal ref
  -> IO ()
removeUser (C vm) k destroyResource vc = do
  cr <- MVar.takeMVar vc
  let newCount = users cr - 1
      cr' = cr { users = newCount }
  case cr' of
    -- We're the last ones around, uninitialise.
    CountedResource { users = 0, resource = Initialised r } -> do
      let uninitialise = MVar.putMVar vc cr' { resource = Uninitialised }
          -- Destroy the resource if we can. If we fail, uninitialise it
          -- anyway and re-throw the exception.
          destroy = do
            destroyResource r `onException` uninitialise
            uninitialise
      -- We were the last ones around and whether we managed to
      -- destroy the resource or not, we want to remove the internal
      -- MVar from the resource map if we're still the last ones.
      destroy `finally` cleanFromResourceMap

    -- Resource is uninitialised or there are some other users
    -- around, simply replace the content with updated user counter.
    _ -> MVar.putMVar vc cr'
  where
    cleanFromResourceMap = MVar.modifyMVar_ vm $ \m -> case Data.ConcurrentResourceMap.lookup k m of
      -- The resource is not even in the map anymore. This could
      -- happen if since we decreased the count, a new user came in,
      -- increased the count, finished and cleaned up before we did.
      -- Seems unlikely but not impossible: in this case, there's
      -- nothing left for us to do.
      Nothing -> return m
      Just rvar -> do
        m'cr <- MVar.tryTakeMVar rvar
        case m'cr of
          -- We took the lock and see that there are still no
          -- remaining users. We must be the last ones and must be the
          -- only user holding a useful reference to the MVar as it
          -- can only be created and passed from addUser and it's
          -- internal API. Remove it from the map.
          Just CountedResource { users = 0 } ->
            return (Data.ConcurrentResourceMap.delete k m)
          -- We were able to take the MVar but there are some other
          -- users around again: we want to keep the original map. Put
          -- the value back and keep original the original mapping.
          Just cr -> do
            MVar.putMVar rvar cr
            return m
          -- We were either unable to take the MVar which means
          -- someone else has started to use it and so, we shouldn't
          -- delete it. Leave it as-is in the map.
          Nothing -> return m