{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE DuplicateRecordFields #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE StrictData #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE UnboxedTuples #-}
{-# LANGUAGE UnliftedFFITypes #-}

{- |
 Module      :  OpenTelemetry.Util
 Copyright   :  (c) Ian Duncan, 2021
 License     :  BSD-3
 Description :  Convenience functions to simplify common instrumentation needs.
 Maintainer  :  Ian Duncan
 Stability   :  experimental
 Portability :  non-portable (GHC extensions)
-}
module OpenTelemetry.Util (
  constructorName,
  HasConstructor,
  getThreadId,
  bracketError,

  -- * Lock-free IORef modification
  casModifyIORef_,
  casReadModifyIORef_,

  -- * Data structures
  AppendOnlyBoundedCollection,
  emptyAppendOnlyBoundedCollection,
  appendToBoundedCollection,
  appendOnlyBoundedCollectionSize,
  appendOnlyBoundedCollectionValues,
  appendOnlyBoundedCollectionDroppedElementCount,

  -- * Vectors
  chunksOfV,
) where

import Control.Exception (SomeException)
import qualified Control.Exception as EUnsafe
import Control.Monad.IO.Unlift
import Data.IORef (IORef)
import Data.Kind
import qualified Data.Vector as V
import Foreign.C (CInt (..))
import GHC.Base (Addr#)
import GHC.Conc (ThreadId (ThreadId))
import GHC.Exts (unsafeCoerce#)
import qualified GHC.Exts as Exts
import GHC.Generics
import GHC.IO (IO (IO))
import GHC.IORef (IORef (IORef))
import GHC.STRef (STRef (STRef))


{- | Useful for annotating which constructor in an ADT was chosen

 @since 0.1.0.0
-}
constructorName :: (HasConstructor (Rep a), Generic a) => a -> String
constructorName :: forall a. (HasConstructor (Rep a), Generic a) => a -> String
constructorName = Rep a (ZonkAny 0) -> String
forall x. Rep a x -> String
forall (f :: * -> *) x. HasConstructor f => f x -> String
genericConstrName (Rep a (ZonkAny 0) -> String)
-> (a -> Rep a (ZonkAny 0)) -> a -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Rep a (ZonkAny 0)
forall x. a -> Rep a x
forall a x. Generic a => a -> Rep a x
from


{- | Detect a constructor from any datatype which derives 'Generic'

@since 0.0.1.0
-}
class HasConstructor (f :: Type -> Type) where
  genericConstrName :: f x -> String


instance (HasConstructor f) => HasConstructor (D1 c f) where
  genericConstrName :: forall x. D1 c f x -> String
genericConstrName (M1 f x
x) = f x -> String
forall x. f x -> String
forall (f :: * -> *) x. HasConstructor f => f x -> String
genericConstrName f x
x


instance (HasConstructor x, HasConstructor y) => HasConstructor (x :+: y) where
  genericConstrName :: forall x. (:+:) x y x -> String
genericConstrName (L1 x x
l) = x x -> String
forall x. x x -> String
forall (f :: * -> *) x. HasConstructor f => f x -> String
genericConstrName x x
l
  genericConstrName (R1 y x
r) = y x -> String
forall x. y x -> String
forall (f :: * -> *) x. HasConstructor f => f x -> String
genericConstrName y x
r


instance (Constructor c) => HasConstructor (C1 c f) where
  genericConstrName :: forall x. C1 c f x -> String
genericConstrName = M1 C c f x -> String
forall {k} (c :: k) k1 (t :: k -> (k1 -> *) -> k1 -> *)
       (f :: k1 -> *) (a :: k1).
Constructor c =>
t c f a -> String
forall k1 (t :: Meta -> (k1 -> *) -> k1 -> *) (f :: k1 -> *)
       (a :: k1).
t c f a -> String
conName


foreign import ccall unsafe "rts_getThreadId" c_getThreadId :: Addr# -> CInt


{- | Get an int representation of a thread id

@since 0.0.1.0
-}
getThreadId :: ThreadId -> Int
getThreadId :: ThreadId -> Int
getThreadId (ThreadId ThreadId#
tid#) = CInt -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (CInt -> Int) -> CInt -> Int
forall a b. (a -> b) -> a -> b
$ Addr# -> CInt
c_getThreadId (ThreadId# -> Addr#
forall a b. a -> b
unsafeCoerce# ThreadId#
tid#)
{-# INLINE getThreadId #-}


{- Note [NOINLINE on CAS functions]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
These functions MUST be NOINLINE. When inlined, GHC's Core-to-STG pass may
wrap the "new" constructor (the result of @f old@) in an updatable thunk rather
than allocating it directly as a constructor. This happens when the modification
function produces a record with a strict field whose value is a locally-allocated
closure (e.g. a DList function composition @dl . (x :)@ in
'AppendOnlyBoundedCollection'). The Core-to-STG pass conservatively inserts a
@case@ to force the strict field, turning the entire constructor into a thunk.

The thunk pointer has tag 0 (untagged). @casMutVar#@ writes this untagged
thunk to the MutVar. On the next read, the thunk is evaluated; it becomes an
indirection to the actual constructor at a different heap address (tag 1). The
CAS then compares the evaluated constructor pointer (tag 1) against the MutVar
contents (the indirection cell, tag 0). These are different addresses, so the
CAS fails *permanently*: every retry reads the indirection, evaluates to the
constructor, and mismatches again. Result: infinite spin-loop.

With NOINLINE, @let !new = f old@ forces the thunk inside the CAS function
before passing the evaluated (tagged) constructor pointer to @casMutVar#@.
The MutVar always contains a properly tagged constructor pointer, so the CAS
comparison works correctly.

Affects all GHC versions tested (9.4 through 9.12), both aarch64 and x86_64.

Note: we deliberately omit @yield#@ on the CAS failure path. NOINLINE already
makes the function call a GC safe point, and each retry re-evaluates @f old@
which allocates, providing another safe point. @yield#@ was measured to add
significant overhead even on the uncontended success path (likely because GHC
cannot prove it is dead code and its presence inhibits tail-call / loop
optimisation).
-}

{- | CAS-based strict IORef modification that avoids the closure and pair
allocation of 'atomicModifyIORef''.

@atomicModifyIORef' ref (\old -> (f old, ()))@ allocates a closure capturing
the modification function, a @(new, ())@ pair, and (in the GHC RTS) a thunk
indirection that is CAS'd into the MutVar.

This function instead reads the current value, applies @f@ strictly, and
performs a compare-and-swap. On success (the common, uncontended case), zero
intermediate heap objects are allocated beyond the new value itself. On CAS
failure (concurrent modification), it retries. Safe because @f@ is pure.

Use for hot-path span operations (addAttribute, addEvent, setStatus, etc.)
where the IORef is rarely contended and the modification is a cheap record
update.

@since 0.0.1.0
-}
casModifyIORef_ :: IORef a -> (a -> a) -> IO ()
casModifyIORef_ :: forall a. IORef a -> (a -> a) -> IO ()
casModifyIORef_ (IORef (STRef MutVar# RealWorld a
ref#)) a -> a
f = (State# RealWorld -> (# State# RealWorld, () #)) -> IO ()
forall a. (State# RealWorld -> (# State# RealWorld, a #)) -> IO a
IO State# RealWorld -> (# State# RealWorld, () #)
go#
  where
    go# :: State# RealWorld -> (# State# RealWorld, () #)
go# State# RealWorld
s0# =
      case MutVar# RealWorld a
-> State# RealWorld -> (# State# RealWorld, a #)
forall d a. MutVar# d a -> State# d -> (# State# d, a #)
Exts.readMutVar# MutVar# RealWorld a
ref# State# RealWorld
s0# of
        (# State# RealWorld
s1#, a
old #) ->
          let !new :: a
new = a -> a
f a
old
          in case MutVar# RealWorld a
-> a -> a -> State# RealWorld -> (# State# RealWorld, Int#, a #)
forall d a.
MutVar# d a -> a -> a -> State# d -> (# State# d, Int#, a #)
Exts.casMutVar# MutVar# RealWorld a
ref# a
old a
new State# RealWorld
s1# of
               (# State# RealWorld
s2#, Int#
0#, a
_ #) -> (# State# RealWorld
s2#, () #)
               (# State# RealWorld
s2#, Int#
_, a
_ #) -> State# RealWorld -> (# State# RealWorld, () #)
go# State# RealWorld
s2#
{-# NOINLINE casModifyIORef_ #-}


-- NOINLINE is load-bearing: see Note [NOINLINE on CAS functions]

{- | CAS-based IORef modification that also reads the old value before the swap.

Performs a strict read, applies @f@ to decide both the new value and a
pre-swap result, then CAS's the new value in. Returns the old value (before
modification) on success. Retries on CAS failure.

Used by 'endSpan' where we need to atomically set @spanEnd@ and also read the
(unchanged) Tracer field to obtain the processor vector, all without
navigating the Tracer inside an @atomicModifyIORef'@ closure.

@since 0.0.1.0
-}
casReadModifyIORef_ :: IORef a -> (a -> a) -> IO a
casReadModifyIORef_ :: forall a. IORef a -> (a -> a) -> IO a
casReadModifyIORef_ (IORef (STRef MutVar# RealWorld a
ref#)) a -> a
f = (State# RealWorld -> (# State# RealWorld, a #)) -> IO a
forall a. (State# RealWorld -> (# State# RealWorld, a #)) -> IO a
IO State# RealWorld -> (# State# RealWorld, a #)
go#
  where
    go# :: State# RealWorld -> (# State# RealWorld, a #)
go# State# RealWorld
s0# =
      case MutVar# RealWorld a
-> State# RealWorld -> (# State# RealWorld, a #)
forall d a. MutVar# d a -> State# d -> (# State# d, a #)
Exts.readMutVar# MutVar# RealWorld a
ref# State# RealWorld
s0# of
        (# State# RealWorld
s1#, a
old #) ->
          let !new :: a
new = a -> a
f a
old
          in case MutVar# RealWorld a
-> a -> a -> State# RealWorld -> (# State# RealWorld, Int#, a #)
forall d a.
MutVar# d a -> a -> a -> State# d -> (# State# d, Int#, a #)
Exts.casMutVar# MutVar# RealWorld a
ref# a
old a
new State# RealWorld
s1# of
               (# State# RealWorld
s2#, Int#
0#, a
_ #) -> (# State# RealWorld
s2#, a
old #)
               (# State# RealWorld
s2#, Int#
_, a
_ #) -> State# RealWorld -> (# State# RealWorld, a #)
go# State# RealWorld
s2#
{-# NOINLINE casReadModifyIORef_ #-}


-- NOINLINE is load-bearing: see Note [NOINLINE on CAS functions]

{- | Bounded append-only collection.

Two constructors: 'EmptyBounded' carries only the capacity (2 words: info
pointer + unboxed Int), avoiding all allocation for the common case of spans
with 0 events or 0 links. 'BoundedCollection' uses a difference list for
O(1) pure append and O(n) materialization at export time.

Safe with 'atomicModifyIORef'' because all operations are pure.

@since 0.0.1.0
-}
data AppendOnlyBoundedCollection a
  = EmptyBounded
      {-# UNPACK #-} !Int
  | BoundedCollection
      !([a] -> [a])
      {-# UNPACK #-} !Int
      {-# UNPACK #-} !Int
      {-# UNPACK #-} !Int


instance forall a. (Show a) => Show (AppendOnlyBoundedCollection a) where
  showsPrec :: Int -> AppendOnlyBoundedCollection a -> ShowS
showsPrec Int
d AppendOnlyBoundedCollection a
c =
    let vec :: Vector a
vec = AppendOnlyBoundedCollection a -> Vector a
forall a. AppendOnlyBoundedCollection a -> Vector a
appendOnlyBoundedCollectionValues AppendOnlyBoundedCollection a
c
    in Bool -> ShowS -> ShowS
showParen (Int
d Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
10) (ShowS -> ShowS) -> ShowS -> ShowS
forall a b. (a -> b) -> a -> b
$
         String -> ShowS
showString String
"AppendOnlyBoundedCollection {collection = "
           ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector a -> ShowS
forall a. Show a => a -> ShowS
shows Vector a
vec
           ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> ShowS
showString String
", maxSize = "
           ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> ShowS
forall a. Show a => a -> ShowS
shows (AppendOnlyBoundedCollection a -> Int
forall a. AppendOnlyBoundedCollection a -> Int
appendOnlyBoundedCollectionMaxSize AppendOnlyBoundedCollection a
c)
           ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> ShowS
showString String
", dropped = "
           ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> ShowS
forall a. Show a => a -> ShowS
shows (AppendOnlyBoundedCollection a -> Int
forall a. AppendOnlyBoundedCollection a -> Int
appendOnlyBoundedCollectionDroppedElementCount AppendOnlyBoundedCollection a
c)
           ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> ShowS
showString String
"}"


{- | Initialize a bounded collection that admits a maximum size

@since 0.0.1.0
-}
emptyAppendOnlyBoundedCollection
  :: Int
  -- ^ Maximum size
  -> AppendOnlyBoundedCollection a
emptyAppendOnlyBoundedCollection :: forall a. Int -> AppendOnlyBoundedCollection a
emptyAppendOnlyBoundedCollection !Int
s = Int -> AppendOnlyBoundedCollection a
forall a. Int -> AppendOnlyBoundedCollection a
EmptyBounded Int
s
{-# INLINE emptyAppendOnlyBoundedCollection #-}


{- | O(n). Materializes the difference list into a 'V.Vector' via 'V.fromListN'.
Called once per span at export time.

@since 0.0.1.0
-}
appendOnlyBoundedCollectionValues :: AppendOnlyBoundedCollection a -> V.Vector a
appendOnlyBoundedCollectionValues :: forall a. AppendOnlyBoundedCollection a -> Vector a
appendOnlyBoundedCollectionValues (EmptyBounded Int
_) = Vector a
forall a. Vector a
V.empty
appendOnlyBoundedCollectionValues (BoundedCollection [a] -> [a]
dl Int
sz Int
_ Int
_) =
  Int -> [a] -> Vector a
forall a. Int -> [a] -> Vector a
V.fromListN Int
sz ([a] -> [a]
dl [])
{-# INLINE appendOnlyBoundedCollectionValues #-}


-- | @since 0.0.1.0
appendOnlyBoundedCollectionSize :: AppendOnlyBoundedCollection a -> Int
appendOnlyBoundedCollectionSize :: forall a. AppendOnlyBoundedCollection a -> Int
appendOnlyBoundedCollectionSize (EmptyBounded Int
_) = Int
0
appendOnlyBoundedCollectionSize (BoundedCollection [a] -> [a]
_ Int
sz Int
_ Int
_) = Int
sz
{-# INLINE appendOnlyBoundedCollectionSize #-}


appendOnlyBoundedCollectionMaxSize :: AppendOnlyBoundedCollection a -> Int
appendOnlyBoundedCollectionMaxSize :: forall a. AppendOnlyBoundedCollection a -> Int
appendOnlyBoundedCollectionMaxSize (EmptyBounded Int
ms) = Int
ms
appendOnlyBoundedCollectionMaxSize (BoundedCollection [a] -> [a]
_ Int
_ Int
ms Int
_) = Int
ms
{-# INLINE appendOnlyBoundedCollectionMaxSize #-}


-- | @since 0.0.1.0
appendOnlyBoundedCollectionDroppedElementCount :: AppendOnlyBoundedCollection a -> Int
appendOnlyBoundedCollectionDroppedElementCount :: forall a. AppendOnlyBoundedCollection a -> Int
appendOnlyBoundedCollectionDroppedElementCount (EmptyBounded Int
_) = Int
0
appendOnlyBoundedCollectionDroppedElementCount (BoundedCollection [a] -> [a]
_ Int
_ Int
_ Int
d) = Int
d
{-# INLINE appendOnlyBoundedCollectionDroppedElementCount #-}


{- | Append an element. O(1): transitions from 'EmptyBounded' to
'BoundedCollection' on first append, then difference-list composition.
Returns the collection unchanged (with incremented drop count) when full.

@since 0.0.1.0
-}
appendToBoundedCollection :: AppendOnlyBoundedCollection a -> a -> AppendOnlyBoundedCollection a
appendToBoundedCollection :: forall a.
AppendOnlyBoundedCollection a -> a -> AppendOnlyBoundedCollection a
appendToBoundedCollection (EmptyBounded Int
ms) a
x
  | Int
ms Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0 = ([a] -> [a]) -> Int -> Int -> Int -> AppendOnlyBoundedCollection a
forall a.
([a] -> [a]) -> Int -> Int -> Int -> AppendOnlyBoundedCollection a
BoundedCollection [a] -> [a]
forall a. a -> a
id Int
0 Int
ms Int
1
  | Bool
otherwise = ([a] -> [a]) -> Int -> Int -> Int -> AppendOnlyBoundedCollection a
forall a.
([a] -> [a]) -> Int -> Int -> Int -> AppendOnlyBoundedCollection a
BoundedCollection (a
x a -> [a] -> [a]
forall a. a -> [a] -> [a]
:) Int
1 Int
ms Int
0
appendToBoundedCollection (BoundedCollection [a] -> [a]
dl Int
sz Int
ms Int
d) a
x
  | Int
sz Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
ms = ([a] -> [a]) -> Int -> Int -> Int -> AppendOnlyBoundedCollection a
forall a.
([a] -> [a]) -> Int -> Int -> Int -> AppendOnlyBoundedCollection a
BoundedCollection [a] -> [a]
dl Int
sz Int
ms (Int
d Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
  | Bool
otherwise = ([a] -> [a]) -> Int -> Int -> Int -> AppendOnlyBoundedCollection a
forall a.
([a] -> [a]) -> Int -> Int -> Int -> AppendOnlyBoundedCollection a
BoundedCollection ([a] -> [a]
dl ([a] -> [a]) -> ([a] -> [a]) -> [a] -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a
x a -> [a] -> [a]
forall a. a -> [a] -> [a]
:)) (Int
sz Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int
ms Int
d
{-# INLINE appendToBoundedCollection #-}


{- | Split a vector into chunks of at most @n@ elements. Used by batch processors.

@since 0.4.0.0
-}
chunksOfV :: Int -> V.Vector a -> [V.Vector a]
chunksOfV :: forall a. Int -> Vector a -> [Vector a]
chunksOfV Int
n Vector a
v
  | Vector a -> Bool
forall a. Vector a -> Bool
V.null Vector a
v = []
  | Bool
otherwise =
      let (Vector a
chunk, Vector a
rest) = Int -> Vector a -> (Vector a, Vector a)
forall a. Int -> Vector a -> (Vector a, Vector a)
V.splitAt Int
n Vector a
v
      in Vector a
chunk Vector a -> [Vector a] -> [Vector a]
forall a. a -> [a] -> [a]
: Int -> Vector a -> [Vector a]
forall a. Int -> Vector a -> [Vector a]
chunksOfV Int
n Vector a
rest


{- | Like 'Context.Exception.bracket', but provides the @after@ function with information about
 uncaught exceptions.

 @since 0.1.0.0
-}
bracketError :: (MonadUnliftIO m) => m a -> (Maybe SomeException -> a -> m b) -> (a -> m c) -> m c
{-# INLINEABLE bracketError #-}
bracketError :: forall (m :: * -> *) a b c.
MonadUnliftIO m =>
m a -> (Maybe SomeException -> a -> m b) -> (a -> m c) -> m c
bracketError m a
before Maybe SomeException -> a -> m b
after a -> m c
thing = ((forall a. m a -> IO a) -> IO c) -> m c
forall b. ((forall a. m a -> IO a) -> IO b) -> m b
forall (m :: * -> *) b.
MonadUnliftIO m =>
((forall a. m a -> IO a) -> IO b) -> m b
withRunInIO (((forall a. m a -> IO a) -> IO c) -> m c)
-> ((forall a. m a -> IO a) -> IO c) -> m c
forall a b. (a -> b) -> a -> b
$ \forall a. m a -> IO a
run -> ((forall a. IO a -> IO a) -> IO c) -> IO c
forall b. ((forall a. IO a -> IO a) -> IO b) -> IO b
EUnsafe.mask (((forall a. IO a -> IO a) -> IO c) -> IO c)
-> ((forall a. IO a -> IO a) -> IO c) -> IO c
forall a b. (a -> b) -> a -> b
$ \forall a. IO a -> IO a
restore -> do
  x <- m a -> IO a
forall a. m a -> IO a
run m a
before
  y <-
    restore (run $ thing x) `EUnsafe.catch` \(SomeException
e1 :: SomeException) -> do
      _ :: Either SomeException b <-
        IO b -> IO (Either SomeException b)
forall e a. Exception e => IO a -> IO (Either e a)
EUnsafe.try (IO b -> IO (Either SomeException b))
-> IO b -> IO (Either SomeException b)
forall a b. (a -> b) -> a -> b
$ IO b -> IO b
forall a. IO a -> IO a
EUnsafe.uninterruptibleMask_ (IO b -> IO b) -> IO b -> IO b
forall a b. (a -> b) -> a -> b
$ m b -> IO b
forall a. m a -> IO a
run (m b -> IO b) -> m b -> IO b
forall a b. (a -> b) -> a -> b
$ Maybe SomeException -> a -> m b
after (SomeException -> Maybe SomeException
forall a. a -> Maybe a
Just SomeException
e1) a
x
      EUnsafe.throwIO e1
  _ <- EUnsafe.uninterruptibleMask_ $ run $ after Nothing x
  return y