{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE OverloadedLists #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE TypeApplications #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}
module OpenTelemetry.Instrumentation.Persistent (
wrapSqlBackend,
wrapSqlBackend',
) where
import Control.Monad
import Control.Monad.IO.Class
import Control.Monad.Reader
import Data.Acquire.Internal
import Data.IORef
import qualified Data.HashMap.Strict as H
import Data.Maybe (fromMaybe)
import Data.Text (Text)
import qualified Data.Text as T
import qualified Data.Vault.Strict as Vault
import Database.Persist.Sql
import Database.Persist.SqlBackend (MkSqlBackendArgs (connRDBMS), emptySqlBackendHooks, getConnVault, getRDBMS, modifyConnVault, setConnHooks)
import Database.Persist.SqlBackend.Internal
import OpenTelemetry.Attributes (Attributes)
import OpenTelemetry.Common
import OpenTelemetry.Context
import OpenTelemetry.Context.ThreadLocal (adjustContext, getContext)
import OpenTelemetry.Resource
import OpenTelemetry.Trace.Core
import OpenTelemetry.Trace.Monad (MonadTracer (..))
import System.Clock
import System.IO.Unsafe (unsafePerformIO)
import UnliftIO.Exception
instance {-# OVERLAPS #-} (MonadTracer m) => MonadTracer (ReaderT SqlBackend m) where
getTracer :: ReaderT SqlBackend m Tracer
getTracer = forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift forall (m :: * -> *). MonadTracer m => m Tracer
OpenTelemetry.Trace.Monad.getTracer
instance {-# OVERLAPS #-} (MonadTracer m) => MonadTracer (ReaderT SqlReadBackend m) where
getTracer :: ReaderT SqlReadBackend m Tracer
getTracer = forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift forall (m :: * -> *). MonadTracer m => m Tracer
OpenTelemetry.Trace.Monad.getTracer
instance {-# OVERLAPS #-} (MonadTracer m) => MonadTracer (ReaderT SqlWriteBackend m) where
getTracer :: ReaderT SqlWriteBackend m Tracer
getTracer = forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift forall (m :: * -> *). MonadTracer m => m Tracer
OpenTelemetry.Trace.Monad.getTracer
originalConnectionKey :: Vault.Key SqlBackend
originalConnectionKey :: Key SqlBackend
originalConnectionKey = forall a. IO a -> a
unsafePerformIO forall a. IO (Key a)
Vault.newKey
{-# NOINLINE originalConnectionKey #-}
insertOriginalConnection :: SqlBackend -> SqlBackend -> SqlBackend
insertOriginalConnection :: SqlBackend -> SqlBackend -> SqlBackend
insertOriginalConnection SqlBackend
conn SqlBackend
original = (Vault -> Vault) -> SqlBackend -> SqlBackend
modifyConnVault (forall a. Key a -> a -> Vault -> Vault
Vault.insert Key SqlBackend
originalConnectionKey SqlBackend
original) SqlBackend
conn
lookupOriginalConnection :: SqlBackend -> Maybe SqlBackend
lookupOriginalConnection :: SqlBackend -> Maybe SqlBackend
lookupOriginalConnection = forall a. Key a -> Vault -> Maybe a
Vault.lookup Key SqlBackend
originalConnectionKey forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall backend (m :: * -> *).
(BackendCompatible SqlBackend backend, MonadReader backend m) =>
m Vault
getConnVault
connectionLevelAttributesKey :: Vault.Key (H.HashMap Text Attribute)
connectionLevelAttributesKey :: Key (HashMap Text Attribute)
connectionLevelAttributesKey = forall a. IO a -> a
unsafePerformIO forall a. IO (Key a)
Vault.newKey
{-# NOINLINE connectionLevelAttributesKey #-}
wrapSqlBackend ::
MonadIO m =>
H.HashMap Text Attribute ->
SqlBackend ->
m SqlBackend
wrapSqlBackend :: forall (m :: * -> *).
MonadIO m =>
HashMap Text Attribute -> SqlBackend -> m SqlBackend
wrapSqlBackend HashMap Text Attribute
attrs SqlBackend
conn_ = do
TracerProvider
tp <- forall (m :: * -> *). MonadIO m => m TracerProvider
getGlobalTracerProvider
forall (m :: * -> *).
MonadIO m =>
TracerProvider
-> HashMap Text Attribute -> SqlBackend -> m SqlBackend
wrapSqlBackend' TracerProvider
tp HashMap Text Attribute
attrs SqlBackend
conn_
wrapSqlBackend' :: MonadIO m =>
TracerProvider ->
H.HashMap Text Attribute ->
SqlBackend ->
m SqlBackend
wrapSqlBackend' :: forall (m :: * -> *).
MonadIO m =>
TracerProvider
-> HashMap Text Attribute -> SqlBackend -> m SqlBackend
wrapSqlBackend' TracerProvider
tp HashMap Text Attribute
attrs SqlBackend
conn_ = do
let conn :: SqlBackend
conn = forall a. a -> Maybe a -> a
Data.Maybe.fromMaybe SqlBackend
conn_ (SqlBackend -> Maybe SqlBackend
lookupOriginalConnection SqlBackend
conn_)
IORef (Maybe Span)
connParentSpan <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. a -> IO (IORef a)
newIORef forall a. Maybe a
Nothing
IORef (Maybe Span)
connSpanInFlight <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. a -> IO (IORef a)
newIORef forall a. Maybe a
Nothing
let t :: Tracer
t = TracerProvider -> InstrumentationLibrary -> TracerOptions -> Tracer
makeTracer TracerProvider
tp InstrumentationLibrary
"hs-opentelemetry-persistent" TracerOptions
tracerOptions
let hooks :: SqlBackendHooks
hooks =
SqlBackendHooks
emptySqlBackendHooks
{ hookGetStatement :: SqlBackend -> Text -> Statement -> IO Statement
hookGetStatement = \SqlBackend
conn Text
sql Statement
stmt -> do
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$
Statement
{ stmtQuery :: forall (m :: * -> *).
MonadIO m =>
[PersistValue] -> Acquire (ConduitM () [PersistValue] m ())
stmtQuery = \[PersistValue]
ps -> do
Context
ctxt <- forall (m :: * -> *). MonadIO m => m Context
getContext
let spanCreator :: IO (Maybe Span, Span)
spanCreator = do
Span
s <-
forall (m :: * -> *).
(MonadIO m, HasCallStack) =>
Tracer -> Context -> Text -> SpanArguments -> m Span
createSpan
Tracer
t
Context
ctxt
Text
sql
(SpanArguments
defaultSpanArguments {kind :: SpanKind
kind = SpanKind
Client, attributes :: HashMap Text Attribute
attributes = forall k v.
(Eq k, Hashable k) =>
k -> v -> HashMap k v -> HashMap k v
H.insert Text
"db.statement" (forall a. ToAttribute a => a -> Attribute
toAttribute Text
sql) HashMap Text Attribute
attrs})
forall (m :: * -> *). MonadIO m => (Context -> Context) -> m ()
adjustContext (Span -> Context -> Context
insertSpan Span
s)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Context -> Maybe Span
lookupSpan Context
ctxt, Span
s)
spanCleanup :: (Maybe Span, Span) -> m ()
spanCleanup (Maybe Span
parent, Span
s) = do
Span
s forall (m :: * -> *). MonadIO m => Span -> Maybe Timestamp -> m ()
`endSpan` forall a. Maybe a
Nothing
forall (m :: * -> *). MonadIO m => (Context -> Context) -> m ()
adjustContext forall a b. (a -> b) -> a -> b
$ \Context
ctx ->
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (Context -> Context
removeSpan Context
ctx) (Span -> Context -> Context
`insertSpan` Context
ctx) Maybe Span
parent
(Maybe Span
p, Span
child) <- forall a. IO a -> (a -> IO ()) -> Acquire a
mkAcquire IO (Maybe Span, Span)
spanCreator forall {m :: * -> *}. MonadIO m => (Maybe Span, Span) -> m ()
spanCleanup
forall (m :: * -> *). MonadIO m => Span -> SqlBackend -> m ()
annotateBasics Span
child SqlBackend
conn
case Statement
-> forall (m :: * -> *).
MonadIO m =>
[PersistValue] -> Acquire (ConduitM () [PersistValue] m ())
stmtQuery Statement
stmt [PersistValue]
ps of
Acquire (forall b. IO b -> IO b)
-> IO (Allocated (ConduitM () [PersistValue] m ()))
stmtQueryAcquireF -> forall a.
((forall b. IO b -> IO b) -> IO (Allocated a)) -> Acquire a
Acquire forall a b. (a -> b) -> a -> b
$ \forall b. IO b -> IO b
f ->
forall (m :: * -> *) a.
MonadUnliftIO m =>
(SomeException -> m a) -> m a -> m a
handleAny
( \(SomeException e
err) -> do
forall (m :: * -> *) e.
(MonadIO m, Exception e) =>
Span -> HashMap Text Attribute -> Maybe Timestamp -> e -> m ()
recordException Span
child [(Text
"exception.escaped", forall a. ToAttribute a => a -> Attribute
toAttribute Bool
True)] forall a. Maybe a
Nothing e
err
forall (m :: * -> *). MonadIO m => Span -> Maybe Timestamp -> m ()
endSpan Span
child forall a. Maybe a
Nothing
forall (m :: * -> *) e a. (MonadIO m, Exception e) => e -> m a
throwIO e
err
)
((forall b. IO b -> IO b)
-> IO (Allocated (ConduitM () [PersistValue] m ()))
stmtQueryAcquireF forall b. IO b -> IO b
f)
, stmtExecute :: [PersistValue] -> IO Int64
stmtExecute = \[PersistValue]
ps -> do
forall (m :: * -> *) a.
(MonadUnliftIO m, HasCallStack) =>
Tracer -> Text -> SpanArguments -> (Span -> m a) -> m a
inSpan' Tracer
t Text
sql (SpanArguments
defaultSpanArguments {kind :: SpanKind
kind = SpanKind
Client, attributes :: HashMap Text Attribute
attributes = forall k v.
(Eq k, Hashable k) =>
k -> v -> HashMap k v -> HashMap k v
H.insert Text
"db.statement" (forall a. ToAttribute a => a -> Attribute
toAttribute Text
sql) HashMap Text Attribute
attrs}) forall a b. (a -> b) -> a -> b
$ \Span
s -> do
forall (m :: * -> *). MonadIO m => Span -> SqlBackend -> m ()
annotateBasics Span
s SqlBackend
conn
Statement -> [PersistValue] -> IO Int64
stmtExecute Statement
stmt [PersistValue]
ps
, stmtReset :: IO ()
stmtReset = Statement -> IO ()
stmtReset Statement
stmt
, stmtFinalize :: IO ()
stmtFinalize = Statement -> IO ()
stmtFinalize Statement
stmt
}
}
conn' :: SqlBackend
conn' =
SqlBackend
conn
{ connHooks :: SqlBackendHooks
connHooks = SqlBackendHooks
hooks
, connBegin :: (Text -> IO Statement) -> Maybe IsolationLevel -> IO ()
connBegin = \Text -> IO Statement
f Maybe IsolationLevel
mIso -> do
Context
ctxt <- forall (m :: * -> *). MonadIO m => m Context
getContext
Span
s <- forall (m :: * -> *).
(MonadIO m, HasCallStack) =>
Tracer -> Context -> Text -> SpanArguments -> m Span
createSpan Tracer
t Context
ctxt Text
"transaction" (SpanArguments
defaultSpanArguments {kind :: SpanKind
kind = SpanKind
Client, attributes :: HashMap Text Attribute
attributes = HashMap Text Attribute
attrs})
forall (m :: * -> *). MonadIO m => Span -> SqlBackend -> m ()
annotateBasics Span
s SqlBackend
conn
forall a. IORef a -> a -> IO ()
writeIORef IORef (Maybe Span)
connSpanInFlight (forall a. a -> Maybe a
Just Span
s)
forall a. IORef a -> a -> IO ()
writeIORef IORef (Maybe Span)
connParentSpan (Context -> Maybe Span
lookupSpan Context
ctxt)
forall (m :: * -> *). MonadIO m => (Context -> Context) -> m ()
adjustContext (Span -> Context -> Context
insertSpan Span
s)
case Maybe IsolationLevel
mIso of
Maybe IsolationLevel
Nothing -> forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
Just IsolationLevel
iso -> forall (m :: * -> *) a.
(MonadIO m, ToAttribute a) =>
Span -> Text -> a -> m ()
addAttribute Span
s Text
"db.transaction.isolation" forall a b. (a -> b) -> a -> b
$ case IsolationLevel
iso of
IsolationLevel
ReadUncommitted -> Text
"read uncommitted" :: Text
IsolationLevel
ReadCommitted -> Text
"read committed"
IsolationLevel
RepeatableRead -> Text
"repeatable read"
IsolationLevel
Serializable -> Text
"serializable"
SqlBackend
-> (Text -> IO Statement) -> Maybe IsolationLevel -> IO ()
connBegin SqlBackend
conn Text -> IO Statement
f Maybe IsolationLevel
mIso
, connCommit :: (Text -> IO Statement) -> IO ()
connCommit = \Text -> IO Statement
f -> do
Maybe Span
spanInFlight <- forall a. IORef a -> IO a
readIORef IORef (Maybe Span)
connSpanInFlight
Maybe Span
parentSpan <- forall a. IORef a -> IO a
readIORef IORef (Maybe Span)
connParentSpan
let act :: IO ()
act = do
(Timestamp TimeSpec
tsStart) <- forall (m :: * -> *). MonadIO m => m Timestamp
getTimestamp
Either SomeException ()
result <- forall (m :: * -> *) a.
MonadUnliftIO m =>
m a -> m (Either SomeException a)
tryAny forall a b. (a -> b) -> a -> b
$ SqlBackend -> (Text -> IO Statement) -> IO ()
connCommit SqlBackend
conn Text -> IO Statement
f
(Timestamp TimeSpec
tsEnd) <- forall (m :: * -> *). MonadIO m => m Timestamp
getTimestamp
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ Maybe Span
spanInFlight forall a b. (a -> b) -> a -> b
$ \Span
s -> do
forall (m :: * -> *).
MonadIO m =>
Span -> HashMap Text Attribute -> m ()
addAttributes
Span
s
[ (Text
"db.transaction.outcome", forall a. ToAttribute a => a -> Attribute
toAttribute (Text
"committed" :: Text))
, (Text
"db.transaction.commit_duration_ns", forall a. ToAttribute a => a -> Attribute
toAttribute forall a b. (a -> b) -> a -> b
$ forall a b. (Integral a, Num b) => a -> b
fromIntegral @Integer @Int forall a b. (a -> b) -> a -> b
$ TimeSpec -> Integer
toNanoSecs (TimeSpec -> TimeSpec -> TimeSpec
diffTimeSpec TimeSpec
tsStart TimeSpec
tsEnd) forall a. Integral a => a -> a -> a
`div` Integer
1000)
]
forall (m :: * -> *). MonadIO m => Span -> Maybe Timestamp -> m ()
endSpan Span
s forall a. Maybe a
Nothing
case Either SomeException ()
result of
Left (SomeException e
err) -> do
forall (m :: * -> *) e.
(MonadIO m, Exception e) =>
Span -> HashMap Text Attribute -> Maybe Timestamp -> e -> m ()
recordException Span
s [(Text
"exception.escaped", forall a. ToAttribute a => a -> Attribute
toAttribute Bool
True)] forall a. Maybe a
Nothing e
err
forall (m :: * -> *) e a. (MonadIO m, Exception e) => e -> m a
throwIO e
err
Right ()
_ -> forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
IO ()
act forall (m :: * -> *) a b. MonadUnliftIO m => m a -> m b -> m a
`finally` do
forall (m :: * -> *). MonadIO m => (Context -> Context) -> m ()
adjustContext forall a b. (a -> b) -> a -> b
$ \Context
ctx ->
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (Context -> Context
removeSpan Context
ctx) (Span -> Context -> Context
`insertSpan` Context
ctx) Maybe Span
parentSpan
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ Maybe Span
spanInFlight forall a b. (a -> b) -> a -> b
$ \Span
s -> forall (m :: * -> *). MonadIO m => Span -> Maybe Timestamp -> m ()
endSpan Span
s forall a. Maybe a
Nothing
, connRollback :: (Text -> IO Statement) -> IO ()
connRollback = \Text -> IO Statement
f -> do
Maybe Span
spanInFlight <- forall a. IORef a -> IO a
readIORef IORef (Maybe Span)
connSpanInFlight
Maybe Span
parentSpan <- forall a. IORef a -> IO a
readIORef IORef (Maybe Span)
connParentSpan
let act :: IO ()
act = do
(Timestamp TimeSpec
tsStart) <- forall (m :: * -> *). MonadIO m => m Timestamp
getTimestamp
Either SomeException ()
result <- forall (m :: * -> *) a.
MonadUnliftIO m =>
m a -> m (Either SomeException a)
tryAny forall a b. (a -> b) -> a -> b
$ SqlBackend -> (Text -> IO Statement) -> IO ()
connRollback SqlBackend
conn Text -> IO Statement
f
e :: Timestamp
e@(Timestamp TimeSpec
tsEnd) <- forall (m :: * -> *). MonadIO m => m Timestamp
getTimestamp
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ Maybe Span
spanInFlight forall a b. (a -> b) -> a -> b
$ \Span
s -> do
forall (m :: * -> *).
MonadIO m =>
Span -> HashMap Text Attribute -> m ()
addAttributes
Span
s
[ (Text
"db.transaction.outcome", forall a. ToAttribute a => a -> Attribute
toAttribute (Text
"rolled back" :: Text))
, (Text
"db.transaction.commit_duration_microseconds", forall a. ToAttribute a => a -> Attribute
toAttribute forall a b. (a -> b) -> a -> b
$ forall a b. (Integral a, Num b) => a -> b
fromIntegral @Integer @Int forall a b. (a -> b) -> a -> b
$ TimeSpec -> Integer
toNanoSecs (TimeSpec -> TimeSpec -> TimeSpec
diffTimeSpec TimeSpec
tsStart TimeSpec
tsEnd forall a. Integral a => a -> a -> a
`div` TimeSpec
1000))
]
forall (m :: * -> *). MonadIO m => Span -> Maybe Timestamp -> m ()
endSpan Span
s (forall a. a -> Maybe a
Just Timestamp
e)
case Either SomeException ()
result of
Left (SomeException e
err) -> do
forall (m :: * -> *) e.
(MonadIO m, Exception e) =>
Span -> HashMap Text Attribute -> Maybe Timestamp -> e -> m ()
recordException Span
s [(Text
"exception.escaped", forall a. ToAttribute a => a -> Attribute
toAttribute Bool
True)] forall a. Maybe a
Nothing e
err
forall (m :: * -> *) e a. (MonadIO m, Exception e) => e -> m a
throwIO e
err
Right ()
_ -> forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
IO ()
act forall (m :: * -> *) a b. MonadUnliftIO m => m a -> m b -> m a
`finally` do
forall (m :: * -> *). MonadIO m => (Context -> Context) -> m ()
adjustContext forall a b. (a -> b) -> a -> b
$ \Context
ctx ->
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (Context -> Context
removeSpan Context
ctx) (Span -> Context -> Context
`insertSpan` Context
ctx) Maybe Span
parentSpan
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ Maybe Span
spanInFlight forall a b. (a -> b) -> a -> b
$ \Span
s -> forall (m :: * -> *). MonadIO m => Span -> Maybe Timestamp -> m ()
endSpan Span
s forall a. Maybe a
Nothing
,
connClose :: IO ()
connClose = do
forall (m :: * -> *) a.
(MonadUnliftIO m, HasCallStack) =>
Tracer -> Text -> SpanArguments -> (Span -> m a) -> m a
inSpan' Tracer
t Text
"close connection" (SpanArguments
defaultSpanArguments {kind :: SpanKind
kind = SpanKind
Client, attributes :: HashMap Text Attribute
attributes = HashMap Text Attribute
attrs}) forall a b. (a -> b) -> a -> b
$ \Span
s -> do
forall (m :: * -> *). MonadIO m => Span -> SqlBackend -> m ()
annotateBasics Span
s SqlBackend
conn
SqlBackend -> IO ()
connClose SqlBackend
conn
}
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ SqlBackend -> SqlBackend -> SqlBackend
insertOriginalConnection SqlBackend
conn' SqlBackend
conn
annotateBasics :: (MonadIO m) => Span -> SqlBackend -> m ()
annotateBasics :: forall (m :: * -> *). MonadIO m => Span -> SqlBackend -> m ()
annotateBasics Span
span SqlBackend
conn = do
forall (m :: * -> *).
MonadIO m =>
Span -> HashMap Text Attribute -> m ()
addAttributes
Span
span
[ (Text
"db.system", forall a. ToAttribute a => a -> Attribute
toAttribute forall a b. (a -> b) -> a -> b
$ forall backend (m :: * -> *).
(BackendCompatible SqlBackend backend, MonadReader backend m) =>
m Text
getRDBMS SqlBackend
conn)
]