module Hasura.Tracing.TraceId
  ( -- * TraceId
    TraceId,
    randomTraceId,
    traceIdFromBytes,
    traceIdToBytes,
    traceIdFromHex,
    traceIdToHex,

    -- * SpanId
    SpanId,
    randomSpanId,
    spanIdFromBytes,
    spanIdToBytes,
    spanIdFromHex,
    spanIdToHex,
  )
where

import Data.Bits ((.|.))
import Data.ByteString (ByteString)
import Data.ByteString qualified as ByteString
import Data.ByteString.Base16 qualified as Base16
import Data.Serialize qualified as Serialize
import Hasura.Prelude
import System.Random.Stateful qualified as Random

--------------------------------------------------------------------------------
-- TraceId

-- | 128-bit trace identifiers.
--
-- 'TraceId's are guaranteed to have at least one non-zero bit.
data TraceId
  = TraceId
      {-# UNPACK #-} !Word64
      {-# UNPACK #-} !Word64
  deriving (TraceId -> TraceId -> Bool
(TraceId -> TraceId -> Bool)
-> (TraceId -> TraceId -> Bool) -> Eq TraceId
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: TraceId -> TraceId -> Bool
== :: TraceId -> TraceId -> Bool
$c/= :: TraceId -> TraceId -> Bool
/= :: TraceId -> TraceId -> Bool
Eq)

-- 128 bits
traceIdBytes :: Int
traceIdBytes :: Int
traceIdBytes = Int
16

randomTraceId :: (MonadIO m) => m TraceId
randomTraceId :: forall (m :: * -> *). MonadIO m => m TraceId
randomTraceId = IO TraceId -> m TraceId
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO do
  (Word64
w1, Word64
w2) <-
    ((StdGen -> ((Word64, Word64), StdGen))
 -> AtomicGenM StdGen -> IO (Word64, Word64))
-> AtomicGenM StdGen
-> (StdGen -> ((Word64, Word64), StdGen))
-> IO (Word64, Word64)
forall a b c. (a -> b -> c) -> b -> a -> c
flip (StdGen -> ((Word64, Word64), StdGen))
-> AtomicGenM StdGen -> IO (Word64, Word64)
forall (m :: * -> *) g a.
MonadIO m =>
(g -> (a, g)) -> AtomicGenM g -> m a
Random.applyAtomicGen AtomicGenM StdGen
Random.globalStdGen ((StdGen -> ((Word64, Word64), StdGen)) -> IO (Word64, Word64))
-> (StdGen -> ((Word64, Word64), StdGen)) -> IO (Word64, Word64)
forall a b. (a -> b) -> a -> b
$ \StdGen
gen0 ->
      let (!Word64
w1, !StdGen
gen1) = StdGen -> (Word64, StdGen)
forall g. RandomGen g => g -> (Word64, g)
forall a g. (Random a, RandomGen g) => g -> (a, g)
Random.random StdGen
gen0
          (!Word64
w2, !StdGen
gen2) = StdGen -> (Word64, StdGen)
forall g. RandomGen g => g -> (Word64, g)
forall a g. (Random a, RandomGen g) => g -> (a, g)
Random.random StdGen
gen1
       in ((Word64
w1, Word64
w2), StdGen
gen2)
  if Word64
w1 Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
.|. Word64
w2 Word64 -> Word64 -> Bool
forall a. Eq a => a -> a -> Bool
== Word64
0
    then IO TraceId
forall (m :: * -> *). MonadIO m => m TraceId
randomTraceId
    else TraceId -> IO TraceId
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (TraceId -> IO TraceId) -> TraceId -> IO TraceId
forall a b. (a -> b) -> a -> b
$ Word64 -> Word64 -> TraceId
TraceId Word64
w1 Word64
w2

-- | Create a 'TraceId' from a 'ByteString'.
--
-- Fails if the 'ByteString' is not exactly 16 bytes long, or if it contains
-- only zero bytes.
traceIdFromBytes :: ByteString -> Maybe TraceId
traceIdFromBytes :: ByteString -> Maybe TraceId
traceIdFromBytes ByteString
bs = do
  Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Bool -> Maybe ()) -> Bool -> Maybe ()
forall a b. (a -> b) -> a -> b
$ ByteString -> Int
ByteString.length ByteString
bs Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
traceIdBytes
  (Word64
w1, Word64
w2) <-
    Either String (Word64, Word64) -> Maybe (Word64, Word64)
forall a b. Either a b -> Maybe b
eitherToMaybe
      (Either String (Word64, Word64) -> Maybe (Word64, Word64))
-> Either String (Word64, Word64) -> Maybe (Word64, Word64)
forall a b. (a -> b) -> a -> b
$ (Get (Word64, Word64)
 -> ByteString -> Either String (Word64, Word64))
-> ByteString
-> Get (Word64, Word64)
-> Either String (Word64, Word64)
forall a b c. (a -> b -> c) -> b -> a -> c
flip Get (Word64, Word64)
-> ByteString -> Either String (Word64, Word64)
forall a. Get a -> ByteString -> Either String a
Serialize.runGet ByteString
bs
      (Get (Word64, Word64) -> Either String (Word64, Word64))
-> Get (Word64, Word64) -> Either String (Word64, Word64)
forall a b. (a -> b) -> a -> b
$ (,)
      (Word64 -> Word64 -> (Word64, Word64))
-> Get Word64 -> Get (Word64 -> (Word64, Word64))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Get Word64
Serialize.getWord64be
      Get (Word64 -> (Word64, Word64))
-> Get Word64 -> Get (Word64, Word64)
forall a b. Get (a -> b) -> Get a -> Get b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Get Word64
Serialize.getWord64be
  Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Bool -> Maybe ()) -> Bool -> Maybe ()
forall a b. (a -> b) -> a -> b
$ Word64
w1 Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
.|. Word64
w2 Word64 -> Word64 -> Bool
forall a. Eq a => a -> a -> Bool
/= Word64
0
  TraceId -> Maybe TraceId
forall a. a -> Maybe a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (TraceId -> Maybe TraceId) -> TraceId -> Maybe TraceId
forall a b. (a -> b) -> a -> b
$ Word64 -> Word64 -> TraceId
TraceId Word64
w1 Word64
w2

-- | Convert a 'TraceId' to a 'ByteString' of 16 bytes.
traceIdToBytes :: TraceId -> ByteString
traceIdToBytes :: TraceId -> ByteString
traceIdToBytes (TraceId Word64
w1 Word64
w2) =
  Put -> ByteString
Serialize.runPut (Put -> ByteString) -> Put -> ByteString
forall a b. (a -> b) -> a -> b
$ Putter Word64
Serialize.putWord64be Word64
w1 Put -> Put -> Put
forall a b. PutM a -> PutM b -> PutM b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Putter Word64
Serialize.putWord64be Word64
w2

-- | Create a 'TraceId' from a 'ByteString' of hex characters.
--
-- Fails if the 'ByteString' is not exactly 32 characters long, or if it
-- contains only zero characters.
traceIdFromHex :: ByteString -> Maybe TraceId
traceIdFromHex :: ByteString -> Maybe TraceId
traceIdFromHex = ByteString -> Maybe TraceId
traceIdFromBytes (ByteString -> Maybe TraceId)
-> (ByteString -> Maybe ByteString) -> ByteString -> Maybe TraceId
forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< Either String ByteString -> Maybe ByteString
forall a b. Either a b -> Maybe b
eitherToMaybe (Either String ByteString -> Maybe ByteString)
-> (ByteString -> Either String ByteString)
-> ByteString
-> Maybe ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> Either String ByteString
Base16.decode

-- | Convert a 'TraceId' to a 'ByteString' of 32 lowercase hex characters.
traceIdToHex :: TraceId -> ByteString
traceIdToHex :: TraceId -> ByteString
traceIdToHex = ByteString -> ByteString
Base16.encode (ByteString -> ByteString)
-> (TraceId -> ByteString) -> TraceId -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TraceId -> ByteString
traceIdToBytes

--------------------------------------------------------------------------------
-- SpanId

-- | 64-bit span identifiers
--
-- 'SpanId's are guaranteed to have at least one non-zero bit.
newtype SpanId = SpanId Word64
  deriving (SpanId -> SpanId -> Bool
(SpanId -> SpanId -> Bool)
-> (SpanId -> SpanId -> Bool) -> Eq SpanId
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: SpanId -> SpanId -> Bool
== :: SpanId -> SpanId -> Bool
$c/= :: SpanId -> SpanId -> Bool
/= :: SpanId -> SpanId -> Bool
Eq)

-- 64 bits
spanIdBytes :: Int
spanIdBytes :: Int
spanIdBytes = Int
8

randomSpanId :: (MonadIO m) => m SpanId
randomSpanId :: forall (m :: * -> *). MonadIO m => m SpanId
randomSpanId = IO SpanId -> m SpanId
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO do
  Word64
w <- AtomicGenM StdGen -> IO Word64
forall a g (m :: * -> *). (Uniform a, StatefulGen g m) => g -> m a
forall g (m :: * -> *). StatefulGen g m => g -> m Word64
Random.uniformM AtomicGenM StdGen
Random.globalStdGen
  if Word64
w Word64 -> Word64 -> Bool
forall a. Eq a => a -> a -> Bool
== Word64
0
    then IO SpanId
forall (m :: * -> *). MonadIO m => m SpanId
randomSpanId
    else SpanId -> IO SpanId
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (SpanId -> IO SpanId) -> SpanId -> IO SpanId
forall a b. (a -> b) -> a -> b
$ Word64 -> SpanId
SpanId Word64
w

-- | Create a 'SpanId' from a 'ByteString'.
--
-- Fails if the 'ByteString' is not exactly 8 bytes long, or if it contains
-- only zero bytes.
spanIdFromBytes :: ByteString -> Maybe SpanId
spanIdFromBytes :: ByteString -> Maybe SpanId
spanIdFromBytes ByteString
bs = do
  Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Bool -> Maybe ()) -> Bool -> Maybe ()
forall a b. (a -> b) -> a -> b
$ ByteString -> Int
ByteString.length ByteString
bs Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
spanIdBytes
  Word64
w <- Either String Word64 -> Maybe Word64
forall a b. Either a b -> Maybe b
eitherToMaybe (Either String Word64 -> Maybe Word64)
-> Either String Word64 -> Maybe Word64
forall a b. (a -> b) -> a -> b
$ Get Word64 -> ByteString -> Either String Word64
forall a. Get a -> ByteString -> Either String a
Serialize.runGet Get Word64
Serialize.getWord64be ByteString
bs
  Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Bool -> Maybe ()) -> Bool -> Maybe ()
forall a b. (a -> b) -> a -> b
$ Word64
w Word64 -> Word64 -> Bool
forall a. Eq a => a -> a -> Bool
/= Word64
0
  SpanId -> Maybe SpanId
forall a. a -> Maybe a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (SpanId -> Maybe SpanId) -> SpanId -> Maybe SpanId
forall a b. (a -> b) -> a -> b
$ Word64 -> SpanId
SpanId Word64
w

-- | Convert a 'SpanId' to a 'ByteString' of 8 bytes.
spanIdToBytes :: SpanId -> ByteString
spanIdToBytes :: SpanId -> ByteString
spanIdToBytes (SpanId Word64
w) = Put -> ByteString
Serialize.runPut (Put -> ByteString) -> Put -> ByteString
forall a b. (a -> b) -> a -> b
$ Putter Word64
Serialize.putWord64be Word64
w

-- | Create a 'SpanId' from a 'ByteString' of hex characters.
--
-- Fails if the 'ByteString' is not exactly 16 characters long, or if it
-- contains only zero characters.
spanIdFromHex :: ByteString -> Maybe SpanId
spanIdFromHex :: ByteString -> Maybe SpanId
spanIdFromHex = ByteString -> Maybe SpanId
spanIdFromBytes (ByteString -> Maybe SpanId)
-> (ByteString -> Maybe ByteString) -> ByteString -> Maybe SpanId
forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< Either String ByteString -> Maybe ByteString
forall a b. Either a b -> Maybe b
eitherToMaybe (Either String ByteString -> Maybe ByteString)
-> (ByteString -> Either String ByteString)
-> ByteString
-> Maybe ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> Either String ByteString
Base16.decode

-- | Convert a 'SpanId' to a 'ByteString' of 16 lowercase hex characters.
spanIdToHex :: SpanId -> ByteString
spanIdToHex :: SpanId -> ByteString
spanIdToHex = ByteString -> ByteString
Base16.encode (ByteString -> ByteString)
-> (SpanId -> ByteString) -> SpanId -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SpanId -> ByteString
spanIdToBytes