{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RankNTypes #-}
-----------------------------------------------------------------------------
-- |
-- Module      :  OpenTelemetry.Context
-- Copyright   :  (c) Ian Duncan, 2021
-- License     :  BSD-3
-- Description :  Carrier for execution-scoped values across API boundaries
-- Maintainer  :  Ian Duncan
-- Stability   :  experimental
-- Portability :  non-portable (GHC extensions)
--
-- The ability to correlate events across service boundaries is one of the principle concepts behind distributed tracing. To find these correlations, components in a distributed system need to be able to collect, store, and transfer metadata referred to as context.
--
-- A context will often have information identifying the current span and trace, and can contain arbitrary correlations as key-value pairs.
--
-- Propagation is the means by which context is bundled and transferred in and across services, often via HTTP headers.
--
-- Together, context and propagation represent the engine behind distributed tracing.
--
-----------------------------------------------------------------------------
module OpenTelemetry.Context
  ( Key(keyName)
  , newKey
  , Context
  , HasContext(..)
  , empty
  , lookup
  , insert
  -- , insertWith
  , adjust
  , delete
  , union
  , insertSpan
  , lookupSpan
  , removeSpan
  , insertBaggage
  , lookupBaggage
  , removeBaggage
  ) where
import Control.Monad.IO.Class
import Data.Maybe
import Data.Text (Text)
import qualified Data.Vault.Strict as V
import OpenTelemetry.Baggage (Baggage)
import OpenTelemetry.Context.Types
import OpenTelemetry.Internal.Trace.Types
import Prelude hiding (lookup)
import System.IO.Unsafe

newKey :: MonadIO m => Text -> m (Key a)
newKey :: Text -> m (Key a)
newKey Text
n = IO (Key a) -> m (Key a)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (Text -> Key a -> Key a
forall a. Text -> Key a -> Key a
Key Text
n (Key a -> Key a) -> IO (Key a) -> IO (Key a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IO (Key a)
forall a. IO (Key a)
V.newKey)

class HasContext s where
  contextL :: Lens' s Context

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

lookup :: Key a -> Context -> Maybe a
lookup :: Key a -> Context -> Maybe a
lookup (Key Text
_ Key a
k) (Context Vault
v) = Key a -> Vault -> Maybe a
forall a. Key a -> Vault -> Maybe a
V.lookup Key a
k Vault
v

insert :: Key a -> a -> Context -> Context
insert :: Key a -> a -> Context -> Context
insert (Key Text
_ Key a
k) a
x (Context Vault
v) = 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
V.insert Key a
k a
x Vault
v

-- insertWith 
--   :: (a -> a -> a) 
--   -- ^ new value -> old value -> result
--   -> Key a -> a -> Context -> Context
-- insertWith f (Key _ k) x (Context v) = Context $ case V.lookup k of
--   Nothing -> V.insert k x v
--   Just ox -> V.insert k (f x ox) v

adjust :: (a -> a) -> Key a -> Context -> Context
adjust :: (a -> a) -> Key a -> Context -> Context
adjust a -> a
f (Key Text
_ Key a
k) (Context Vault
v) = Vault -> Context
Context (Vault -> Context) -> Vault -> Context
forall a b. (a -> b) -> a -> b
$ (a -> a) -> Key a -> Vault -> Vault
forall a. (a -> a) -> Key a -> Vault -> Vault
V.adjust a -> a
f Key a
k Vault
v

delete :: Key a -> Context -> Context
delete :: Key a -> Context -> Context
delete (Key Text
_ Key a
k) (Context Vault
v) = Vault -> Context
Context (Vault -> Context) -> Vault -> Context
forall a b. (a -> b) -> a -> b
$ Key a -> Vault -> Vault
forall a. Key a -> Vault -> Vault
V.delete Key a
k Vault
v

union :: Context -> Context -> Context
union :: Context -> Context -> Context
union (Context Vault
v1) (Context Vault
v2) = Vault -> Context
Context (Vault -> Context) -> Vault -> Context
forall a b. (a -> b) -> a -> b
$ Vault -> Vault -> Vault
V.union Vault
v1 Vault
v2

spanKey :: Key Span
spanKey :: Key Span
spanKey = IO (Key Span) -> Key Span
forall a. IO a -> a
unsafePerformIO (IO (Key Span) -> Key Span) -> IO (Key Span) -> Key Span
forall a b. (a -> b) -> a -> b
$ Text -> IO (Key Span)
forall (m :: * -> *) a. MonadIO m => Text -> m (Key a)
newKey Text
"span"
{-# NOINLINE spanKey #-}

lookupSpan :: Context -> Maybe Span
lookupSpan :: Context -> Maybe Span
lookupSpan = Key Span -> Context -> Maybe Span
forall a. Key a -> Context -> Maybe a
lookup Key Span
spanKey

insertSpan :: Span -> Context -> Context
insertSpan :: Span -> Context -> Context
insertSpan = Key Span -> Span -> Context -> Context
forall a. Key a -> a -> Context -> Context
insert Key Span
spanKey

removeSpan :: Context -> Context 
removeSpan :: Context -> Context
removeSpan = Key Span -> Context -> Context
forall a. Key a -> Context -> Context
delete Key Span
spanKey

baggageKey :: Key Baggage
baggageKey :: Key Baggage
baggageKey = IO (Key Baggage) -> Key Baggage
forall a. IO a -> a
unsafePerformIO (IO (Key Baggage) -> Key Baggage)
-> IO (Key Baggage) -> Key Baggage
forall a b. (a -> b) -> a -> b
$ Text -> IO (Key Baggage)
forall (m :: * -> *) a. MonadIO m => Text -> m (Key a)
newKey Text
"baggage"
{-# NOINLINE baggageKey #-}

lookupBaggage :: Context -> Maybe Baggage
lookupBaggage :: Context -> Maybe Baggage
lookupBaggage = Key Baggage -> Context -> Maybe Baggage
forall a. Key a -> Context -> Maybe a
lookup Key Baggage
baggageKey

insertBaggage :: Baggage -> Context -> Context
insertBaggage :: Baggage -> Context -> Context
insertBaggage Baggage
b Context
c = case Key Baggage -> Context -> Maybe Baggage
forall a. Key a -> Context -> Maybe a
lookup Key Baggage
baggageKey Context
c of
  Maybe Baggage
Nothing -> Key Baggage -> Baggage -> Context -> Context
forall a. Key a -> a -> Context -> Context
insert Key Baggage
baggageKey Baggage
b Context
c
  Just Baggage
b' -> Key Baggage -> Baggage -> Context -> Context
forall a. Key a -> a -> Context -> Context
insert Key Baggage
baggageKey (Baggage
b Baggage -> Baggage -> Baggage
forall a. Semigroup a => a -> a -> a
<> Baggage
b') Context
c

removeBaggage :: Context -> Context
removeBaggage :: Context -> Context
removeBaggage = Key Baggage -> Context -> Context
forall a. Key a -> Context -> Context
delete Key Baggage
baggageKey