{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE OverloadedLists #-}
{-# LANGUAGE RecordWildCards #-}
module OpenTelemetry.Trace.Sampler (
Sampler (..),
SamplingResult (..),
SamplingDecision (..),
ParentBasedOptions (..),
shouldSample,
getDescription,
alwaysOn,
alwaysOff,
traceIdRatioBased,
parentBased,
parentBasedOptions,
alwaysRecord,
) where
import Data.Bits
import qualified Data.HashMap.Strict as H
import Data.Text (Text)
import qualified Data.Text.Lazy as TL
import Data.Text.Lazy.Builder (toLazyText)
import Data.Text.Lazy.Builder.RealFloat (realFloat)
import Data.Word (Word64, byteSwap64)
import GHC.ByteOrder (ByteOrder (..), targetByteOrder)
import OpenTelemetry.Attributes (toAttribute)
import OpenTelemetry.Context
import OpenTelemetry.Internal.Common.Types (InstrumentationLibrary)
import OpenTelemetry.Internal.Trace.Id (TraceId (..))
import OpenTelemetry.Internal.Trace.Types
import OpenTelemetry.Trace.Id
import OpenTelemetry.Trace.TraceState as TraceState
alwaysOn :: Sampler
alwaysOn :: Sampler
alwaysOn = Sampler
AlwaysOnSampler
alwaysOff :: Sampler
alwaysOff :: Sampler
alwaysOff = Sampler
AlwaysOffSampler
traceIdRatioBased :: Double -> Sampler
traceIdRatioBased :: Double -> Sampler
traceIdRatioBased Double
fraction =
Double -> Word64 -> Attribute -> Sampler
TraceIdRatioSampler Double
safeFraction Word64
traceIdUpperBound Attribute
sampleRate
where
safeFraction :: Double
safeFraction = Double -> Double -> Double
forall a. Ord a => a -> a -> a
max Double
0 (Double -> Double -> Double
forall a. Ord a => a -> a -> a
min Double
1 Double
fraction)
sampleRate :: Attribute
sampleRate =
if Double
safeFraction Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> Double
0
then Int -> Attribute
forall a. ToAttribute a => a -> Attribute
toAttribute ((Double -> Int
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
round (Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
safeFraction)) :: Int)
else Int -> Attribute
forall a. ToAttribute a => a -> Attribute
toAttribute (Int
0 :: Int)
traceIdUpperBound :: Word64
traceIdUpperBound = Double -> Word64
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
floor (Double
safeFraction Double -> Double -> Double
forall a. Num a => a -> a -> a
* Word64 -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral ((Word64
1 :: Word64) Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
`shiftL` Int
63)) :: Word64
parentBasedOptions
:: Sampler
-> ParentBasedOptions
parentBasedOptions :: Sampler -> ParentBasedOptions
parentBasedOptions Sampler
root =
ParentBasedOptions
{ rootSampler :: Sampler
rootSampler = Sampler
root
, remoteParentSampled :: Sampler
remoteParentSampled = Sampler
alwaysOn
, remoteParentNotSampled :: Sampler
remoteParentNotSampled = Sampler
alwaysOff
, localParentSampled :: Sampler
localParentSampled = Sampler
alwaysOn
, localParentNotSampled :: Sampler
localParentNotSampled = Sampler
alwaysOff
}
parentBased :: ParentBasedOptions -> Sampler
parentBased :: ParentBasedOptions -> Sampler
parentBased = ParentBasedOptions -> Sampler
ParentBasedSampler
alwaysRecord :: Sampler -> Sampler
alwaysRecord :: Sampler -> Sampler
alwaysRecord = Sampler -> Sampler
AlwaysRecordSampler
shouldSample :: Sampler -> Context -> TraceId -> Text -> SpanArguments -> InstrumentationLibrary -> IO SamplingDecision
shouldSample :: Sampler
-> Context
-> TraceId
-> Text
-> SpanArguments
-> InstrumentationLibrary
-> IO SamplingDecision
shouldSample Sampler
AlwaysOnSampler Context
ctxt TraceId
_tid Text
_name SpanArguments
_args InstrumentationLibrary
_scope =
let !ts :: TraceState
ts = Context -> TraceState
parentTraceState Context
ctxt
in SamplingDecision -> IO SamplingDecision
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (SamplingDecision -> IO SamplingDecision)
-> SamplingDecision -> IO SamplingDecision
forall a b. (a -> b) -> a -> b
$! SamplingResult -> AttributeMap -> TraceState -> SamplingDecision
SamplingDecision SamplingResult
RecordAndSample AttributeMap
forall k v. HashMap k v
H.empty TraceState
ts
shouldSample Sampler
AlwaysOffSampler Context
ctxt TraceId
_tid Text
_name SpanArguments
_args InstrumentationLibrary
_scope =
let !ts :: TraceState
ts = Context -> TraceState
parentTraceState Context
ctxt
in SamplingDecision -> IO SamplingDecision
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (SamplingDecision -> IO SamplingDecision)
-> SamplingDecision -> IO SamplingDecision
forall a b. (a -> b) -> a -> b
$! SamplingResult -> AttributeMap -> TraceState -> SamplingDecision
SamplingDecision SamplingResult
Drop AttributeMap
forall k v. HashMap k v
H.empty TraceState
ts
shouldSample Sampler
sampler Context
ctxt TraceId
tid Text
name SpanArguments
args InstrumentationLibrary
scope =
Sampler
-> Context
-> TraceId
-> Text
-> SpanArguments
-> InstrumentationLibrary
-> IO SamplingDecision
shouldSampleComplex Sampler
sampler Context
ctxt TraceId
tid Text
name SpanArguments
args InstrumentationLibrary
scope
{-# INLINE shouldSample #-}
shouldSampleComplex :: Sampler -> Context -> TraceId -> Text -> SpanArguments -> InstrumentationLibrary -> IO SamplingDecision
shouldSampleComplex :: Sampler
-> Context
-> TraceId
-> Text
-> SpanArguments
-> InstrumentationLibrary
-> IO SamplingDecision
shouldSampleComplex (TraceIdRatioSampler Double
frac Word64
upperBound Attribute
sampleRateAttr) Context
ctxt TraceId
tid Text
_name SpanArguments
_args InstrumentationLibrary
_scope =
let !ts :: TraceState
ts = Context -> TraceState
parentTraceState Context
ctxt
in if Double
frac Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
>= Double
1
then SamplingDecision -> IO SamplingDecision
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (SamplingDecision -> IO SamplingDecision)
-> SamplingDecision -> IO SamplingDecision
forall a b. (a -> b) -> a -> b
$! SamplingResult -> AttributeMap -> TraceState -> SamplingDecision
SamplingDecision SamplingResult
RecordAndSample (Text -> Attribute -> AttributeMap
forall k v. Hashable k => k -> v -> HashMap k v
H.singleton Text
"sampleRate" Attribute
sampleRateAttr) TraceState
ts
else
let !(TraceId Word64
_ Word64
lo) = TraceId
tid
!loBE :: Word64
loBE = case ByteOrder
targetByteOrder of
ByteOrder
BigEndian -> Word64
lo
ByteOrder
LittleEndian -> Word64 -> Word64
byteSwap64 Word64
lo
!x :: Word64
x = Word64
loBE Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
`shiftR` Int
1
in if Word64
x Word64 -> Word64 -> Bool
forall a. Ord a => a -> a -> Bool
< Word64
upperBound
then SamplingDecision -> IO SamplingDecision
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (SamplingDecision -> IO SamplingDecision)
-> SamplingDecision -> IO SamplingDecision
forall a b. (a -> b) -> a -> b
$! SamplingResult -> AttributeMap -> TraceState -> SamplingDecision
SamplingDecision SamplingResult
RecordAndSample (Text -> Attribute -> AttributeMap
forall k v. Hashable k => k -> v -> HashMap k v
H.singleton Text
"sampleRate" Attribute
sampleRateAttr) TraceState
ts
else SamplingDecision -> IO SamplingDecision
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (SamplingDecision -> IO SamplingDecision)
-> SamplingDecision -> IO SamplingDecision
forall a b. (a -> b) -> a -> b
$! SamplingResult -> AttributeMap -> TraceState -> SamplingDecision
SamplingDecision SamplingResult
Drop AttributeMap
forall k v. HashMap k v
H.empty TraceState
ts
shouldSampleComplex (ParentBasedSampler ParentBasedOptions {Sampler
rootSampler :: ParentBasedOptions -> Sampler
remoteParentSampled :: ParentBasedOptions -> Sampler
remoteParentNotSampled :: ParentBasedOptions -> Sampler
localParentSampled :: ParentBasedOptions -> Sampler
localParentNotSampled :: ParentBasedOptions -> Sampler
rootSampler :: Sampler
remoteParentSampled :: Sampler
remoteParentNotSampled :: Sampler
localParentSampled :: Sampler
localParentNotSampled :: Sampler
..}) Context
ctxt TraceId
tid Text
name SpanArguments
args InstrumentationLibrary
scope =
case Context -> Maybe SpanContext
parentSpanContext Context
ctxt of
Maybe SpanContext
Nothing -> Sampler
-> Context
-> TraceId
-> Text
-> SpanArguments
-> InstrumentationLibrary
-> IO SamplingDecision
shouldSample Sampler
rootSampler Context
ctxt TraceId
tid Text
name SpanArguments
args InstrumentationLibrary
scope
Just SpanContext
sc ->
if SpanContext -> Bool
OpenTelemetry.Internal.Trace.Types.isRemote SpanContext
sc
then
if TraceFlags -> Bool
isSampled (SpanContext -> TraceFlags
traceFlags SpanContext
sc)
then Sampler
-> Context
-> TraceId
-> Text
-> SpanArguments
-> InstrumentationLibrary
-> IO SamplingDecision
shouldSample Sampler
remoteParentSampled Context
ctxt TraceId
tid Text
name SpanArguments
args InstrumentationLibrary
scope
else Sampler
-> Context
-> TraceId
-> Text
-> SpanArguments
-> InstrumentationLibrary
-> IO SamplingDecision
shouldSample Sampler
remoteParentNotSampled Context
ctxt TraceId
tid Text
name SpanArguments
args InstrumentationLibrary
scope
else
if TraceFlags -> Bool
isSampled (SpanContext -> TraceFlags
traceFlags SpanContext
sc)
then Sampler
-> Context
-> TraceId
-> Text
-> SpanArguments
-> InstrumentationLibrary
-> IO SamplingDecision
shouldSample Sampler
localParentSampled Context
ctxt TraceId
tid Text
name SpanArguments
args InstrumentationLibrary
scope
else Sampler
-> Context
-> TraceId
-> Text
-> SpanArguments
-> InstrumentationLibrary
-> IO SamplingDecision
shouldSample Sampler
localParentNotSampled Context
ctxt TraceId
tid Text
name SpanArguments
args InstrumentationLibrary
scope
shouldSampleComplex (AlwaysRecordSampler Sampler
inner) Context
ctxt TraceId
tid Text
name SpanArguments
args InstrumentationLibrary
scope = do
decision <- Sampler
-> Context
-> TraceId
-> Text
-> SpanArguments
-> InstrumentationLibrary
-> IO SamplingDecision
shouldSample Sampler
inner Context
ctxt TraceId
tid Text
name SpanArguments
args InstrumentationLibrary
scope
let !outcome' = case SamplingDecision -> SamplingResult
samplingOutcome SamplingDecision
decision of
SamplingResult
Drop -> SamplingResult
RecordOnly
SamplingResult
other -> SamplingResult
other
pure $! decision {samplingOutcome = outcome'}
shouldSampleComplex (CustomSampler Text
_ Context
-> TraceId
-> Text
-> SpanArguments
-> InstrumentationLibrary
-> IO SamplingDecision
f) Context
ctxt TraceId
tid Text
name SpanArguments
args InstrumentationLibrary
scope = Context
-> TraceId
-> Text
-> SpanArguments
-> InstrumentationLibrary
-> IO SamplingDecision
f Context
ctxt TraceId
tid Text
name SpanArguments
args InstrumentationLibrary
scope
shouldSampleComplex Sampler
AlwaysOnSampler Context
ctxt TraceId
tid Text
name SpanArguments
args InstrumentationLibrary
scope = Sampler
-> Context
-> TraceId
-> Text
-> SpanArguments
-> InstrumentationLibrary
-> IO SamplingDecision
shouldSample Sampler
AlwaysOnSampler Context
ctxt TraceId
tid Text
name SpanArguments
args InstrumentationLibrary
scope
shouldSampleComplex Sampler
AlwaysOffSampler Context
ctxt TraceId
tid Text
name SpanArguments
args InstrumentationLibrary
scope = Sampler
-> Context
-> TraceId
-> Text
-> SpanArguments
-> InstrumentationLibrary
-> IO SamplingDecision
shouldSample Sampler
AlwaysOffSampler Context
ctxt TraceId
tid Text
name SpanArguments
args InstrumentationLibrary
scope
{-# NOINLINE shouldSampleComplex #-}
getDescription :: Sampler -> Text
getDescription :: Sampler -> Text
getDescription Sampler
AlwaysOnSampler = Text
"AlwaysOnSampler"
getDescription Sampler
AlwaysOffSampler = Text
"AlwaysOffSampler"
getDescription (TraceIdRatioSampler Double
frac Word64
_ Attribute
_) =
Text
"TraceIdRatioBased{" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> LazyText -> Text
TL.toStrict (Builder -> LazyText
toLazyText (Double -> Builder
forall a. RealFloat a => a -> Builder
realFloat Double
frac)) Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"}"
getDescription (ParentBasedSampler ParentBasedOptions {Sampler
rootSampler :: ParentBasedOptions -> Sampler
remoteParentSampled :: ParentBasedOptions -> Sampler
remoteParentNotSampled :: ParentBasedOptions -> Sampler
localParentSampled :: ParentBasedOptions -> Sampler
localParentNotSampled :: ParentBasedOptions -> Sampler
rootSampler :: Sampler
remoteParentSampled :: Sampler
remoteParentNotSampled :: Sampler
localParentSampled :: Sampler
localParentNotSampled :: Sampler
..}) =
Text
"ParentBased{root="
Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Sampler -> Text
getDescription Sampler
rootSampler
Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
", remoteParentSampled="
Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Sampler -> Text
getDescription Sampler
remoteParentSampled
Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
", remoteParentNotSampled="
Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Sampler -> Text
getDescription Sampler
remoteParentNotSampled
Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
", localParentSampled="
Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Sampler -> Text
getDescription Sampler
localParentSampled
Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
", localParentNotSampled="
Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Sampler -> Text
getDescription Sampler
localParentNotSampled
Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"}"
getDescription (AlwaysRecordSampler Sampler
inner) =
Text
"AlwaysRecord{" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Sampler -> Text
getDescription Sampler
inner Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"}"
getDescription (CustomSampler Text
desc Context
-> TraceId
-> Text
-> SpanArguments
-> InstrumentationLibrary
-> IO SamplingDecision
_) = Text
desc
parentSpanContext :: Context -> Maybe SpanContext
parentSpanContext :: Context -> Maybe SpanContext
parentSpanContext Context
ctxt = case Context -> Maybe Span
lookupSpan Context
ctxt of
Maybe Span
Nothing -> Maybe SpanContext
forall a. Maybe a
Nothing
Just (Span ImmutableSpan
imm) -> SpanContext -> Maybe SpanContext
forall a. a -> Maybe a
Just (ImmutableSpan -> SpanContext
spanContext ImmutableSpan
imm)
Just (FrozenSpan SpanContext
sc) -> SpanContext -> Maybe SpanContext
forall a. a -> Maybe a
Just SpanContext
sc
Just (Dropped SpanContext
sc) -> SpanContext -> Maybe SpanContext
forall a. a -> Maybe a
Just SpanContext
sc
{-# INLINE parentSpanContext #-}
parentTraceState :: Context -> TraceState
parentTraceState :: Context -> TraceState
parentTraceState Context
ctxt = case Context -> Maybe SpanContext
parentSpanContext Context
ctxt of
Maybe SpanContext
Nothing -> TraceState
TraceState.empty
Just SpanContext
sc -> SpanContext -> TraceState
traceState SpanContext
sc
{-# INLINE parentTraceState #-}