rewrite over

postgrest.cabal

@@ -58,6 +58,7 @@ library
  PostgREST.Query.QueryBuilder
  PostgREST.Query.SqlFragment
  PostgREST.Query.Statements
+ PostgREST.OpenTelemetry
  PostgREST.Plan
  PostgREST.Plan.CallPlan
  PostgREST.Plan.MutatePlan
@@ -103,6 +104,7 @@ library
  , heredoc >= 0.2 && < 0.3
  , http-types >= 0.12.2 && < 0.13
  , hs-opentelemetry-sdk >= 0.0.3.6 && < 0.0.4
+ , hs-opentelemetry-instrumentation-wai
  , insert-ordered-containers >= 0.2.2 && < 0.3
  , interpolatedstring-perl6 >= 1 && < 1.1
  , jose >= 0.8.5.1 && < 0.12

src/PostgREST/AppState.hs

@@ -16,6 +16,7 @@ module PostgREST.AppState
  , getJwtCache
  , getSocketREST
  , getSocketAdmin
+ , getOTelTracer
  , init
  , initSockets
  , initWithPool
@@ -74,6 +75,7 @@ import PostgREST.Unix (createAndBindDomainSocket)
 
 import Data.Streaming.Network (bindPortTCP, bindRandomPortTCP)
 import Data.String (IsString (..))
+import OpenTelemetry.Trace (Tracer)
 import Protolude
 
 data AuthResult = AuthResult
@@ -112,19 +114,21 @@ data AppState = AppState
  , stateSocketREST :: NS.Socket
  -- | Network socket for the admin UI
  , stateSocketAdmin :: Maybe NS.Socket
+ -- | OpenTelemetry tracer
+ , oTelTracer :: Tracer
  }
 
 type AppSockets = (NS.Socket, Maybe NS.Socket)
 
-init :: AppConfig -> IO AppState
-init conf = do
+init :: AppConfig -> Tracer -> IO AppState
+init conf tracer = do
  pool <- initPool conf
  (sock, adminSock) <- initSockets conf
- state' <- initWithPool (sock, adminSock) pool conf
- pure state' { stateSocketREST = sock, stateSocketAdmin = adminSock }
+ state' <- initWithPool (sock, adminSock) pool tracer conf
+ pure state' { stateSocketREST = sock, stateSocketAdmin = adminSock}
 
-initWithPool :: AppSockets -> SQL.Pool -> AppConfig -> IO AppState
-initWithPool (sock, adminSock) pool conf = do
+initWithPool :: AppSockets -> SQL.Pool -> Tracer -> AppConfig -> IO AppState
+initWithPool (sock, adminSock) pool tracer conf = do
  appState <- AppState pool
  <$> newIORef minimumPgVersion -- assume we're in a supported version when starting, this will be corrected on a later step
  <*> newIORef Nothing
@@ -140,6 +144,7 @@ initWithPool (sock, adminSock) pool conf = do
  <*> C.newCache Nothing
  <*> pure sock
  <*> pure adminSock
+ <*> pure tracer
 
 
  debLogTimeout <-
@@ -268,6 +273,9 @@ getSocketREST = stateSocketREST
 getSocketAdmin :: AppState -> Maybe NS.Socket
 getSocketAdmin = stateSocketAdmin
 
+getOTelTracer :: AppState -> Tracer 
+getOTelTracer = oTelTracer
+
 -- | Log to stderr with local time
 logWithZTime :: AppState -> Text -> IO ()
 logWithZTime appState txt = do

src/PostgREST/CLI.hs

@@ -17,42 +17,35 @@ import qualified Options.Applicative as O
 import Data.Text.IO (hPutStrLn)
 import Text.Heredoc (str)
 
-import PostgREST.AppState (AppState)
-import PostgREST.Config (AppConfig (..))
-import PostgREST.SchemaCache (querySchemaCache)
-import PostgREST.Version (prettyVersion)
+import PostgREST.AppState (AppState)
+import PostgREST.Config (AppConfig (..))
+import PostgREST.OpenTelemetry (withTracer)
+import PostgREST.SchemaCache (querySchemaCache)
+import PostgREST.Version (prettyVersion)
 
 import qualified PostgREST.App as App
 import qualified PostgREST.AppState as AppState
 import qualified PostgREST.Config as Config
 
 import Protolude hiding (hPutStrLn)
 
-import OpenTelemetry.Trace
-
 main :: CLI -> IO ()
-main CLI{cliCommand, cliPath} = withTracer $ \_tracer -> do
+main CLI{cliCommand, cliPath} = withTracer "PostgREST" $ \tracer -> do
  conf@AppConfig{..} <-
  either panic identity <$> Config.readAppConfig mempty cliPath Nothing mempty mempty
 
  -- Per https://github.com/PostgREST/postgrest/issues/268, we want to
  -- explicitly close the connections to PostgreSQL on shutdown.
  -- 'AppState.destroy' takes care of that.
  bracket
- (AppState.init conf)
+ (AppState.init conf tracer)
  AppState.destroy
  (\appState -> case cliCommand of
  CmdDumpConfig -> do
  when configDbConfig $ AppState.reReadConfig True appState
  putStr . Config.toText =<< AppState.getConfig appState
  CmdDumpSchema -> putStrLn =<< dumpSchema appState
  CmdRun -> App.run appState)
- where
- withTracer :: ((TracerOptions -> Tracer) -> IO c) -> IO c 
- withTracer f = bracket 
- initializeGlobalTracerProvider
- shutdownTracerProvider
- (\tracerProvider -> f $ makeTracer tracerProvider "PostgREST")
 
 -- | Dump SchemaCache schema to JSON
 dumpSchema :: AppState -> IO LBS.ByteString

src/PostgREST/OpenTelemetry.hs

@@ -0,0 +1,128 @@
+{-# LANGUAGE ScopedTypeVariables, OverloadedLists #-}
+module PostgREST.OpenTelemetry (withTracer) where
+
+import OpenTelemetry.Trace (InstrumentationLibrary (..), Tracer,
+ initializeGlobalTracerProvider,
+ makeTracer, shutdownTracerProvider,
+ tracerOptions, SpanArguments (..), Span, createSpanWithoutCallStack, setStatus, SpanStatus (..), recordException, endSpan)
+import PostgREST.Version (prettyVersion)
+import qualified Data.HashMap.Strict as H
+import Protolude
+import OpenTelemetry.Attributes (Attribute, ToAttribute (..))
+import qualified Data.Text as T
+import Data.String (String)
+import GHC.Stack
+import OpenTelemetry.Context.ThreadLocal (getContext, adjustContext)
+import OpenTelemetry.Context (insertSpan, lookupSpan, removeSpan)
+import qualified Control.Exception as EUnsafe
+
+
+withTracer :: Text -> (Tracer -> IO c) -> IO c
+withTracer label f = bracket
+ initializeGlobalTracerProvider
+ shutdownTracerProvider
+ (\tracerProvider -> f $ makeTracer tracerProvider instrumentationLibrary tracerOptions)
+ where
+ instrumentationLibrary = InstrumentationLibrary {libraryName = label, libraryVersion = decodeUtf8 prettyVersion}
+
+
+ownCodeAttributes :: (HasCallStack) => H.HashMap Text Attribute
+ownCodeAttributes = case getCallStack callStack of
+ _ : caller : _ -> srcAttributes caller
+ _ -> mempty
+
+
+callerAttributes :: (HasCallStack) => H.HashMap Text Attribute
+callerAttributes = case getCallStack callStack of
+ _ : _ : caller : _ -> srcAttributes caller
+ _ -> mempty
+
+
+srcAttributes :: (String, GHC.Stack.SrcLoc) -> H.HashMap Text Attribute
+srcAttributes (fn, loc) = H.fromList
+ [ ("code.function", toAttribute $ T.pack fn)
+ , ("code.namespace", toAttribute $ T.pack $ srcLocModule loc)
+ , ("code.filepath", toAttribute $ T.pack $ srcLocFile loc)
+ , ("code.lineno", toAttribute $ srcLocStartLine loc)
+ , ("code.package", toAttribute $ T.pack $ srcLocPackage loc)
+ ]
+
+{- | The simplest function for annotating code with trace information.
+
+ @since 0.0.1.0
+-}
+inSpan
+ :: (MonadIO m, HasCallStack)
+ => Tracer
+ -> Text
+ -- ^ The name of the span. This may be updated later via 'updateName'
+ -> SpanArguments
+ -- ^ Additional options for creating the span, such as 'SpanKind',
+ -- span links, starting attributes, etc.
+ -> m a
+ -- ^ The action to perform. 'inSpan' will record the time spent on the
+ -- action without forcing strict evaluation of the result. Any uncaught
+ -- exceptions will be recorded and rethrown.
+ -> m a
+inSpan t n args m = inSpan'' t n (args {attributes = H.union (attributes args) callerAttributes}) (const m)
+
+
+inSpan'
+ :: (MonadIO m, HasCallStack)
+ => Tracer
+ -> Text
+ -- ^ The name of the span. This may be updated later via 'updateName'
+ -> SpanArguments
+ -> (Span -> m a)
+ -> m a
+inSpan' t n args = inSpan'' t n (args {attributes = H.union (attributes args) callerAttributes})
+
+
+inSpan''
+ :: (MonadIO m, HasCallStack)
+ => Tracer
+ -> Text
+ -- ^ The name of the span. This may be updated later via 'updateName'
+ -> SpanArguments
+ -> (Span -> m a)
+ -> m a
+inSpan'' t n args f = do
+ bracketError
+ ( liftIO $ do
+ ctx <- getContext
+ s <- createSpanWithoutCallStack t ctx n args
+ adjustContext (insertSpan s)
+ pure (lookupSpan ctx, s)
+ )
+ ( \e (parent, s) -> liftIO $ do
+ forM_ e $ \(SomeException inner) -> do
+ setStatus s $ Error $ T.pack $ displayException inner
+ recordException s [("exception.escaped", toAttribute True)] Nothing inner
+ endSpan s Nothing
+ adjustContext $ \ctx ->
+ maybe (removeSpan ctx) (`insertSpan` ctx) parent
+ )
+ (\(_, s) -> f s)
+
+{- | Like 'Context.Exception.bracket', but provides the @after@ function with information about
+ uncaught exceptions.
+
+ @since 0.1.0.0
+-}
+bracketError :: (MonadIO m) => m a -> (Maybe SomeException -> a -> m b) -> (a -> m c) -> m c
+bracketError before after thing = withRunInIO $ \run -> EUnsafe.mask $ \restore -> do
+ x <- run before
+ res1 <- EUnsafe.try $ restore $ run $ thing x
+ case res1 of
+ Left (e1 :: SomeException) -> do
+ -- explicitly ignore exceptions from after. We know that
+ -- no async exceptions were thrown there, so therefore
+ -- the stronger exception must come from thing
+ --
+ -- https://github.com/fpco/safe-exceptions/issues/2
+ _ :: Either SomeException b <-
+ EUnsafe.try $ EUnsafe.uninterruptibleMask_ $ run $ after (Just e1) x
+ EUnsafe.throwIO e1
+ Right y -> do
+ _ <- EUnsafe.uninterruptibleMask_ $ run $ after Nothing x
+ return y

stack.yaml

@@ -31,5 +31,6 @@ extra-deps:
  - 'propagators/b3'
  - 'propagators/w3c'
  - otlp
+ - 'instrumentation/wai'
  - thread-utils-context-0.3.0.4@sha256:e763da1c6cab3b6d378fb670ca74aa9bf03c9b61b6fcf7628c56363fb0e3e71e,1671
  - thread-utils-finalizers-0.1.1.0@sha256:24944b71d9f1d01695a5908b4a3b44838fab870883114a323336d537995e0a5b,1381

stack.yaml.lock

@@ -124,6 +124,19 @@ packages:
  original:
  subdir: otlp
  url: https://github.com/iand675/hs-opentelemetry/archive/54251bd15565e3a07cda0cecfa544d209ae63df6.tar.gz
+- completed:
+ name: hs-opentelemetry-instrumentation-wai
+ pantry-tree:
+ sha256: 8caff8dcafdc6503610a30903f566be6f5125bd3fa38ae71c7fc1255e75f5db5
+ size: 513
+ sha256: eba9c66b5e90e4b4f4a90119053a75b68c0901454a992b463f17192600d034a9
+ size: 357729
+ subdir: instrumentation/wai
+ url: https://github.com/iand675/hs-opentelemetry/archive/54251bd15565e3a07cda0cecfa544d209ae63df6.tar.gz
+ version: 0.1.0.0
+ original:
+ subdir: instrumentation/wai
+ url: https://github.com/iand675/hs-opentelemetry/archive/54251bd15565e3a07cda0cecfa544d209ae63df6.tar.gz
 - completed:
  hackage: thread-utils-context-0.3.0.4@sha256:e763da1c6cab3b6d378fb670ca74aa9bf03c9b61b6fcf7628c56363fb0e3e71e,1671
  pantry-tree:

test/spec/Main.hs

@@ -10,6 +10,7 @@ import Test.Hspec
 import PostgREST.App (postgrest)
 import PostgREST.Config (AppConfig (..))
 import PostgREST.Config.Database (queryPgVersion)
+import PostgREST.OpenTelemetry (withTracer)
 import PostgREST.SchemaCache (querySchemaCache)
 import Protolude hiding (toList, toS)
 import SpecHelper
@@ -78,16 +79,16 @@ main = do
 
  let
  -- For tests that run with the same refSchemaCache
- app config = do
- appState <- AppState.initWithPool sockets pool config
+ app config = withTracer "PostgREST.Spec" $ \tracer -> do
+ appState <- AppState.initWithPool sockets pool tracer config
  AppState.putPgVersion appState actualPgVersion
  AppState.putSchemaCache appState (Just baseSchemaCache)
  return ((), postgrest config appState $ pure ())
 
  -- For tests that run with a different SchemaCache(depends on configSchemas)
- appDbs config = do
+ appDbs config = withTracer "PostgREST.Spec" $ \tracer -> do
  customSchemaCache <- loadSchemaCache pool config
- appState <- AppState.initWithPool sockets pool config
+ appState <- AppState.initWithPool sockets pool tracer config
  AppState.putPgVersion appState actualPgVersion
  AppState.putSchemaCache appState (Just customSchemaCache)
  return ((), postgrest config appState $ pure ())