{-# LANGUAGE UnliftedFFITypes #-}

-----------------------------------------------------------------------------

-----------------------------------------------------------------------------

{- |
 Module      :  OpenTelemetry.Context.ThreadLocal
 Copyright   :  (c) Ian Duncan, 2021
 License     :  BSD-3
 Description :  State management for 'OpenTelemetry.Context.Context' on a per-thread basis.
 Maintainer  :  Ian Duncan
 Stability   :  experimental
 Portability :  non-portable (GHC extensions)

 Thread-local contexts may be attached as implicit state at a per-Haskell-thread
 level.

 This module uses a fair amount of GHC internals to enable performing
 lookups of context for any threads that are alive. Caution should be
 taken for consumers of this module to not retain ThreadId references
 indefinitely, as that could delay cleanup of thread-local state.

 Thread-local contexts have the following semantics:

 - A value 'attach'ed to a 'ThreadId' will remain alive at least as long
   as the 'ThreadId'.
 - A value may be detached from a 'ThreadId' via 'detach' by the
   library consumer without detriment.
 - No guarantees are made about when a value will be garbage-collected
   once all references to 'ThreadId' have been dropped. However, this simply
   means in practice that any unused contexts will cleaned up upon the next
   garbage collection and may not be actively freed when the program exits.
-}
module OpenTelemetry.Context.ThreadLocal (
  -- * Thread-local context
  getContext,
  lookupContext,
  attachContext,
  detachContext,
  adjustContext,

  -- ** Generalized thread-local context functions

  -- You should not use these without using some sort of specific cross-thread coordination mechanism,
  -- as there is no guarantee of what work the remote thread has done yet.
  lookupContextOnThread,
  attachContextOnThread,
  detachContextFromThread,
  adjustContextOnThread,

  -- ** Debugging tools
  threadContextMap,
) where

import Control.Concurrent
-- import Control.Concurrent.Async
import Control.Concurrent.Thread.Storage
-- import Control.Monad

import Control.Monad (void)
import Control.Monad.IO.Class
import Data.Maybe (fromMaybe)
import OpenTelemetry.Context (Context, empty)
import System.IO.Unsafe
import Prelude hiding (lookup)


type ThreadContextMap = ThreadStorageMap Context


{- | This is a global variable that is used to store the thread-local context map.
 It is not intended to be used directly for production purposes, but is exposed for debugging purposes.
-}
threadContextMap :: ThreadContextMap
threadContextMap :: ThreadContextMap
threadContextMap = forall a. IO a -> a
unsafePerformIO forall (m :: * -> *) a. MonadIO m => m (ThreadStorageMap a)
newThreadStorageMap
{-# NOINLINE threadContextMap #-}


{- | Retrieve a stored 'Context' for the current thread, or an empty context if none exists.

 Warning: this can easily cause disconnected traces if libraries don't explicitly set the
 context on forked threads.

 @since 0.0.1.0
-}
getContext :: MonadIO m => m Context
getContext :: forall (m :: * -> *). MonadIO m => m Context
getContext = forall a. a -> Maybe a -> a
fromMaybe Context
empty forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *). MonadIO m => m (Maybe Context)
lookupContext


{- | Retrieve a stored 'Context' for the current thread, if it exists.

 @since 0.0.1.0
-}
lookupContext :: MonadIO m => m (Maybe Context)
lookupContext :: forall (m :: * -> *). MonadIO m => m (Maybe Context)
lookupContext = forall (m :: * -> *) a.
MonadIO m =>
ThreadStorageMap a -> m (Maybe a)
lookup ThreadContextMap
threadContextMap


{- | Retrieve a stored 'Context' for the provided 'ThreadId', if it exists.

 @since 0.0.1.0
-}
lookupContextOnThread :: MonadIO m => ThreadId -> m (Maybe Context)
lookupContextOnThread :: forall (m :: * -> *). MonadIO m => ThreadId -> m (Maybe Context)
lookupContextOnThread = forall (m :: * -> *) a.
MonadIO m =>
ThreadStorageMap a -> ThreadId -> m (Maybe a)
lookupOnThread ThreadContextMap
threadContextMap


{- | Store a given 'Context' for the current thread, returning any context previously stored.

 @since 0.0.1.0
-}
attachContext :: MonadIO m => Context -> m (Maybe Context)
attachContext :: forall (m :: * -> *). MonadIO m => Context -> m (Maybe Context)
attachContext = forall (m :: * -> *) a.
MonadIO m =>
ThreadStorageMap a -> a -> m (Maybe a)
attach ThreadContextMap
threadContextMap


{- | Store a given 'Context' for the provided 'ThreadId', returning any context previously stored.

 @since 0.0.1.0
-}
attachContextOnThread :: MonadIO m => ThreadId -> Context -> m (Maybe Context)
attachContextOnThread :: forall (m :: * -> *).
MonadIO m =>
ThreadId -> Context -> m (Maybe Context)
attachContextOnThread = forall (m :: * -> *) a.
MonadIO m =>
ThreadStorageMap a -> ThreadId -> a -> m (Maybe a)
attachOnThread ThreadContextMap
threadContextMap


{- | Remove a stored 'Context' for the current thread, returning any context previously stored.

The detach functions don't generally need to be called manually, because finalizers will automatically
clean up contexts when a thread has completed and been garbage collected. If you are replacing a context
on a long-lived thread by detaching and attaching, use `adjustContext (const newContext)` instead to avoid
registering additional finalizer functions to be called on thread exit.

 @since 0.0.1.0
-}
detachContext :: MonadIO m => m (Maybe Context)
detachContext :: forall (m :: * -> *). MonadIO m => m (Maybe Context)
detachContext = forall (m :: * -> *) a.
MonadIO m =>
ThreadStorageMap a -> m (Maybe a)
detach ThreadContextMap
threadContextMap


{- | Remove a stored 'Context' for the provided 'ThreadId', returning any context previously stored.

The detach functions don't generally need to be called manually, because finalizers will automatically
clean up contexts when a thread has completed and been garbage collected. If you are replacing a context
on a long-lived thread by detaching and attaching, use `adjustContext (const newContext)` instead to avoid
registering additional finalizer functions to be called on thread exit.

 @since 0.0.1.0
-}
detachContextFromThread :: MonadIO m => ThreadId -> m (Maybe Context)
detachContextFromThread :: forall (m :: * -> *). MonadIO m => ThreadId -> m (Maybe Context)
detachContextFromThread = forall (m :: * -> *) a.
MonadIO m =>
ThreadStorageMap a -> ThreadId -> m (Maybe a)
detachFromThread ThreadContextMap
threadContextMap


{- | Alter the context on the current thread using the provided function.

If there is not a context associated with the current thread, the function will
be applied to an empty context and the result will be stored

 @since 0.0.1.0
-}
adjustContext :: MonadIO m => (Context -> Context) -> m ()
adjustContext :: forall (m :: * -> *). MonadIO m => (Context -> Context) -> m ()
adjustContext Context -> Context
f = forall (m :: * -> *) a b.
MonadIO m =>
ThreadStorageMap a -> (Maybe a -> (Maybe a, b)) -> m b
update ThreadContextMap
threadContextMap forall a b. (a -> b) -> a -> b
$ \Maybe Context
mctx ->
  (forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ Context -> Context
f forall a b. (a -> b) -> a -> b
$ forall a. a -> Maybe a -> a
fromMaybe Context
empty Maybe Context
mctx, ())


{- | Alter the context

If there is not a context associated with the provided thread, the function will
be applied to an empty context and the result will be stored

 @since 0.0.1.0
-}
adjustContextOnThread :: MonadIO m => ThreadId -> (Context -> Context) -> m ()
adjustContextOnThread :: forall (m :: * -> *).
MonadIO m =>
ThreadId -> (Context -> Context) -> m ()
adjustContextOnThread ThreadId
tid Context -> Context
f = forall (m :: * -> *) a b.
MonadIO m =>
ThreadStorageMap a -> ThreadId -> (Maybe a -> (Maybe a, b)) -> m b
updateOnThread ThreadContextMap
threadContextMap ThreadId
tid forall a b. (a -> b) -> a -> b
$ \Maybe Context
mctx ->
  (forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ Context -> Context
f forall a b. (a -> b) -> a -> b
$ forall a. a -> Maybe a -> a
fromMaybe Context
empty Maybe Context
mctx, ())