warp-effectful-1.1.1: src/ThreadEnvs.hs

{-# OPTIONS_GHC -Wno-name-shadowing #-}

module ThreadEnvs where

import Control.Concurrent (ThreadId, myThreadId)
import Control.Exception (bracket_)
import Data.IORef (IORef, atomicModifyIORef', newIORef, readIORef)
import Data.Map.Strict (Map)
import Data.Map.Strict qualified as Map
import Effectful
import Effectful.Dispatch.Static (unEff, unsafeEff)
import Effectful.Internal.Env (Env, cloneEnv)
import Prelude

-- | Per-thread 'Env' cache backing the unlifting function handed to Warp.
data ThreadEnvs es = ThreadEnvs
    { owner :: ThreadId
    , ownerEnv :: Env es
    , template :: Env es
    , cache :: IORef (Map ThreadId (Env es))
    }

new :: Eff es (ThreadEnvs es)
new = unsafeEff \ownerEnv -> do
    owner <- myThreadId
    template <- cloneEnv ownerEnv
    cache <- newIORef mempty
    pure ThreadEnvs{..}

-- | Give the calling thread its own environment for the duration of the action.
with :: ThreadEnvs es -> IO a -> IO a
with ThreadEnvs{..} action = do
    tid <- myThreadId
    es <- cloneEnv template
    bracket_
        (atomicModifyIORef' cache \cache -> (Map.insert tid es cache, ()))
        (atomicModifyIORef' cache \cache -> (Map.delete tid cache, ()))
        action

-- | Unlift using the calling thread's own environment.
unlift :: ThreadEnvs es -> Eff es r -> IO r
unlift ThreadEnvs{..} m = do
    tid <- myThreadId
    if tid == owner
        then unEff m ownerEnv
        else do
            cache <- readIORef cache
            unEff m =<< maybe (cloneEnv template) pure (Map.lookup tid cache)