graphql-engine/server/src-lib/Hasura/Tracing/Sampling.hs
Toan Nguyen f915c7d1a2 server: support w3c traceparent context
PR-URL: https://github.com/hasura/graphql-engine-mono/pull/10218
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
Co-authored-by: Rob Dominguez <24390149+robertjdominguez@users.noreply.github.com>
GitOrigin-RevId: d3dbea6220fd2127ab76c0a240fc4725ca5d6aac
2023-09-13 13:42:30 +00:00

90 lines
2.8 KiB
Haskell

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
deriving (Show, Eq)
-- | 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 = \case
SamplingDefer -> Nothing
SamplingDeny -> Just "0"
SamplingAccept -> Just "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 = \case
Nothing -> SamplingDefer
Just "0" -> SamplingDeny
Just "1" -> SamplingAccept
Just _ -> 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 = pure SampleNever
sampleAlways :: SamplingPolicy
sampleAlways = pure SampleAlways
-- @sampleRandomly p@ returns `SampleAlways` with probability @p@ and
-- `SampleNever` with probability @1 - p@.
sampleRandomly :: Double -> SamplingPolicy
sampleRandomly samplingProbability
| samplingProbability <= 0 = pure SampleNever
| samplingProbability >= 1 = pure SampleAlways
| otherwise = do
x <- Random.uniformRM (0, 1) Random.globalStdGen
pure $ if x < samplingProbability then SampleAlways else SampleNever
-- Like @sampleRandomly@, but with the probability expressed as the denominator
-- N of the fraction 1/N.
sampleOneInN :: Refined Positive Int -> SamplingPolicy
sampleOneInN denominator
| n == 1 = pure SampleAlways
| otherwise = do
x <- Random.uniformRM (0, n - 1) Random.globalStdGen
pure $ if x == 0 then SampleAlways else SampleNever
where
n = unrefine denominator