Span.hs

module Effectful.OpenTelemetry.Tracing.Span where

import Data.Aeson.Types (Pair, ToJSON (..), Value, object, (.=))
import Data.Function ((&))
import Data.Maybe (catMaybes)
import Data.Sequence (Seq)
import Data.Text (Text)
import Data.Text qualified as Text
import Effectful.OpenTelemetry.Protocol.Attributes (Attributes)
import Effectful.OpenTelemetry.Protocol.Export qualified as Export
import Effectful.OpenTelemetry.Protocol.Resource qualified as Resource
import Effectful.OpenTelemetry.Protocol.Scope qualified as Scope
import Effectful.OpenTelemetry.Timestamp (Timestamp (..))
import Effectful.OpenTelemetry.Tracing.Span.Context (Context (..))
import Effectful.OpenTelemetry.Tracing.Span.Event (Event)
import Effectful.OpenTelemetry.Tracing.Span.ID qualified as Span (ID)
import Effectful.OpenTelemetry.Tracing.Span.Kind (Kind (..))
import Effectful.OpenTelemetry.Tracing.Span.Status
import GHC.Generics (Generic)
import Network.GRPC.HTTP2.Proto3Wire (RPC (..))
import Prettyprinter (Pretty (..))
import Prettyprinter.Extra (PrettyAnn (..))
import Prettyprinter.Extra qualified as Pretty
import Prettyprinter.Render.Terminal (AnsiStyle)
import Proto3.Wire.Encode.Class qualified as Proto
import Prelude

-- | A single unit of work or operation in a trace. Spans are the building blocks of traces.
--
-- See <https://opentelemetry.io/docs/concepts/signals/traces/#spans the OpenTelemetry spec>.
data Span = Span
    { name :: Text
    , parentSpanId :: Maybe Span.ID
    , context :: Context
    , kind :: Kind
    , startTime :: Timestamp
    , endTime :: Maybe Timestamp
    , attributes :: Attributes
    , events :: Seq Event
    , status :: Status
    }
    deriving stock (Generic, Eq, Show)

instance ToJSON Span where
    toJSON Span{context = Context{..}, ..} =
        object $
            [ "traceId" .= traceId
            , "spanId" .= spanId
            , "traceState" .= traceState
            , "flags" .= traceFlags
            , "name" .= name
            , "kind" .= kind
            , "startTimeUnixNano" .= startTime
            , "attributes" .= attributes
            , "events" .= events
            , "status" .= status
            ]
                <> catMaybes @Pair
                    [ ("parentSpanId" .=) <$> parentSpanId
                    , ("endTimeUnixNano" .=) <$> endTime
                    ]

instance Proto.Encode Span where
    encode Span{context = Context{..}, ..} =
        mconcat
            [ Proto.encodeField 1 traceId
            , Proto.encodeField 2 spanId
            , Proto.encodeField 3 traceState
            , foldMap (Proto.encodeField 4) parentSpanId
            , Proto.encodeField 5 name
            , Proto.encodeField 6 kind
            , Proto.encodeField 7 startTime
            , foldMap (Proto.encodeField 8) endTime
            , Proto.encodeField 9 attributes
            , -- 10: dropped_attributes_count: Not supported
              Proto.encodeField 11 events
            , -- 12: dropped_events_count: Not supported
              -- TODO: 13: links
              -- 14: dropped_links_count: Not supported
              Proto.encodeField 15 status
            , Proto.encodeField 16 traceFlags
            ]

instance PrettyAnn AnsiStyle Span where
    prettyAnn Span{..} =
        Pretty.unwords . filter (not . Pretty.null) $
            [ prettyAnn startTime
            , endTime & maybe mempty \(nanos -> subtract startTime.nanos -> duration) ->
                pretty . (<> "ms") . Text.show $ duration `div` 1_000_000
            , pretty name
            , prettyAnn attributes
            , prettyAnn status
            ]

instance Export.Request Span where
    exportHttpPathComponents = ["v1", "traces"]
    exportGrpcRPC =
        RPC
            { pkg = "opentelemetry.proto.collector.trace.v1"
            , srv = "TraceService"
            , meth = "Export"
            }
    exportJson = object . pure . ("resourceSpans" .=) . fmap resourceSpans
      where
        resourceSpans :: Resource.Items Span -> Value
        resourceSpans Resource.Items{resource, scopeItems} =
            object
                [ "resource" .= resource
                , "scopeSpans" .= fmap scopeSpans scopeItems
                ]
        scopeSpans :: Scope.Items Span -> Value
        scopeSpans Scope.Items{..} =
            object
                [ "scope" .= scope
                , "spans" .= items
                ]

    transportSignalEnvName = "TRACES"
    batchEnvPrefix = Just "BSP"

    -- https://opentelemetry.io/docs/specs/otel/trace/sdk/#batching-processor
    defaultConfig =
        Export.Config
            { batch =
                Just
                    Export.BatchConfig
                        { maxQueueSize = 2_048
                        , scheduledDelayMs = 5_000
                        , maxBatchSize = 512
                        }
            , exportTimeoutMs = 30_000
            }