-----------------------------------------------------------------------------
-- |
-- Module      :  OpenTelemetry.Trace.Sampler
-- Copyright   :  (c) Ian Duncan, 2021
-- License     :  BSD-3
-- Description :  Sampling strategies for reducing tracing overhead
-- Maintainer  :  Ian Duncan
-- Stability   :  experimental
-- Portability :  non-portable (GHC extensions)
--
-- This module provides several built-in sampling strategies, as well as the ability to define custom samplers.
--
-- Sampling is the concept of selecting a few elements from a large collection and learning about the entire collection by extrapolating from the selected set. It’s widely used throughout the world whenever trying to tackle a problem of scale: for example, a survey assumes that by asking a small group of people a set of questions, you can learn something about the opinions of the entire populace.
--
-- While it’s nice to believe that every event is precious, the reality of monitoring high volume production infrastructure is that there are some attributes to events that make them more interesting than the rest. Failures are often more interesting than successes! Rare events are more interesting than common events! Capturing some traffic from all customers can be better than capturing all traffic from some customers.
--
-- Sampling as a basic technique for instrumentation is no different—by recording information about a representative subset of requests flowing through a system, you can learn about the overall performance of the system. And as with surveys and air monitoring, the way you choose your representative set (the sample set) can greatly influence the accuracy of your results.
--
-- Sampling is widespread in observability systems because it lowers the cost of producing, collecting, and analyzing data in systems anywhere cost is a concern. Developers and operators in an observability system apply or attach key=value properties to observability data–spans and metrics–and we use these properties to investigate hypotheses about our systems after the fact. It is interesting to look at how sampling impacts our ability to analyze observability data, using key=value restrictions for some keys and grouping the output based on other keys.
--
-- Sampling schemes let observability systems collect examples of data that are not merely exemplary, but also representative. Sampling schemes compute a set of representative items and, in doing so, score each item with what is commonly called the item's "sampling rate." A sampling rate of 10 indicates that the item represents an estimated 10 individuals in the original data set.
-----------------------------------------------------------------------------
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.Trace.Id
import OpenTelemetry.Context
import OpenTelemetry.Internal.Trace.Types
import OpenTelemetry.Trace.TraceState as TraceState
import OpenTelemetry.Attributes (toAttribute)

-- | Returns @RecordAndSample@ always.
--
-- Description returns AlwaysOnSampler.
--
-- @since 0.1.0.0
alwaysOn :: Sampler
alwaysOn :: Sampler
alwaysOn = Sampler :: Text
-> (Context
    -> TraceId
    -> Text
    -> SpanArguments
    -> IO (SamplingResult, [(Text, Attribute)], TraceState))
-> Sampler
Sampler
  { getDescription :: Text
getDescription = Text
"AlwaysOnSampler"
  , shouldSample :: Context
-> TraceId
-> Text
-> SpanArguments
-> IO (SamplingResult, [(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)
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, [(Text, Attribute)], TraceState)
-> IO (SamplingResult, [(Text, Attribute)], TraceState)
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)
  }

-- | Returns @RecordAndSample@ always.
--
-- Description returns AlwaysOffSampler.
--
-- @since 0.1.0.0
alwaysOff :: Sampler
alwaysOff :: Sampler
alwaysOff = Sampler :: Text
-> (Context
    -> TraceId
    -> Text
    -> SpanArguments
    -> IO (SamplingResult, [(Text, Attribute)], TraceState))
-> Sampler
Sampler
  { getDescription :: Text
getDescription = Text
"AlwaysOffSampler"
  , shouldSample :: Context
-> TraceId
-> Text
-> SpanArguments
-> IO (SamplingResult, [(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)
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, [(Text, Attribute)], TraceState)
-> IO (SamplingResult, [(Text, Attribute)], TraceState)
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)
  }

-- | The TraceIdRatioBased ignores the parent SampledFlag. To respect the parent SampledFlag, 
-- the TraceIdRatioBased should be used as a delegate of the @parentBased@ sampler specified below.
-- 
-- Description returns a string of the form "TraceIdRatioBased{RATIO}" with RATIO replaced with the Sampler 
-- instance's trace sampling ratio represented as a decimal number.
--
-- @since 0.1.0.0
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 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 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 :: Text
-> (Context
    -> TraceId
    -> Text
    -> SpanArguments
    -> IO (SamplingResult, [(Text, Attribute)], TraceState))
-> 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, [(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)
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, [(Text, Attribute)], TraceState)
-> IO (SamplingResult, [(Text, Attribute)], TraceState)
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, [(Text, Attribute)], TraceState)
-> IO (SamplingResult, [(Text, Attribute)], TraceState)
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)
      }

-- | This is a composite sampler. ParentBased helps distinguish between the following cases:
--
-- No parent (root span).
--
-- Remote parent (SpanContext.IsRemote() == true) with SampledFlag equals true
-- 
-- Remote parent (SpanContext.IsRemote() == true) with SampledFlag equals false
--
-- Local parent (SpanContext.IsRemote() == false) with SampledFlag equals true
--
-- Local parent (SpanContext.IsRemote() == false) with SampledFlag equals false
--
-- @since 0.1.0.0
data ParentBasedOptions = ParentBasedOptions 
  { ParentBasedOptions -> Sampler
rootSampler :: Sampler
  -- ^ Sampler called for spans with no parent (root spans)
  , ParentBasedOptions -> Sampler
remoteParentSampled :: Sampler
  -- ^ default: alwaysOn
  , ParentBasedOptions -> Sampler
remoteParentNotSampled :: Sampler
  -- ^ default: alwaysOff
  , ParentBasedOptions -> Sampler
localParentSampled :: Sampler
  -- ^ default: alwaysOn
  , ParentBasedOptions -> Sampler
localParentNotSampled :: Sampler
  -- ^ default: alwaysOff
  }

-- | A smart constructor for 'ParentBasedOptions' with reasonable starting
-- defaults.
--
-- @since 0.1.0.0
parentBasedOptions 
  :: Sampler 
  -- ^ Root sampler
  -> ParentBasedOptions
parentBasedOptions :: Sampler -> ParentBasedOptions
parentBasedOptions Sampler
root = ParentBasedOptions :: Sampler
-> Sampler -> Sampler -> Sampler -> Sampler -> ParentBasedOptions
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
  }

-- | A sampler which behaves differently based on the incoming sampling decision. 
--
-- In general, this will sample spans that have parents that were sampled, and will not sample spans whose parents were not sampled.
--
-- @since 0.1.0.0
parentBased :: ParentBasedOptions -> Sampler
parentBased :: ParentBasedOptions -> Sampler
parentBased ParentBasedOptions{Sampler
localParentNotSampled :: Sampler
localParentSampled :: Sampler
remoteParentNotSampled :: Sampler
remoteParentSampled :: Sampler
rootSampler :: Sampler
localParentNotSampled :: ParentBasedOptions -> Sampler
localParentSampled :: ParentBasedOptions -> Sampler
remoteParentNotSampled :: ParentBasedOptions -> Sampler
remoteParentSampled :: ParentBasedOptions -> Sampler
rootSampler :: ParentBasedOptions -> Sampler
..} = Sampler :: Text
-> (Context
    -> TraceId
    -> Text
    -> SpanArguments
    -> IO (SamplingResult, [(Text, Attribute)], TraceState))
-> 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, [(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)
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, [(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, [(Text, Attribute)], TraceState)
shouldSample Sampler
remoteParentSampled Context
ctx TraceId
tid Text
name SpanArguments
csa
            else Sampler
-> Context
-> TraceId
-> Text
-> SpanArguments
-> IO (SamplingResult, [(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, [(Text, Attribute)], TraceState)
shouldSample Sampler
localParentSampled Context
ctx TraceId
tid Text
name SpanArguments
csa
            else Sampler
-> Context
-> TraceId
-> Text
-> SpanArguments
-> IO (SamplingResult, [(Text, Attribute)], TraceState)
shouldSample Sampler
localParentNotSampled Context
ctx TraceId
tid Text
name SpanArguments
csa
  }