module Hasura.Tracing.Sampling
  ( -- * SamplingState
    SamplingState (..),
    samplingStateToHeader,
    samplingStateFromHeader,

    -- * SamplingDecision
    SamplingDecision (..),

    -- * SamplingPolicy
    SamplingPolicy,
    sampleNever,
    sampleAlways,
    sampleRandomly,
    sampleOneInN,
  )
where

import Hasura.Prelude
import Refined (Positive, Refined, unrefine)
import System.Random.Stateful qualified as Random

--------------------------------------------------------------------------------
-- SamplingState

-- | B3 propagation sampling state.
--
-- Debug sampling state not represented.
data SamplingState = SamplingDefer | SamplingDeny | SamplingAccept

-- | Convert a sampling state to a value for the X-B3-Sampled header. A return
-- value of Nothing indicates that the header should not be set.
samplingStateToHeader :: (IsString s) => SamplingState -> Maybe s
samplingStateToHeader :: forall s. IsString s => SamplingState -> Maybe s
samplingStateToHeader = \case
  SamplingState
SamplingDefer -> Maybe s
forall a. Maybe a
Nothing
  SamplingState
SamplingDeny -> s -> Maybe s
forall a. a -> Maybe a
Just s
"0"
  SamplingState
SamplingAccept -> s -> Maybe s
forall a. a -> Maybe a
Just s
"1"

-- | Convert a X-B3-Sampled header value to a sampling state. An input of
-- Nothing indicates that the header was not set.
samplingStateFromHeader :: (IsString s, Eq s) => Maybe s -> SamplingState
samplingStateFromHeader :: forall s. (IsString s, Eq s) => Maybe s -> SamplingState
samplingStateFromHeader = \case
  Maybe s
Nothing -> SamplingState
SamplingDefer
  Just s
"0" -> SamplingState
SamplingDeny
  Just s
"1" -> SamplingState
SamplingAccept
  Just s
_ -> SamplingState
SamplingDefer

--------------------------------------------------------------------------------
-- SamplingDecision

-- | A local decision about whether or not to sample spans.
data SamplingDecision = SampleNever | SampleAlways

--------------------------------------------------------------------------------
-- SamplingPolicy

-- | An IO action for deciding whether or not to sample a trace.
--
-- Currently restricted to deny access to the B3 sampling state, but we may
-- want to be more flexible in the future.
type SamplingPolicy = IO SamplingDecision

sampleNever :: SamplingPolicy
sampleNever :: SamplingPolicy
sampleNever = SamplingDecision -> SamplingPolicy
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SamplingDecision
SampleNever

sampleAlways :: SamplingPolicy
sampleAlways :: SamplingPolicy
sampleAlways = SamplingDecision -> SamplingPolicy
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SamplingDecision
SampleAlways

-- @sampleRandomly p@ returns `SampleAlways` with probability @p@ and
-- `SampleNever` with probability @1 - p@.
sampleRandomly :: Double -> SamplingPolicy
sampleRandomly :: Double -> SamplingPolicy
sampleRandomly Double
samplingProbability
  | Double
samplingProbability Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Double
0 = SamplingDecision -> SamplingPolicy
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SamplingDecision
SampleNever
  | Double
samplingProbability Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
>= Double
1 = SamplingDecision -> SamplingPolicy
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SamplingDecision
SampleAlways
  | Bool
otherwise = do
      Double
x <- (Double, Double) -> AtomicGenM StdGen -> IO Double
forall a g (m :: * -> *).
(UniformRange a, StatefulGen g m) =>
(a, a) -> g -> m a
forall g (m :: * -> *).
StatefulGen g m =>
(Double, Double) -> g -> m Double
Random.uniformRM (Double
0, Double
1) AtomicGenM StdGen
Random.globalStdGen
      SamplingDecision -> SamplingPolicy
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (SamplingDecision -> SamplingPolicy)
-> SamplingDecision -> SamplingPolicy
forall a b. (a -> b) -> a -> b
$ if Double
x Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
samplingProbability then SamplingDecision
SampleAlways else SamplingDecision
SampleNever

-- Like @sampleRandomly@, but with the probability expressed as the denominator
-- N of the fraction 1/N.
sampleOneInN :: Refined Positive Int -> SamplingPolicy
sampleOneInN :: Refined Positive Int -> SamplingPolicy
sampleOneInN Refined Positive Int
denominator
  | Int
n Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
1 = SamplingDecision -> SamplingPolicy
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SamplingDecision
SampleAlways
  | Bool
otherwise = do
      Int
x <- (Int, Int) -> AtomicGenM StdGen -> IO Int
forall a g (m :: * -> *).
(UniformRange a, StatefulGen g m) =>
(a, a) -> g -> m a
forall g (m :: * -> *). StatefulGen g m => (Int, Int) -> g -> m Int
Random.uniformRM (Int
0, Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) AtomicGenM StdGen
Random.globalStdGen
      SamplingDecision -> SamplingPolicy
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (SamplingDecision -> SamplingPolicy)
-> SamplingDecision -> SamplingPolicy
forall a b. (a -> b) -> a -> b
$ if Int
x Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 then SamplingDecision
SampleAlways else SamplingDecision
SampleNever
  where
    n :: Int
n = Refined Positive Int -> Int
forall {k} (p :: k) x. Refined p x -> x
unrefine Refined Positive Int
denominator