Publishing.hs

module Atelier.Effects.Publishing
    ( runPubSub
    , runPubSub_
    )
where

import Data.Time (UTCTime)
import Effectful.Dispatch.Dynamic (interpret, interpretWith, interpretWith_, interpret_, localSeqUnlift)

import Atelier.Effects.Chan (Chan)
import Atelier.Effects.Clock (Clock)
import Atelier.Effects.Monitoring.Tracing (SpanContext, Tracing)
import Atelier.Effects.Publishing.Pub (Pub (..))
import Atelier.Effects.Publishing.Sub (Sub (..))

import Atelier.Effects.Chan qualified as Chan
import Atelier.Effects.Clock qualified as Clock
import Atelier.Effects.Monitoring.Tracing qualified as Tracing


-- | Internal wrapper for events with trace context
data TracedEvent event = TracedEvent
    { event :: event
    , timestamp :: UTCTime
    , publisherSpanContext :: Maybe SpanContext
    }


-- | Runs 'Pub' and 'Sub' effects with an internal channel for a specific event
-- type. Automatically captures span context from the publisher and creates
-- linked spans in listeners.
runPubSub
    :: forall event es a
     . ( Chan :> es
       , Clock :> es
       , Tracing :> es
       )
    => Eff (Pub event : Sub event : es) a -> Eff es a
runPubSub action = do
    (inChan, _) <- Chan.newChan @(TracedEvent event)

    let handlePub eff = interpretWith_ eff \case
            Publish event -> do
                timestamp <- Clock.currentTime
                -- Capture the current span context from the publisher
                publisherSpanContext <- Tracing.getSpanContext
                Chan.writeChan inChan TracedEvent {event, timestamp, publisherSpanContext}

        handleSub eff = interpretWith eff \env -> \case
            ListenWith onSubscribed listener -> localSeqUnlift env \unlift -> do
                chan <- Chan.dupChan inChan
                unlift onSubscribed
                forever do
                    TracedEvent {event, timestamp, publisherSpanContext} <- Chan.readChan chan
                    Tracing.withLinkPropagation publisherSpanContext $ unlift $ listener timestamp event

    handleSub . handlePub $ action


-- | Runs 'Pub' and 'Sub' effects with an internal channel for a specific event
-- type.
runPubSub_
    :: forall event es a
     . (Chan :> es, Clock :> es)
    => Eff (Pub event : Sub event : es) a -> Eff es a
runPubSub_ action = do
    (inChan, _) <- Chan.newChan @(UTCTime, event)
    let handlePub = interpret_ \case
            Publish event -> do
                timestamp <- Clock.currentTime
                Chan.writeChan inChan (timestamp, event)

        handleSub = interpret \env -> \case
            ListenWith onSubscribed listener -> localSeqUnlift env \unlift -> do
                chan <- Chan.dupChan inChan
                unlift onSubscribed
                forever do
                    (timestamp, event) <- Chan.readChan chan
                    unlift $ listener timestamp event

    handleSub . handlePub $ action