{-# 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
  ) where
import OpenTelemetry.Context (Context, empty)
import Control.Concurrent
-- import Control.Concurrent.Async
import Control.Concurrent.Thread.Storage
import Control.Monad.IO.Class
import Data.Maybe (fromMaybe)
-- import Control.Monad
import System.IO.Unsafe
import Prelude hiding (lookup)
import Control.Monad (void)

type ThreadContextMap = ThreadStorageMap Context

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 #-}

-- | 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 :: m Context
getContext = Context -> Maybe Context -> Context
forall a. a -> Maybe a -> a
fromMaybe Context
empty (Maybe Context -> Context) -> m (Maybe Context) -> m Context
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> m (Maybe Context)
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 :: m (Maybe Context)
lookupContext = ThreadContextMap -> m (Maybe Context)
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 :: ThreadId -> m (Maybe Context)
lookupContextOnThread = ThreadContextMap -> ThreadId -> m (Maybe Context)
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 :: Context -> m (Maybe Context)
attachContext = ThreadContextMap -> Context -> m (Maybe Context)
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 :: ThreadId -> Context -> m (Maybe Context)
attachContextOnThread = ThreadContextMap -> ThreadId -> Context -> m (Maybe Context)
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.
--
-- @since 0.0.1.0
detachContext :: MonadIO m => m (Maybe Context)
detachContext :: m (Maybe Context)
detachContext = ThreadContextMap -> m (Maybe Context)
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.
--
-- @since 0.0.1.0
detachContextFromThread :: MonadIO m => ThreadId -> m (Maybe Context)
detachContextFromThread :: ThreadId -> m (Maybe Context)
detachContextFromThread = ThreadContextMap -> ThreadId -> m (Maybe Context)
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
--
-- @since 0.0.1.0
adjustContext :: MonadIO m => (Context -> Context) -> m ()
adjustContext :: (Context -> Context) -> m ()
adjustContext Context -> Context
f = IO () -> m ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ do
  Context
ctxt <- IO Context
forall (m :: * -> *). MonadIO m => m Context
getContext
  IO (Maybe Context) -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO (Maybe Context) -> IO ()) -> IO (Maybe Context) -> IO ()
forall a b. (a -> b) -> a -> b
$ Context -> IO (Maybe Context)
forall (m :: * -> *). MonadIO m => Context -> m (Maybe Context)
attachContext (Context -> IO (Maybe Context)) -> Context -> IO (Maybe Context)
forall a b. (a -> b) -> a -> b
$ Context -> Context
f Context
ctxt

-- | Alter the context
--
-- @since 0.0.1.0
adjustContextOnThread :: MonadIO m => ThreadId -> (Context -> Context) -> m ()
adjustContextOnThread :: ThreadId -> (Context -> Context) -> m ()
adjustContextOnThread = ThreadContextMap -> ThreadId -> (Context -> Context) -> m ()
forall (m :: * -> *) a.
MonadIO m =>
ThreadStorageMap a -> ThreadId -> (a -> a) -> m ()
adjustOnThread ThreadContextMap
threadContextMap