{-# LANGUAGE OverloadedLists #-}
module OpenTelemetry.Trace.Sampler (
Sampler (..),
SamplingResult (..),
parentBased,
parentBasedOptions,
ParentBasedOptions (..),
traceIdRatioBased,
alwaysOn,
alwaysOff,
) where
import Data.Binary.Get
import Data.Bits
import qualified Data.ByteString as B
import qualified Data.ByteString.Lazy as L
import Data.Text
import Data.Word (Word64)
import OpenTelemetry.Attributes (toAttribute)
import OpenTelemetry.Context
import OpenTelemetry.Internal.Trace.Types
import OpenTelemetry.Trace.Id
import OpenTelemetry.Trace.TraceState as TraceState
alwaysOn :: Sampler
alwaysOn :: Sampler
alwaysOn =
Sampler
{ getDescription :: Text
getDescription = Text
"AlwaysOnSampler"
, shouldSample :: Context
-> TraceId
-> Text
-> SpanArguments
-> IO (SamplingResult, HashMap Text Attribute, TraceState)
shouldSample = \Context
ctxt TraceId
_ Text
_ SpanArguments
_ -> do
Maybe SpanContext
mspanCtxt <- Maybe (IO SpanContext) -> IO (Maybe SpanContext)
forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
forall (m :: * -> *) a. Monad m => Maybe (m a) -> m (Maybe a)
sequence (Span -> IO SpanContext
forall (m :: * -> *). MonadIO m => Span -> m SpanContext
getSpanContext (Span -> IO SpanContext) -> Maybe Span -> Maybe (IO SpanContext)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Context -> Maybe Span
lookupSpan Context
ctxt)
(SamplingResult, HashMap Text Attribute, TraceState)
-> IO (SamplingResult, HashMap Text Attribute, TraceState)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (SamplingResult
RecordAndSample, [], TraceState
-> (SpanContext -> TraceState) -> Maybe SpanContext -> TraceState
forall b a. b -> (a -> b) -> Maybe a -> b
maybe TraceState
TraceState.empty SpanContext -> TraceState
traceState Maybe SpanContext
mspanCtxt)
}
alwaysOff :: Sampler
alwaysOff :: Sampler
alwaysOff =
Sampler
{ getDescription :: Text
getDescription = Text
"AlwaysOffSampler"
, shouldSample :: Context
-> TraceId
-> Text
-> SpanArguments
-> IO (SamplingResult, HashMap Text Attribute, TraceState)
shouldSample = \Context
ctxt TraceId
_ Text
_ SpanArguments
_ -> do
Maybe SpanContext
mspanCtxt <- Maybe (IO SpanContext) -> IO (Maybe SpanContext)
forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
forall (m :: * -> *) a. Monad m => Maybe (m a) -> m (Maybe a)
sequence (Span -> IO SpanContext
forall (m :: * -> *). MonadIO m => Span -> m SpanContext
getSpanContext (Span -> IO SpanContext) -> Maybe Span -> Maybe (IO SpanContext)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Context -> Maybe Span
lookupSpan Context
ctxt)
(SamplingResult, HashMap Text Attribute, TraceState)
-> IO (SamplingResult, HashMap Text Attribute, TraceState)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (SamplingResult
Drop, [], TraceState
-> (SpanContext -> TraceState) -> Maybe SpanContext -> TraceState
forall b a. b -> (a -> b) -> Maybe a -> b
maybe TraceState
TraceState.empty SpanContext -> TraceState
traceState Maybe SpanContext
mspanCtxt)
}
traceIdRatioBased :: Double -> Sampler
traceIdRatioBased :: Double -> Sampler
traceIdRatioBased Double
fraction =
if Double
fraction Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
>= Double
1
then Sampler
alwaysOn
else Sampler
sampler
where
safeFraction :: Double
safeFraction = Double -> Double -> Double
forall a. Ord a => a -> a -> a
max Double
fraction Double
0
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
fraction 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
sampler :: Sampler
sampler =
Sampler
{ getDescription :: Text
getDescription = Text
"TraceIdRatioBased{" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> String -> Text
pack (Double -> String
forall a. Show a => a -> String
show Double
fraction) Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"}"
, shouldSample :: Context
-> TraceId
-> Text
-> SpanArguments
-> IO (SamplingResult, HashMap Text Attribute, TraceState)
shouldSample = \Context
ctxt TraceId
tid Text
_ SpanArguments
_ -> do
Maybe SpanContext
mspanCtxt <- Maybe (IO SpanContext) -> IO (Maybe SpanContext)
forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
forall (m :: * -> *) a. Monad m => Maybe (m a) -> m (Maybe a)
sequence (Span -> IO SpanContext
forall (m :: * -> *). MonadIO m => Span -> m SpanContext
getSpanContext (Span -> IO SpanContext) -> Maybe Span -> Maybe (IO SpanContext)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Context -> Maybe Span
lookupSpan Context
ctxt)
let x :: Word64
x = Get Word64 -> ByteString -> Word64
forall a. Get a -> ByteString -> a
runGet Get Word64
getWord64be (ByteString -> ByteString
L.fromStrict (ByteString -> ByteString) -> ByteString -> ByteString
forall a b. (a -> b) -> a -> b
$ Int -> ByteString -> ByteString
B.take Int
8 (ByteString -> ByteString) -> ByteString -> ByteString
forall a b. (a -> b) -> a -> b
$ TraceId -> ByteString
traceIdBytes TraceId
tid) Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
`shiftR` Int
1
if Word64
x Word64 -> Word64 -> Bool
forall a. Ord a => a -> a -> Bool
< Word64
traceIdUpperBound
then do
(SamplingResult, HashMap Text Attribute, TraceState)
-> IO (SamplingResult, HashMap Text Attribute, TraceState)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (SamplingResult
RecordAndSample, [(Text
"sampleRate", Attribute
sampleRate)], TraceState
-> (SpanContext -> TraceState) -> Maybe SpanContext -> TraceState
forall b a. b -> (a -> b) -> Maybe a -> b
maybe TraceState
TraceState.empty SpanContext -> TraceState
traceState Maybe SpanContext
mspanCtxt)
else (SamplingResult, HashMap Text Attribute, TraceState)
-> IO (SamplingResult, HashMap Text Attribute, TraceState)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (SamplingResult
Drop, [], TraceState
-> (SpanContext -> TraceState) -> Maybe SpanContext -> TraceState
forall b a. b -> (a -> b) -> Maybe a -> b
maybe TraceState
TraceState.empty SpanContext -> TraceState
traceState Maybe SpanContext
mspanCtxt)
}
data ParentBasedOptions = ParentBasedOptions
{ ParentBasedOptions -> Sampler
rootSampler :: Sampler
, ParentBasedOptions -> Sampler
remoteParentSampled :: Sampler
, ParentBasedOptions -> Sampler
remoteParentNotSampled :: Sampler
, ParentBasedOptions -> Sampler
localParentSampled :: Sampler
, ParentBasedOptions -> Sampler
localParentNotSampled :: Sampler
}
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
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
..} =
Sampler
{ getDescription :: Text
getDescription =
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
"}"
, shouldSample :: Context
-> TraceId
-> Text
-> SpanArguments
-> IO (SamplingResult, HashMap Text Attribute, TraceState)
shouldSample = \Context
ctx TraceId
tid Text
name SpanArguments
csa -> do
Maybe SpanContext
mspanCtxt <- Maybe (IO SpanContext) -> IO (Maybe SpanContext)
forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
forall (m :: * -> *) a. Monad m => Maybe (m a) -> m (Maybe a)
sequence (Span -> IO SpanContext
forall (m :: * -> *). MonadIO m => Span -> m SpanContext
getSpanContext (Span -> IO SpanContext) -> Maybe Span -> Maybe (IO SpanContext)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Context -> Maybe Span
lookupSpan Context
ctx)
case Maybe SpanContext
mspanCtxt of
Maybe SpanContext
Nothing -> Sampler
-> Context
-> TraceId
-> Text
-> SpanArguments
-> IO (SamplingResult, HashMap Text Attribute, TraceState)
shouldSample Sampler
rootSampler Context
ctx TraceId
tid Text
name SpanArguments
csa
Just SpanContext
root ->
if SpanContext -> Bool
OpenTelemetry.Internal.Trace.Types.isRemote SpanContext
root
then
if TraceFlags -> Bool
isSampled (TraceFlags -> Bool) -> TraceFlags -> Bool
forall a b. (a -> b) -> a -> b
$ SpanContext -> TraceFlags
traceFlags SpanContext
root
then Sampler
-> Context
-> TraceId
-> Text
-> SpanArguments
-> IO (SamplingResult, HashMap Text Attribute, TraceState)
shouldSample Sampler
remoteParentSampled Context
ctx TraceId
tid Text
name SpanArguments
csa
else Sampler
-> Context
-> TraceId
-> Text
-> SpanArguments
-> IO (SamplingResult, HashMap Text Attribute, TraceState)
shouldSample Sampler
remoteParentNotSampled Context
ctx TraceId
tid Text
name SpanArguments
csa
else
if TraceFlags -> Bool
isSampled (TraceFlags -> Bool) -> TraceFlags -> Bool
forall a b. (a -> b) -> a -> b
$ SpanContext -> TraceFlags
traceFlags SpanContext
root
then Sampler
-> Context
-> TraceId
-> Text
-> SpanArguments
-> IO (SamplingResult, HashMap Text Attribute, TraceState)
shouldSample Sampler
localParentSampled Context
ctx TraceId
tid Text
name SpanArguments
csa
else Sampler
-> Context
-> TraceId
-> Text
-> SpanArguments
-> IO (SamplingResult, HashMap Text Attribute, TraceState)
shouldSample Sampler
localParentNotSampled Context
ctx TraceId
tid Text
name SpanArguments
csa
}