{-|
Module: OpenTracing.Sampling

Distributed traces are sampled, meaning they (and the spans that make them up) are
selected to either be reported or not.

This module defines a few different ways to determine if a given trace should be
sampled.
-}

{-# LANGUAGE RankNTypes      #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE StrictData      #-}
{-# LANGUAGE ViewPatterns    #-}

module OpenTracing.Sampling
    ( Sampler(..)
    , constSampler
    , probSampler
    , rateLimitSampler
    )
where

import Control.Monad.IO.Class
import Data.IORef
import Data.Text              (Text)
import OpenTracing.Types      (TraceID (..))
import System.Clock

-- | A `Sampler` is an algorithm for determine if a trace should be reported.
newtype Sampler = Sampler
    { Sampler
-> forall (m :: * -> *). MonadIO m => TraceID -> Text -> m Bool
runSampler :: forall m. MonadIO m => TraceID -> Text -> m Bool
      -- ^ Run a sampler, providing it a trace id and the operation of the span.
    }

-- | A `Sampler` that always returns the given value. Useful for debug environments.
constSampler :: Bool -> Sampler
constSampler :: Bool -> Sampler
constSampler Bool
x = (forall (m :: * -> *). MonadIO m => TraceID -> Text -> m Bool)
-> Sampler
Sampler forall a b. (a -> b) -> a -> b
$ \TraceID
_ Text
_ -> forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
x

-- | A `Sampler` that randomly chooses to report a given percentage of traces. The
-- source of randomness is the ID of the trace.
probSampler
  :: Double -- ^ A probability percentage, between 0.0 and 1.0
  -> Sampler
probSampler :: Double -> Sampler
probSampler (forall a. Ord a => a -> a -> a
max Double
0.0 forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Ord a => a -> a -> a
min Double
1.0 -> Double
rate) = (forall (m :: * -> *). MonadIO m => TraceID -> Text -> m Bool)
-> Sampler
Sampler forall a b. (a -> b) -> a -> b
$ \TraceID
trace Text
_ ->
    forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ Word64
boundary forall a. Ord a => a -> a -> Bool
>= TraceID -> Word64
traceIdLo TraceID
trace
  where
    boundary :: Word64
boundary = forall a b. (RealFrac a, Integral b) => a -> b
round forall a b. (a -> b) -> a -> b
$ Double
maxRand forall a. Num a => a -> a -> a
* Double
rate
    maxRand :: Double
maxRand  = Double
0x7fffffffffffffff

-- | A `Sampler` that will report the given number of traces per second.
rateLimitSampler
  :: Double -- ^ Traces per second
  -> IO Sampler
rateLimitSampler :: Double -> IO Sampler
rateLimitSampler Double
tps = do
    RateLimiter
lim <- Double -> Double -> IO RateLimiter
newRateLimiter Double
tps (forall a. Ord a => a -> a -> a
max Double
1.0 Double
tps)
    forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ (forall (m :: * -> *). MonadIO m => TraceID -> Text -> m Bool)
-> Sampler
Sampler forall a b. (a -> b) -> a -> b
$ \TraceID
_ Text
_ -> forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ RateLimiter -> Double -> IO Bool
haveCredit RateLimiter
lim Double
1.0

data RateLimiter = RateLimiter
    { RateLimiter -> Double
creds      :: Double
    , RateLimiter -> IORef Double
balance    :: IORef Double
    , RateLimiter -> Double
maxBalance :: Double
    , RateLimiter -> IORef TimeSpec
lastTick   :: IORef TimeSpec
    , RateLimiter -> IO TimeSpec
timeNow    :: IO TimeSpec
    }

newRateLimiter :: Double -> Double -> IO RateLimiter
newRateLimiter :: Double -> Double -> IO RateLimiter
newRateLimiter Double
creds Double
maxb = Double
-> IORef Double
-> Double
-> IORef TimeSpec
-> IO TimeSpec
-> RateLimiter
RateLimiter
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (f :: * -> *) a. Applicative f => a -> f a
pure Double
creds
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall a. a -> IO (IORef a)
newIORef Double
maxb
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (f :: * -> *) a. Applicative f => a -> f a
pure Double
maxb
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (IO TimeSpec
tnow forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= forall a. a -> IO (IORef a)
newIORef)
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (f :: * -> *) a. Applicative f => a -> f a
pure IO TimeSpec
tnow
 where
    tnow :: IO TimeSpec
tnow = Clock -> IO TimeSpec
getTime Clock
Monotonic

haveCredit :: RateLimiter -> Double -> IO Bool
haveCredit :: RateLimiter -> Double -> IO Bool
haveCredit RateLimiter{Double
IO TimeSpec
IORef Double
IORef TimeSpec
timeNow :: IO TimeSpec
lastTick :: IORef TimeSpec
maxBalance :: Double
balance :: IORef Double
creds :: Double
timeNow :: RateLimiter -> IO TimeSpec
lastTick :: RateLimiter -> IORef TimeSpec
maxBalance :: RateLimiter -> Double
balance :: RateLimiter -> IORef Double
creds :: RateLimiter -> Double
..} Double
cost = do
    TimeSpec
now     <- IO TimeSpec
timeNow
    (TimeSpec
lst,TimeSpec
t) <- forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' IORef TimeSpec
lastTick forall a b. (a -> b) -> a -> b
$ \TimeSpec
x ->
                    if TimeSpec
now forall a. Ord a => a -> a -> Bool
> TimeSpec
x then (TimeSpec
now,(TimeSpec
x,TimeSpec
now)) else (TimeSpec
x,(TimeSpec
x,TimeSpec
x))

    forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' IORef Double
balance forall a b. (a -> b) -> a -> b
$ \Double
bal -> do
        let elapsed :: TimeSpec
elapsed = TimeSpec -> TimeSpec -> TimeSpec
diffTimeSpec TimeSpec
lst TimeSpec
t
        let bal' :: Double
bal'    = forall a. Ord a => a -> a -> a
min Double
maxBalance (Double
bal forall a. Num a => a -> a -> a
+ (forall a b. (Real a, Fractional b) => a -> b
realToFrac (TimeSpec -> Int64
sec TimeSpec
elapsed) forall a. Num a => a -> a -> a
* Double
creds))
        if Double
bal' forall a. Ord a => a -> a -> Bool
>= Double
cost then
            (Double
bal' forall a. Num a => a -> a -> a
- Double
cost, Bool
True)
        else
            (Double
bal', Bool
False)