Flags.hs

module Effectful.OpenTelemetry.Tracing.Trace.Flags
    ( Flags (..)
    , Remote (..)
    , toHex
    , toWord8
    , toWord32
    , fromWord32
    )
where

import Control.Monad (mzero)
import Data.Aeson.Types (FromJSON (..), ToJSON (..), Value (..))
import Data.Bits (Bits (..))
import Data.Bool (bool)
import Data.Text (Text)
import Data.Text qualified as Text
import Data.Word (Word32, Word8)
import GHC.Generics (Generic)
import Numeric (readHex, showHex)
import Proto3.Wire.Encode.Class qualified as Proto
import Prelude

-- | Describes whether a span's parent is remote.
data Remote
    = Unknown -- 0x
    | IsNotRemote -- 10
    | IsRemote -- 11
    deriving stock (Show, Eq, Bounded, Enum)

-- | Per-trace flags propagated alongside the 'Trace.ID' and 'Span.ID'.
data Flags = Flags
    { sampled :: Bool
    -- ^ The caller may have recorded trace data. When unset, the
    -- caller did not record trace data out-of-band.
    , random :: Bool
    -- ^ At least the right-most 7 bytes of the 'Trace.ID' have been
    -- selected randomly (or pseudo-randomly) with uniform distribution.
    , remote :: Remote
    }
    deriving stock (Generic, Show, Eq)

sampledBit, randomBit, hasIsRemoteBit, isRemoteBit :: Int
sampledBit = 0
randomBit = 1
hasIsRemoteBit = 8
isRemoteBit = 9

toWord32 :: Flags -> Word32
toWord32 Flags{..} =
    foldr
        (.|.)
        zeroBits
        [ setBit' sampledBit sampled
        , setBit' randomBit random
        , setBit' hasIsRemoteBit $ remote /= Unknown
        , setBit' isRemoteBit $ remote == IsRemote
        ]
  where
    setBit' :: Int -> Bool -> Word32
    setBit' = bool zeroBits . bit

fromWord32 :: Word32 -> Flags
fromWord32 bits =
    Flags
        { sampled = testBit bits sampledBit
        , random = testBit bits randomBit
        , remote = case (testBit bits hasIsRemoteBit, testBit bits isRemoteBit) of
            (False, _) -> Unknown
            (_, False) -> IsNotRemote
            (_, True) -> IsRemote
        }

instance Monoid Flags where
    mempty = fromWord32 zeroBits

instance Semigroup Flags where
    a <> b = fromWord32 $ toWord32 a .|. toWord32 b

toWord8 :: Flags -> Word8
toWord8 Flags{..} = setBit' sampledBit sampled .|. setBit' randomBit random
  where
    setBit' :: Int -> Bool -> Word8
    setBit' = bool zeroBits . bit

-- | Encode the @trace-flags@ field as the two hex characters it is specified as.
toHex :: Flags -> Text
toHex flags = Text.justifyRight 2 '0' . Text.pack $ showHex (toWord8 flags) ""

instance FromJSON Flags where
    parseJSON = \case
        Number n -> pure . fromWord32 . round $ n
        String s -> case readHex (Text.unpack s) of
            [(w, "")] -> pure $ fromWord32 w
            _ -> mzero
        _ -> mzero

instance ToJSON Flags where
    toJSON = Number . fromIntegral . toWord32

instance {-# OVERLAPPING #-} Proto.EncodeField Flags where
    encodeField n = Proto.encodeField n . toWord32