{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE StrictData #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
module OTel.API.Context.Core.Internal
  ( -- * Disclaimer
    -- $disclaimer
    Context(..)
  , emptyContext
  , lookupContext
  , insertContext
  , ContextKey(..)
  , contextKeyName
  , unsafeNewContextKey
  , attachContextValueUsing
  , getAttachedContextValueUsing
  , getAttachedContextUsing
  , ContextBackend(..)
  , unsafeNewContextBackend
  , SomeContextBackend(..)
  , ContextBackendRegistry(..)
  , emptyContextBackendRegistry
  , registerContextBackend
  , defaultContextBackendRegistry
  ) where

import Context (Store)
import Control.Monad.Catch (MonadMask, MonadThrow)
import Control.Monad.IO.Class (MonadIO(liftIO))
import Data.HashMap.Strict (HashMap)
import Data.IORef (IORef)
import Data.Proxy (Proxy(..))
import Data.Text (Text)
import Data.Typeable (Typeable)
import Data.Unique.Really (Unique)
import Data.Vault.Strict (Vault)
import Prelude
import System.IO.Unsafe (unsafePerformIO)
import qualified Context
import qualified Data.HashMap.Strict as HashMap
import qualified Data.IORef as IORef
import qualified Data.Text as Text
import qualified Data.Traversable as Traversable
import qualified Data.Typeable as Typeable
import qualified Data.Unique.Really as Unique
import qualified Data.Vault.Strict as Vault

newtype Context = Context
  { Context -> Vault
unContext :: Vault
  }

emptyContext :: Context
emptyContext :: Context
emptyContext = Vault -> Context
Context Vault
Vault.empty

lookupContext :: ContextKey a -> Context -> Maybe a
lookupContext :: forall a. ContextKey a -> Context -> Maybe a
lookupContext ContextKey a
contextKey Context
context = Key a -> Vault -> Maybe a
forall a. Key a -> Vault -> Maybe a
Vault.lookup Key a
vaultKey Vault
vault
  where
  ContextKey { contextKeyVaultKey :: forall a. ContextKey a -> Key a
contextKeyVaultKey = Key a
vaultKey } = ContextKey a
contextKey
  Context { unContext :: Context -> Vault
unContext = Vault
vault } = Context
context

insertContext :: ContextKey a -> a -> Context -> Context
insertContext :: forall a. ContextKey a -> a -> Context -> Context
insertContext ContextKey a
contextKey a
value Context
context =
  Vault -> Context
Context (Vault -> Context) -> Vault -> Context
forall a b. (a -> b) -> a -> b
$ Key a -> a -> Vault -> Vault
forall a. Key a -> a -> Vault -> Vault
Vault.insert Key a
vaultKey a
value Vault
vault
  where
  ContextKey { contextKeyVaultKey :: forall a. ContextKey a -> Key a
contextKeyVaultKey = Key a
vaultKey } = ContextKey a
contextKey
  Context { unContext :: Context -> Vault
unContext  = Vault
vault } = Context
context

data ContextKey a = ContextKey
  { forall a. ContextKey a -> Text
contextKeyDebugName :: Text
  , forall a. ContextKey a -> Key a
contextKeyVaultKey :: Vault.Key a
  }

contextKeyName :: ContextKey a -> Text
contextKeyName :: forall a. ContextKey a -> Text
contextKeyName = ContextKey a -> Text
forall a. ContextKey a -> Text
contextKeyDebugName

unsafeNewContextKey :: forall m a. (MonadIO m) => Text -> m (ContextKey a)
unsafeNewContextKey :: forall (m :: * -> *) a. MonadIO m => Text -> m (ContextKey a)
unsafeNewContextKey Text
contextKeyDebugName = do
  Key a
contextKeyVaultKey <- IO (Key a) -> m (Key a)
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO IO (Key a)
forall a. IO (Key a)
Vault.newKey
  ContextKey a -> m (ContextKey a)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ContextKey
    { Text
contextKeyDebugName :: Text
contextKeyDebugName :: Text
contextKeyDebugName
    , Key a
contextKeyVaultKey :: Key a
contextKeyVaultKey :: Key a
contextKeyVaultKey
    }

attachContextValueUsing
  :: forall m a b
   . (MonadIO m, MonadMask m)
  => ContextBackend a
  -> a
  -> m b
  -> m b
attachContextValueUsing :: forall (m :: * -> *) a b.
(MonadIO m, MonadMask m) =>
ContextBackend a -> a -> m b -> m b
attachContextValueUsing ContextBackend a
contextBackend a
value =
  Store Context -> (Context -> Context) -> m b -> m b
forall (m :: * -> *) ctx a.
(MonadIO m, MonadMask m) =>
Store ctx -> (ctx -> ctx) -> m a -> m a
Context.adjust (ContextBackend a -> Store Context
forall a. ContextBackend a -> Store Context
contextBackendStore ContextBackend a
contextBackend)
    ((Context -> Context) -> m b -> m b)
-> (Context -> Context) -> m b -> m b
forall a b. (a -> b) -> a -> b
$ ContextKey a -> a -> Context -> Context
forall a. ContextKey a -> a -> Context -> Context
insertContext (ContextBackend a -> ContextKey a
forall a. ContextBackend a -> ContextKey a
contextBackendValueKey ContextBackend a
contextBackend) a
value

getAttachedContextValueUsing
  :: forall m a
   . (MonadIO m, MonadThrow m)
  => ContextBackend a
  -> m (Maybe a)
getAttachedContextValueUsing :: forall (m :: * -> *) a.
(MonadIO m, MonadThrow m) =>
ContextBackend a -> m (Maybe a)
getAttachedContextValueUsing ContextBackend a
contextBackend = do
  Context
context <- Store Context -> m Context
forall (m :: * -> *) ctx.
(MonadIO m, MonadThrow m) =>
Store ctx -> m ctx
Context.mine (Store Context -> m Context) -> Store Context -> m Context
forall a b. (a -> b) -> a -> b
$ ContextBackend a -> Store Context
forall a. ContextBackend a -> Store Context
contextBackendStore ContextBackend a
contextBackend
  Maybe a -> m (Maybe a)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe a -> m (Maybe a)) -> Maybe a -> m (Maybe a)
forall a b. (a -> b) -> a -> b
$ ContextKey a -> Context -> Maybe a
forall a. ContextKey a -> Context -> Maybe a
lookupContext (ContextBackend a -> ContextKey a
forall a. ContextBackend a -> ContextKey a
contextBackendValueKey ContextBackend a
contextBackend) Context
context

getAttachedContextUsing
  :: forall m a
   . (MonadIO m, MonadThrow m)
  => ContextBackend a
  -> m Context
getAttachedContextUsing :: forall (m :: * -> *) a.
(MonadIO m, MonadThrow m) =>
ContextBackend a -> m Context
getAttachedContextUsing ContextBackend a
contextBackend = do
  HashMap Unique SomeContextBackend
someContextBackends <- do
    IO (HashMap Unique SomeContextBackend)
-> m (HashMap Unique SomeContextBackend)
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (HashMap Unique SomeContextBackend)
 -> m (HashMap Unique SomeContextBackend))
-> IO (HashMap Unique SomeContextBackend)
-> m (HashMap Unique SomeContextBackend)
forall a b. (a -> b) -> a -> b
$ IORef (HashMap Unique SomeContextBackend)
-> IO (HashMap Unique SomeContextBackend)
forall a. IORef a -> IO a
IORef.readIORef (IORef (HashMap Unique SomeContextBackend)
 -> IO (HashMap Unique SomeContextBackend))
-> IORef (HashMap Unique SomeContextBackend)
-> IO (HashMap Unique SomeContextBackend)
forall a b. (a -> b) -> a -> b
$ ContextBackendRegistry -> IORef (HashMap Unique SomeContextBackend)
unContextBackendRegistry (ContextBackendRegistry
 -> IORef (HashMap Unique SomeContextBackend))
-> ContextBackendRegistry
-> IORef (HashMap Unique SomeContextBackend)
forall a b. (a -> b) -> a -> b
$ ContextBackend a -> ContextBackendRegistry
forall a. ContextBackend a -> ContextBackendRegistry
contextBackendRegistry ContextBackend a
contextBackend
  ([Vault] -> Context) -> m [Vault] -> m Context
forall a b. (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Vault -> Context
Context (Vault -> Context) -> ([Vault] -> Vault) -> [Vault] -> Context
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Vault] -> Vault
forall a. Monoid a => [a] -> a
mconcat) do
    [SomeContextBackend]
-> (SomeContextBackend -> m Vault) -> m [Vault]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
t a -> (a -> f b) -> f (t b)
Traversable.for (HashMap Unique SomeContextBackend -> [SomeContextBackend]
forall k v. HashMap k v -> [v]
HashMap.elems HashMap Unique SomeContextBackend
someContextBackends) \case
      SomeContextBackend ContextBackend a
registeredContextBackend -> do
        Context
context <- Store Context -> m Context
forall (m :: * -> *) ctx.
(MonadIO m, MonadThrow m) =>
Store ctx -> m ctx
Context.mine (Store Context -> m Context) -> Store Context -> m Context
forall a b. (a -> b) -> a -> b
$ ContextBackend a -> Store Context
forall a. ContextBackend a -> Store Context
contextBackendStore ContextBackend a
registeredContextBackend
        Vault -> m Vault
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Vault -> m Vault) -> Vault -> m Vault
forall a b. (a -> b) -> a -> b
$ Context -> Vault
unContext Context
context

data ContextBackend a = ContextBackend
  { forall a. ContextBackend a -> Store Context
contextBackendStore :: Store Context
  , forall a. ContextBackend a -> ContextKey a
contextBackendValueKey :: ContextKey a
  , forall a. ContextBackend a -> ContextBackendRegistry
contextBackendRegistry :: ContextBackendRegistry
  }

unsafeNewContextBackend :: forall m a. (MonadIO m, Typeable a) => m (ContextBackend a)
unsafeNewContextBackend :: forall (m :: * -> *) a.
(MonadIO m, Typeable a) =>
m (ContextBackend a)
unsafeNewContextBackend = do
  IO (ContextBackend a) -> m (ContextBackend a)
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO do
    ContextBackend a
contextBackend <- do
      ContextKey a
contextBackendValueKey <- do
        Text -> IO (ContextKey a)
forall (m :: * -> *) a. MonadIO m => Text -> m (ContextKey a)
unsafeNewContextKey (Text -> IO (ContextKey a)) -> Text -> IO (ContextKey a)
forall a b. (a -> b) -> a -> b
$ String -> Text
Text.pack (String -> Text) -> String -> Text
forall a b. (a -> b) -> a -> b
$ TypeRep -> String
forall a. Show a => a -> String
show (TypeRep -> String) -> TypeRep -> String
forall a b. (a -> b) -> a -> b
$ Proxy a -> TypeRep
forall {k} (proxy :: k -> *) (a :: k).
Typeable a =>
proxy a -> TypeRep
Typeable.typeRep (Proxy a -> TypeRep) -> Proxy a -> TypeRep
forall a b. (a -> b) -> a -> b
$ forall t. Proxy t
forall {k} (t :: k). Proxy t
Proxy @a
      Store Context
contextBackendStore <- do
        PropagationStrategy -> Maybe Context -> IO (Store Context)
forall (m :: * -> *) ctx.
MonadIO m =>
PropagationStrategy -> Maybe ctx -> m (Store ctx)
Context.newStore PropagationStrategy
Context.noPropagation (Maybe Context -> IO (Store Context))
-> Maybe Context -> IO (Store Context)
forall a b. (a -> b) -> a -> b
$ Context -> Maybe Context
forall a. a -> Maybe a
Just Context
emptyContext
      ContextBackend a -> IO (ContextBackend a)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ContextBackend
        { Store Context
contextBackendStore :: Store Context
contextBackendStore :: Store Context
contextBackendStore
        , ContextKey a
contextBackendValueKey :: ContextKey a
contextBackendValueKey :: ContextKey a
contextBackendValueKey
        , contextBackendRegistry :: ContextBackendRegistry
contextBackendRegistry = ContextBackendRegistry
defaultContextBackendRegistry
        }
    Unique
contextBackendRegistryKey <- IO Unique
Unique.newUnique
    Unique -> ContextBackend a -> ContextBackendRegistry -> IO ()
forall a.
Unique -> ContextBackend a -> ContextBackendRegistry -> IO ()
registerContextBackend Unique
contextBackendRegistryKey ContextBackend a
contextBackend
      (ContextBackendRegistry -> IO ())
-> ContextBackendRegistry -> IO ()
forall a b. (a -> b) -> a -> b
$ ContextBackend a -> ContextBackendRegistry
forall a. ContextBackend a -> ContextBackendRegistry
contextBackendRegistry ContextBackend a
contextBackend
    ContextBackend a -> IO (ContextBackend a)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ContextBackend a
contextBackend

data SomeContextBackend where
  SomeContextBackend :: ContextBackend a -> SomeContextBackend

newtype ContextBackendRegistry = ContextBackendRegistry
  { ContextBackendRegistry -> IORef (HashMap Unique SomeContextBackend)
unContextBackendRegistry :: IORef (HashMap Unique SomeContextBackend)
  }

emptyContextBackendRegistry :: IO ContextBackendRegistry
emptyContextBackendRegistry :: IO ContextBackendRegistry
emptyContextBackendRegistry = do
  (IORef (HashMap Unique SomeContextBackend)
 -> ContextBackendRegistry)
-> IO (IORef (HashMap Unique SomeContextBackend))
-> IO ContextBackendRegistry
forall a b. (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap IORef (HashMap Unique SomeContextBackend) -> ContextBackendRegistry
ContextBackendRegistry (IO (IORef (HashMap Unique SomeContextBackend))
 -> IO ContextBackendRegistry)
-> IO (IORef (HashMap Unique SomeContextBackend))
-> IO ContextBackendRegistry
forall a b. (a -> b) -> a -> b
$ HashMap Unique SomeContextBackend
-> IO (IORef (HashMap Unique SomeContextBackend))
forall a. a -> IO (IORef a)
IORef.newIORef HashMap Unique SomeContextBackend
forall k v. HashMap k v
HashMap.empty

registerContextBackend :: Unique -> ContextBackend a -> ContextBackendRegistry -> IO ()
registerContextBackend :: forall a.
Unique -> ContextBackend a -> ContextBackendRegistry -> IO ()
registerContextBackend Unique
registryKey ContextBackend a
contextBackend ContextBackendRegistry
registry = do
  IORef (HashMap Unique SomeContextBackend)
-> (HashMap Unique SomeContextBackend
    -> (HashMap Unique SomeContextBackend, ()))
-> IO ()
forall a b. IORef a -> (a -> (a, b)) -> IO b
IORef.atomicModifyIORef' IORef (HashMap Unique SomeContextBackend)
ref \HashMap Unique SomeContextBackend
contextBackends ->
    ( Unique
-> SomeContextBackend
-> HashMap Unique SomeContextBackend
-> HashMap Unique SomeContextBackend
forall k v.
(Eq k, Hashable k) =>
k -> v -> HashMap k v -> HashMap k v
HashMap.insert Unique
registryKey (ContextBackend a -> SomeContextBackend
forall a. ContextBackend a -> SomeContextBackend
SomeContextBackend ContextBackend a
contextBackend) HashMap Unique SomeContextBackend
contextBackends
    , ()
    )
  where
  ContextBackendRegistry
    { unContextBackendRegistry :: ContextBackendRegistry -> IORef (HashMap Unique SomeContextBackend)
unContextBackendRegistry = IORef (HashMap Unique SomeContextBackend)
ref
    } = ContextBackendRegistry
registry

defaultContextBackendRegistry :: ContextBackendRegistry
defaultContextBackendRegistry :: ContextBackendRegistry
defaultContextBackendRegistry = IO ContextBackendRegistry -> ContextBackendRegistry
forall a. IO a -> a
unsafePerformIO IO ContextBackendRegistry
emptyContextBackendRegistry
{-# NOINLINE defaultContextBackendRegistry #-}

-- $disclaimer
--
-- In general, changes to this module will not be reflected in the library's
-- version updates. Direct use of this module should be done with utmost care,
-- otherwise invariants will easily be violated.