{-|
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 (m :: * -> *). MonadIO m => TraceID -> Text -> m Bool)
 -> Sampler)
-> (forall (m :: * -> *). MonadIO m => TraceID -> Text -> m Bool)
-> Sampler
forall a b. (a -> b) -> a -> b
$ \TraceID
_ Text
_ -> Bool -> m Bool
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 (Double -> Double -> Double
forall a. Ord a => a -> a -> a
max Double
0.0 (Double -> Double) -> (Double -> Double) -> Double -> Double
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Double -> Double -> Double
forall a. Ord a => a -> a -> a
min Double
1.0 -> Double
rate) = (forall (m :: * -> *). MonadIO m => TraceID -> Text -> m Bool)
-> Sampler
Sampler ((forall (m :: * -> *). MonadIO m => TraceID -> Text -> m Bool)
 -> Sampler)
-> (forall (m :: * -> *). MonadIO m => TraceID -> Text -> m Bool)
-> Sampler
forall a b. (a -> b) -> a -> b
$ \TraceID
trace Text
_ ->
    Bool -> m Bool
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Bool -> m Bool) -> Bool -> m Bool
forall a b. (a -> b) -> a -> b
$ Word64
boundary Word64 -> Word64 -> Bool
forall a. Ord a => a -> a -> Bool
>= TraceID -> Word64
traceIdLo TraceID
trace
  where
    boundary :: Word64
boundary = Double -> Word64
forall a b. (RealFrac a, Integral b) => a -> b
round (Double -> Word64) -> Double -> Word64
forall a b. (a -> b) -> a -> b
$ Double
maxRand Double -> Double -> Double
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 (Double -> Double -> Double
forall a. Ord a => a -> a -> a
max Double
1.0 Double
tps)
    Sampler -> IO Sampler
forall (m :: * -> *) a. Monad m => a -> m a
return (Sampler -> IO Sampler) -> Sampler -> IO Sampler
forall a b. (a -> b) -> a -> b
$ (forall (m :: * -> *). MonadIO m => TraceID -> Text -> m Bool)
-> Sampler
Sampler ((forall (m :: * -> *). MonadIO m => TraceID -> Text -> m Bool)
 -> Sampler)
-> (forall (m :: * -> *). MonadIO m => TraceID -> Text -> m Bool)
-> Sampler
forall a b. (a -> b) -> a -> b
$ \TraceID
_ Text
_ -> IO Bool -> m Bool
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO Bool -> m Bool) -> IO Bool -> m Bool
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
    (Double
 -> IORef Double
 -> Double
 -> IORef TimeSpec
 -> IO TimeSpec
 -> RateLimiter)
-> IO Double
-> IO
     (IORef Double
      -> Double -> IORef TimeSpec -> IO TimeSpec -> RateLimiter)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Double -> IO Double
forall (f :: * -> *) a. Applicative f => a -> f a
pure Double
creds
    IO
  (IORef Double
   -> Double -> IORef TimeSpec -> IO TimeSpec -> RateLimiter)
-> IO (IORef Double)
-> IO (Double -> IORef TimeSpec -> IO TimeSpec -> RateLimiter)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Double -> IO (IORef Double)
forall a. a -> IO (IORef a)
newIORef Double
maxb
    IO (Double -> IORef TimeSpec -> IO TimeSpec -> RateLimiter)
-> IO Double -> IO (IORef TimeSpec -> IO TimeSpec -> RateLimiter)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Double -> IO Double
forall (f :: * -> *) a. Applicative f => a -> f a
pure Double
maxb
    IO (IORef TimeSpec -> IO TimeSpec -> RateLimiter)
-> IO (IORef TimeSpec) -> IO (IO TimeSpec -> RateLimiter)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (IO TimeSpec
tnow IO TimeSpec
-> (TimeSpec -> IO (IORef TimeSpec)) -> IO (IORef TimeSpec)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= TimeSpec -> IO (IORef TimeSpec)
forall a. a -> IO (IORef a)
newIORef)
    IO (IO TimeSpec -> RateLimiter)
-> IO (IO TimeSpec) -> IO RateLimiter
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> IO TimeSpec -> IO (IO TimeSpec)
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) <- IORef TimeSpec
-> (TimeSpec -> (TimeSpec, (TimeSpec, TimeSpec)))
-> IO (TimeSpec, TimeSpec)
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' IORef TimeSpec
lastTick ((TimeSpec -> (TimeSpec, (TimeSpec, TimeSpec)))
 -> IO (TimeSpec, TimeSpec))
-> (TimeSpec -> (TimeSpec, (TimeSpec, TimeSpec)))
-> IO (TimeSpec, TimeSpec)
forall a b. (a -> b) -> a -> b
$ \TimeSpec
x ->
                    if TimeSpec
now TimeSpec -> TimeSpec -> Bool
forall a. Ord a => a -> a -> Bool
> TimeSpec
x then (TimeSpec
now,(TimeSpec
x,TimeSpec
now)) else (TimeSpec
x,(TimeSpec
x,TimeSpec
x))

    IORef Double -> (Double -> (Double, Bool)) -> IO Bool
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' IORef Double
balance ((Double -> (Double, Bool)) -> IO Bool)
-> (Double -> (Double, Bool)) -> IO Bool
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'    = Double -> Double -> Double
forall a. Ord a => a -> a -> a
min Double
maxBalance (Double
bal Double -> Double -> Double
forall a. Num a => a -> a -> a
+ (Int64 -> Double
forall a b. (Real a, Fractional b) => a -> b
realToFrac (TimeSpec -> Int64
sec TimeSpec
elapsed) Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
creds))
        if Double
bal' Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
>= Double
cost then
            (Double
bal' Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
cost, Bool
True)
        else
            (Double
bal', Bool
False)