{-# LANGUAGE BangPatterns #-}
module OpenTelemetry.Context.ThreadLocal (
Token,
getContext,
lookupContext,
attachContext,
detachContext,
adjustContext,
getAndAdjustContext,
getActiveBaggage,
setActiveBaggage,
clearActiveBaggage,
ContextEntry (ceContext),
emptyEntry,
ensureContextRef,
ensureContextRefFast,
lookupContextRefFast,
readContextRef,
writeContextRef,
lookupContextOnThread,
attachContextOnThread,
detachContextFromThread,
adjustContextOnThread,
threadContextMap,
) where
import Control.Concurrent
import Control.Concurrent.Thread.Storage
import Control.Monad (when)
import Control.Monad.IO.Class
import Data.IORef
import Data.Maybe (fromMaybe)
import Data.Word (Word64)
import OpenTelemetry.Baggage (Baggage)
import OpenTelemetry.Context (Context, empty, insertBaggage, lookupBaggage, removeBaggage)
import OpenTelemetry.Internal.Logging (otelLogError)
import System.IO.Unsafe
import Prelude hiding (lookup)
data ContextEntry = ContextEntry
{ ContextEntry -> Context
ceContext :: !Context
, ContextEntry -> Word64
ceTokenId :: {-# UNPACK #-} !Word64
}
emptyEntry :: ContextEntry
emptyEntry :: ContextEntry
emptyEntry = Context -> Word64 -> ContextEntry
ContextEntry Context
empty Word64
0
{-# INLINE emptyEntry #-}
data Token = Token
{ Token -> Word64
_tokenId :: {-# UNPACK #-} !Word64
, Token -> Context
_tokenPreviousContext :: !Context
, Token -> Word64
_tokenPreviousTokenId :: {-# UNPACK #-} !Word64
}
type ThreadContextMap = ThreadStorageMap ContextEntry
tokenCounter :: IORef Word64
tokenCounter :: IORef Word64
tokenCounter = IO (IORef Word64) -> IORef Word64
forall a. IO a -> a
unsafePerformIO (Word64 -> IO (IORef Word64)
forall a. a -> IO (IORef a)
newIORef Word64
0)
{-# NOINLINE tokenCounter #-}
nextTokenId :: IO Word64
nextTokenId :: IO Word64
nextTokenId = IORef Word64 -> (Word64 -> (Word64, Word64)) -> IO Word64
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' IORef Word64
tokenCounter (\Word64
n -> let !n' :: Word64
n' = Word64
n Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
+ Word64
1 in (Word64
n', Word64
n'))
{-# INLINE nextTokenId #-}
threadContextMap :: ThreadContextMap
threadContextMap :: ThreadContextMap
threadContextMap = IO ThreadContextMap -> ThreadContextMap
forall a. IO a -> a
unsafePerformIO IO ThreadContextMap
forall (m :: * -> *) a. MonadIO m => m (ThreadStorageMap a)
newThreadStorageMap
{-# NOINLINE threadContextMap #-}
getContext :: (MonadIO m) => m Context
getContext :: forall (m :: * -> *). MonadIO m => m Context
getContext = do
me <- ThreadContextMap -> m (Maybe ContextEntry)
forall (m :: * -> *) a.
MonadIO m =>
ThreadStorageMap a -> m (Maybe a)
lookup ThreadContextMap
threadContextMap
pure $! case me of
Maybe ContextEntry
Nothing -> Context
empty
Just ContextEntry
entry -> ContextEntry -> Context
ceContext ContextEntry
entry
{-# INLINE getContext #-}
lookupContext :: (MonadIO m) => m (Maybe Context)
lookupContext :: forall (m :: * -> *). MonadIO m => m (Maybe Context)
lookupContext = (Maybe ContextEntry -> Maybe Context)
-> m (Maybe ContextEntry) -> m (Maybe Context)
forall a b. (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((ContextEntry -> Context) -> Maybe ContextEntry -> Maybe Context
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ContextEntry -> Context
ceContext) (ThreadContextMap -> m (Maybe ContextEntry)
forall (m :: * -> *) a.
MonadIO m =>
ThreadStorageMap a -> m (Maybe a)
lookup ThreadContextMap
threadContextMap)
{-# INLINE lookupContext #-}
lookupContextOnThread :: (MonadIO m) => ThreadId -> m (Maybe Context)
lookupContextOnThread :: forall (m :: * -> *). MonadIO m => ThreadId -> m (Maybe Context)
lookupContextOnThread ThreadId
tid = (Maybe ContextEntry -> Maybe Context)
-> m (Maybe ContextEntry) -> m (Maybe Context)
forall a b. (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((ContextEntry -> Context) -> Maybe ContextEntry -> Maybe Context
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ContextEntry -> Context
ceContext) (ThreadContextMap -> ThreadId -> m (Maybe ContextEntry)
forall (m :: * -> *) a.
MonadIO m =>
ThreadStorageMap a -> ThreadId -> m (Maybe a)
lookupOnThread ThreadContextMap
threadContextMap ThreadId
tid)
{-# INLINE lookupContextOnThread #-}
attachContext :: (MonadIO m) => Context -> m Token
attachContext :: forall (m :: * -> *). MonadIO m => Context -> m Token
attachContext Context
newCtx = IO Token -> m Token
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO Token -> m Token) -> IO Token -> m Token
forall a b. (a -> b) -> a -> b
$ do
tokId <- IO Word64
nextTokenId
update threadContextMap $ \Maybe ContextEntry
mentry ->
let !old :: ContextEntry
old = ContextEntry -> Maybe ContextEntry -> ContextEntry
forall a. a -> Maybe a -> a
fromMaybe ContextEntry
emptyEntry Maybe ContextEntry
mentry
!tok :: Token
tok = Word64 -> Context -> Word64 -> Token
Token Word64
tokId (ContextEntry -> Context
ceContext ContextEntry
old) (ContextEntry -> Word64
ceTokenId ContextEntry
old)
!new :: ContextEntry
new = Context -> Word64 -> ContextEntry
ContextEntry Context
newCtx Word64
tokId
in (ContextEntry -> Maybe ContextEntry
forall a. a -> Maybe a
Just ContextEntry
new, Token
tok)
{-# INLINE attachContext #-}
attachContextOnThread :: (MonadIO m) => ThreadId -> Context -> m Token
attachContextOnThread :: forall (m :: * -> *). MonadIO m => ThreadId -> Context -> m Token
attachContextOnThread ThreadId
tid Context
newCtx = IO Token -> m Token
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO Token -> m Token) -> IO Token -> m Token
forall a b. (a -> b) -> a -> b
$ do
tokId <- IO Word64
nextTokenId
updateOnThread threadContextMap tid $ \Maybe ContextEntry
mentry ->
let !old :: ContextEntry
old = ContextEntry -> Maybe ContextEntry -> ContextEntry
forall a. a -> Maybe a -> a
fromMaybe ContextEntry
emptyEntry Maybe ContextEntry
mentry
!tok :: Token
tok = Word64 -> Context -> Word64 -> Token
Token Word64
tokId (ContextEntry -> Context
ceContext ContextEntry
old) (ContextEntry -> Word64
ceTokenId ContextEntry
old)
!new :: ContextEntry
new = Context -> Word64 -> ContextEntry
ContextEntry Context
newCtx Word64
tokId
in (ContextEntry -> Maybe ContextEntry
forall a. a -> Maybe a
Just ContextEntry
new, Token
tok)
{-# INLINE attachContextOnThread #-}
detachContext :: (MonadIO m) => Token -> m ()
detachContext :: forall (m :: * -> *). MonadIO m => Token -> m ()
detachContext (Token Word64
expectedId Context
prevCtx Word64
prevTokenId) = IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ do
mismatch <- ThreadContextMap
-> (Maybe ContextEntry -> (Maybe ContextEntry, Bool)) -> IO Bool
forall (m :: * -> *) a b.
MonadIO m =>
ThreadStorageMap a -> (Maybe a -> (Maybe a, b)) -> m b
update ThreadContextMap
threadContextMap ((Maybe ContextEntry -> (Maybe ContextEntry, Bool)) -> IO Bool)
-> (Maybe ContextEntry -> (Maybe ContextEntry, Bool)) -> IO Bool
forall a b. (a -> b) -> a -> b
$ \Maybe ContextEntry
mentry ->
let !current :: ContextEntry
current = ContextEntry -> Maybe ContextEntry -> ContextEntry
forall a. a -> Maybe a -> a
fromMaybe ContextEntry
emptyEntry Maybe ContextEntry
mentry
!restored :: ContextEntry
restored = Context -> Word64 -> ContextEntry
ContextEntry Context
prevCtx Word64
prevTokenId
in (ContextEntry -> Maybe ContextEntry
forall a. a -> Maybe a
Just ContextEntry
restored, ContextEntry -> Word64
ceTokenId ContextEntry
current Word64 -> Word64 -> Bool
forall a. Eq a => a -> a -> Bool
/= Word64
expectedId)
when mismatch $
otelLogError
"Context detach token mismatch: LIFO ordering violated. \
\This likely indicates a context leak — an attachContext call \
\without a corresponding detachContext in the correct order."
{-# INLINE detachContext #-}
detachContextFromThread :: (MonadIO m) => ThreadId -> Token -> m ()
detachContextFromThread :: forall (m :: * -> *). MonadIO m => ThreadId -> Token -> m ()
detachContextFromThread ThreadId
tid (Token Word64
expectedId Context
prevCtx Word64
prevTokenId) = IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ do
mismatch <- ThreadContextMap
-> ThreadId
-> (Maybe ContextEntry -> (Maybe ContextEntry, Bool))
-> IO Bool
forall (m :: * -> *) a b.
MonadIO m =>
ThreadStorageMap a -> ThreadId -> (Maybe a -> (Maybe a, b)) -> m b
updateOnThread ThreadContextMap
threadContextMap ThreadId
tid ((Maybe ContextEntry -> (Maybe ContextEntry, Bool)) -> IO Bool)
-> (Maybe ContextEntry -> (Maybe ContextEntry, Bool)) -> IO Bool
forall a b. (a -> b) -> a -> b
$ \Maybe ContextEntry
mentry ->
let !current :: ContextEntry
current = ContextEntry -> Maybe ContextEntry -> ContextEntry
forall a. a -> Maybe a -> a
fromMaybe ContextEntry
emptyEntry Maybe ContextEntry
mentry
!restored :: ContextEntry
restored = Context -> Word64 -> ContextEntry
ContextEntry Context
prevCtx Word64
prevTokenId
in (ContextEntry -> Maybe ContextEntry
forall a. a -> Maybe a
Just ContextEntry
restored, ContextEntry -> Word64
ceTokenId ContextEntry
current Word64 -> Word64 -> Bool
forall a. Eq a => a -> a -> Bool
/= Word64
expectedId)
when mismatch $
otelLogError
"Context detach token mismatch on remote thread: LIFO ordering violated."
{-# INLINE detachContextFromThread #-}
adjustContext :: (MonadIO m) => (Context -> Context) -> m ()
adjustContext :: forall (m :: * -> *). MonadIO m => (Context -> Context) -> m ()
adjustContext Context -> Context
f = ThreadContextMap
-> (Maybe ContextEntry -> (Maybe ContextEntry, ())) -> m ()
forall (m :: * -> *) a b.
MonadIO m =>
ThreadStorageMap a -> (Maybe a -> (Maybe a, b)) -> m b
update ThreadContextMap
threadContextMap ((Maybe ContextEntry -> (Maybe ContextEntry, ())) -> m ())
-> (Maybe ContextEntry -> (Maybe ContextEntry, ())) -> m ()
forall a b. (a -> b) -> a -> b
$ \Maybe ContextEntry
mentry ->
let !old :: ContextEntry
old = ContextEntry -> Maybe ContextEntry -> ContextEntry
forall a. a -> Maybe a -> a
fromMaybe ContextEntry
emptyEntry Maybe ContextEntry
mentry
!ctx' :: Context
ctx' = Context -> Context
f (ContextEntry -> Context
ceContext ContextEntry
old)
in (ContextEntry -> Maybe ContextEntry
forall a. a -> Maybe a
Just (ContextEntry
old {ceContext = ctx'}), ())
{-# INLINE adjustContext #-}
getAndAdjustContext :: (MonadIO m) => (Context -> Context) -> m Context
getAndAdjustContext :: forall (m :: * -> *).
MonadIO m =>
(Context -> Context) -> m Context
getAndAdjustContext Context -> Context
f = ThreadContextMap
-> (Maybe ContextEntry -> (Maybe ContextEntry, Context))
-> m Context
forall (m :: * -> *) a b.
MonadIO m =>
ThreadStorageMap a -> (Maybe a -> (Maybe a, b)) -> m b
update ThreadContextMap
threadContextMap ((Maybe ContextEntry -> (Maybe ContextEntry, Context))
-> m Context)
-> (Maybe ContextEntry -> (Maybe ContextEntry, Context))
-> m Context
forall a b. (a -> b) -> a -> b
$ \Maybe ContextEntry
mentry ->
let !old :: ContextEntry
old = ContextEntry -> Maybe ContextEntry -> ContextEntry
forall a. a -> Maybe a -> a
fromMaybe ContextEntry
emptyEntry Maybe ContextEntry
mentry
!ctx :: Context
ctx = ContextEntry -> Context
ceContext ContextEntry
old
!new :: Context
new = Context -> Context
f Context
ctx
in (ContextEntry -> Maybe ContextEntry
forall a. a -> Maybe a
Just (ContextEntry
old {ceContext = new}), Context
ctx)
{-# INLINE getAndAdjustContext #-}
ensureContextRef :: ThreadId -> Int -> IO (IORef ContextEntry)
ensureContextRef :: ThreadId -> Int -> IO (IORef ContextEntry)
ensureContextRef ThreadId
tid Int
tw = ThreadContextMap
-> ThreadId -> Int -> ContextEntry -> IO (IORef ContextEntry)
forall a.
ThreadStorageMap a -> ThreadId -> Int -> a -> IO (IORef a)
ensureRef ThreadContextMap
threadContextMap ThreadId
tid Int
tw ContextEntry
emptyEntry
{-# INLINE ensureContextRef #-}
ensureContextRefFast :: IO (Int, IORef ContextEntry)
ensureContextRefFast :: IO (Int, IORef ContextEntry)
ensureContextRefFast = ThreadContextMap -> ContextEntry -> IO (Int, IORef ContextEntry)
forall a. ThreadStorageMap a -> a -> IO (Int, IORef a)
ensureRefFast ThreadContextMap
threadContextMap ContextEntry
emptyEntry
{-# INLINE ensureContextRefFast #-}
lookupContextRefFast :: IO (Int, Maybe (IORef ContextEntry))
lookupContextRefFast :: IO (Int, Maybe (IORef ContextEntry))
lookupContextRefFast = ThreadContextMap -> IO (Int, Maybe (IORef ContextEntry))
forall a. ThreadStorageMap a -> IO (Int, Maybe (IORef a))
lookupRefFast ThreadContextMap
threadContextMap
{-# INLINE lookupContextRefFast #-}
readContextRef :: IORef ContextEntry -> IO Context
readContextRef :: IORef ContextEntry -> IO Context
readContextRef IORef ContextEntry
ref = ContextEntry -> Context
ceContext (ContextEntry -> Context) -> IO ContextEntry -> IO Context
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IORef ContextEntry -> IO ContextEntry
forall a. IORef a -> IO a
readIORef IORef ContextEntry
ref
{-# INLINE readContextRef #-}
writeContextRef :: IORef ContextEntry -> Context -> IO ()
writeContextRef :: IORef ContextEntry -> Context -> IO ()
writeContextRef IORef ContextEntry
ref Context
ctx = IORef ContextEntry -> (ContextEntry -> ContextEntry) -> IO ()
forall a. IORef a -> (a -> a) -> IO ()
modifyIORef' IORef ContextEntry
ref (\ContextEntry
e -> ContextEntry
e {ceContext = ctx})
{-# INLINE writeContextRef #-}
adjustContextOnThread :: (MonadIO m) => ThreadId -> (Context -> Context) -> m ()
adjustContextOnThread :: forall (m :: * -> *).
MonadIO m =>
ThreadId -> (Context -> Context) -> m ()
adjustContextOnThread ThreadId
tid Context -> Context
f = ThreadContextMap
-> ThreadId
-> (Maybe ContextEntry -> (Maybe ContextEntry, ()))
-> m ()
forall (m :: * -> *) a b.
MonadIO m =>
ThreadStorageMap a -> ThreadId -> (Maybe a -> (Maybe a, b)) -> m b
updateOnThread ThreadContextMap
threadContextMap ThreadId
tid ((Maybe ContextEntry -> (Maybe ContextEntry, ())) -> m ())
-> (Maybe ContextEntry -> (Maybe ContextEntry, ())) -> m ()
forall a b. (a -> b) -> a -> b
$ \Maybe ContextEntry
mentry ->
let !old :: ContextEntry
old = ContextEntry -> Maybe ContextEntry -> ContextEntry
forall a. a -> Maybe a -> a
fromMaybe ContextEntry
emptyEntry Maybe ContextEntry
mentry
!ctx' :: Context
ctx' = Context -> Context
f (ContextEntry -> Context
ceContext ContextEntry
old)
in (ContextEntry -> Maybe ContextEntry
forall a. a -> Maybe a
Just (ContextEntry
old {ceContext = ctx'}), ())
{-# INLINE adjustContextOnThread #-}
getActiveBaggage :: (MonadIO m) => m (Maybe Baggage)
getActiveBaggage :: forall (m :: * -> *). MonadIO m => m (Maybe Baggage)
getActiveBaggage = Context -> Maybe Baggage
lookupBaggage (Context -> Maybe Baggage) -> m Context -> m (Maybe Baggage)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> m Context
forall (m :: * -> *). MonadIO m => m Context
getContext
{-# INLINE getActiveBaggage #-}
setActiveBaggage :: (MonadIO m) => Baggage -> m ()
setActiveBaggage :: forall (m :: * -> *). MonadIO m => Baggage -> m ()
setActiveBaggage Baggage
b = (Context -> Context) -> m ()
forall (m :: * -> *). MonadIO m => (Context -> Context) -> m ()
adjustContext (Baggage -> Context -> Context
insertBaggage Baggage
b)
{-# INLINE setActiveBaggage #-}
clearActiveBaggage :: (MonadIO m) => m ()
clearActiveBaggage :: forall (m :: * -> *). MonadIO m => m ()
clearActiveBaggage = (Context -> Context) -> m ()
forall (m :: * -> *). MonadIO m => (Context -> Context) -> m ()
adjustContext Context -> Context
removeBaggage
{-# INLINE clearActiveBaggage #-}