{-# 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 System.IO.Unsafe
import Prelude hiding (lookup)


newKey :: (MonadIO m) => Text -> m (Key a)
newKey :: forall (m :: * -> *) a. MonadIO m => Text -> m (Key a)
newKey Text
n = IO (Key a) -> m (Key a)
forall a. IO a -> m 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 :: forall a. 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 :: forall a. 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 :: forall a. (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 :: forall a. 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