{-# 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 = forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (forall a. Text -> Key a -> Key a
Key Text
n forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> 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) = 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 forall a b. (a -> b) -> a -> b
$ 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 forall a b. (a -> b) -> a -> b
$ 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 forall a b. (a -> b) -> a -> b
$ 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 forall a b. (a -> b) -> a -> b
$ Vault -> Vault -> Vault
V.union Vault
v1 Vault
v2


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


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


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


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


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


lookupBaggage :: Context -> Maybe Baggage
lookupBaggage :: Context -> Maybe Baggage
lookupBaggage = 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 forall a. Key a -> Context -> Maybe a
lookup Key Baggage
baggageKey Context
c of
  Maybe Baggage
Nothing -> forall a. Key a -> a -> Context -> Context
insert Key Baggage
baggageKey Baggage
b Context
c
  Just Baggage
b' -> forall a. Key a -> a -> Context -> Context
insert Key Baggage
baggageKey (Baggage
b forall a. Semigroup a => a -> a -> a
<> Baggage
b') Context
c


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