feat: produce OpenTelemetry traces with hs-opentelemetry

cabal.project

@@ -2,3 +2,9 @@ packages: postgrest.cabal
 tests: true
 package *
   ghc-options: -split-sections
+
+source-repository-package
+    type: git
+    location: https://github.com/develop7/hs-opentelemetry.git
+    tag: 90c424f5e01cc653cdc7ffe767c05bda55eb048a
+    subdir: sdk api propagators/b3 propagators/w3c propagators/datadog exporters/otlp utils/exceptions instrumentation/wai otlp

postgrest.cabal

@@ -72,6 +72,7 @@ library
                       PostgREST.Query.QueryBuilder
                       PostgREST.Query.SqlFragment
                       PostgREST.Query.Statements
+                      PostgREST.OpenTelemetry
                       PostgREST.Plan
                       PostgREST.Plan.CallPlan
                       PostgREST.Plan.MutatePlan
@@ -115,6 +116,9 @@ library
                     , hasql-transaction         >= 1.0.1 && < 1.1
                     , 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
+                    , hs-opentelemetry-utils-exceptions
                     , insert-ordered-containers >= 0.2.2 && < 0.3
                     , iproute                   >= 1.7.0 && < 1.8
                     , jose-jwt                  >= 0.9.6 && < 0.11
@@ -262,6 +266,7 @@ test-suite spec
                     , hasql-pool        >= 1.0.1 && < 1.1
                     , hasql-transaction >= 1.0.1 && < 1.1
                     , heredoc           >= 0.2 && < 0.3
+                    , hs-opentelemetry-sdk      >= 0.0.3.6 && < 0.0.4
                     , hspec             >= 2.3 && < 2.12
                     , hspec-wai         >= 0.10 && < 0.12
                     , hspec-wai-json    >= 0.10 && < 0.12

src/PostgREST/App.hs

@@ -41,24 +41,27 @@ import qualified PostgREST.Response   as Response
 import qualified PostgREST.Unix       as Unix (installSignalHandlers)
 
 import PostgREST.ApiRequest           (ApiRequest (..))
-import PostgREST.AppState             (AppState)
+import PostgREST.AppState             (AppState, getOTelTracer)
 import PostgREST.Auth                 (AuthResult (..))
 import PostgREST.Config               (AppConfig (..), LogLevel (..))
 import PostgREST.Config.PgVersion     (PgVersion (..))
-import PostgREST.Error                (Error)
+import PostgREST.Error                (Error (..))
 import PostgREST.Network              (resolveHost)
 import PostgREST.Observation          (Observation (..))
 import PostgREST.Response.Performance (ServerTiming (..),
                                        serverTimingHeader)
 import PostgREST.SchemaCache          (SchemaCache (..))
 import PostgREST.Version              (docsVersion, prettyVersion)
 
-import qualified Data.ByteString.Char8 as BS
-import qualified Data.List             as L
-import qualified Network.HTTP.Types    as HTTP
-import qualified Network.Socket        as NS
-import           Protolude             hiding (Handler)
-import           System.TimeIt         (timeItT)
+import qualified Data.ByteString.Char8             as BS
+import qualified Data.List                         as L
+import qualified Network.HTTP.Types                as HTTP
+import qualified Network.Socket                    as NS
+import           OpenTelemetry.Instrumentation.Wai (newOpenTelemetryWaiMiddleware)
+import           OpenTelemetry.Trace               (defaultSpanArguments)
+import           OpenTelemetry.Utils.Exceptions    (inSpanM)
+import           Protolude                         hiding (Handler)
+import           System.TimeIt                     (timeItT)
 
 type Handler = ExceptT Error
 
@@ -84,7 +87,9 @@ run appState = do
       host <- resolveHost $ AppState.getSocketREST appState
       observer $ AppServerPortObs (fromJust host) port
 
-  Warp.runSettingsSocket (serverSettings conf) (AppState.getSocketREST appState) app
+  oTelMWare <- newOpenTelemetryWaiMiddleware
+
+  Warp.runSettingsSocket (serverSettings conf) (AppState.getSocketREST appState) (oTelMWare app)
 
 serverSettings :: AppConfig -> Warp.Settings
 serverSettings AppConfig{..} =
@@ -102,27 +107,28 @@ postgrest logLevel appState connWorker =
   Logger.middleware logLevel Auth.getRole $
     -- fromJust can be used, because the auth middleware will **always** add
     -- some AuthResult to the vault.
-    \req respond -> case fromJust $ Auth.getResult req of
-      Left err -> respond $ Error.errorResponseFor err
-      Right authResult -> do
-        appConf <- AppState.getConfig appState -- the config must be read again because it can reload
-        maybeSchemaCache <- AppState.getSchemaCache appState
-        pgVer <- AppState.getPgVersion appState
-
-        let
-          eitherResponse :: IO (Either Error Wai.Response)
-          eitherResponse =
-            runExceptT $ postgrestResponse appState appConf maybeSchemaCache pgVer authResult req
-
-        response <- either Error.errorResponseFor identity <$> eitherResponse
-        -- Launch the connWorker when the connection is down.  The postgrest
-        -- function can respond successfully (with a stale schema cache) before
-        -- the connWorker is done.
-        when (isServiceUnavailable response) connWorker
-        resp <- do
-          delay <- AppState.getNextDelay appState
-          return $ addRetryHint delay response
-        respond resp
+    \req respond -> inSpanM (getOTelTracer appState) "respond" defaultSpanArguments $
+      case fromJust $ Auth.getResult req of
+        Left err -> respond $ Error.errorResponseFor err
+        Right authResult -> do
+          appConf <- AppState.getConfig appState -- the config must be read again because it can reload
+          maybeSchemaCache <- AppState.getSchemaCache appState
+          pgVer <- AppState.getPgVersion appState
+
+          let
+            eitherResponse :: IO (Either Error Wai.Response)
+            eitherResponse = inSpanM (getOTelTracer appState) "eitherResponse" defaultSpanArguments $
+              runExceptT $ postgrestResponse appState appConf maybeSchemaCache pgVer authResult req
+
+          response <- either Error.errorResponseFor identity <$> eitherResponse
+          -- Launch the connWorker when the connection is down.  The postgrest
+          -- function can respond successfully (with a stale schema cache) before
+          -- the connWorker is done.
+          when (isServiceUnavailable response) connWorker
+          resp <- do
+            delay <- AppState.getNextDelay appState
+            return $ addRetryHint delay response
+          respond resp
 
 postgrestResponse
   :: AppState.AppState

src/PostgREST/AppState.hs

@@ -16,6 +16,7 @@ module PostgREST.AppState
   , getJwtCache
   , getSocketREST
   , getSocketAdmin
+  , getOTelTracer
   , init
   , initSockets
   , initWithPool
@@ -76,6 +77,7 @@ import PostgREST.Unix                    (createAndBindDomainSocket)
 
 import Data.Streaming.Network (bindPortTCP, bindRandomPortTCP)
 import Data.String            (IsString (..))
+import OpenTelemetry.Trace    (Tracer)
 import Protolude
 
 data AuthResult = AuthResult
@@ -116,6 +118,8 @@ data AppState = AppState
   , stateObserver          :: ObservationHandler
   , stateLogger            :: Logger.LoggerState
   , stateMetrics           :: Metrics.MetricsState
+    -- | OpenTelemetry tracer
+  , oTelTracer             :: Tracer
   }
 
 -- | Schema cache status
@@ -126,8 +130,8 @@ data SchemaCacheStatus
 
 type AppSockets = (NS.Socket, Maybe NS.Socket)
 
-init :: AppConfig -> IO AppState
-init conf@AppConfig{configLogLevel, configDbPoolSize} = do
+init :: AppConfig -> Tracer -> IO AppState
+init conf@AppConfig{configLogLevel, configDbPoolSize} tracer = do
   loggerState  <- Logger.init
   metricsState <- Metrics.init configDbPoolSize
   let observer = liftA2 (>>) (Logger.observationLogger loggerState configLogLevel) (Metrics.observationMetrics metricsState)
@@ -136,11 +140,11 @@ init conf@AppConfig{configLogLevel, configDbPoolSize} = do
 
   pool <- initPool conf observer
   (sock, adminSock) <- initSockets conf
-  state' <- initWithPool (sock, adminSock) pool conf loggerState metricsState observer
+  state' <- initWithPool (sock, adminSock) pool conf loggerState metricsState tracer observer
   pure state' { stateSocketREST = sock, stateSocketAdmin = adminSock}
 
-initWithPool :: AppSockets -> SQL.Pool -> AppConfig -> Logger.LoggerState -> Metrics.MetricsState -> ObservationHandler -> IO AppState
-initWithPool (sock, adminSock) pool conf loggerState metricsState observer = do
+initWithPool :: AppSockets -> SQL.Pool ->  AppConfig -> Logger.LoggerState -> Metrics.MetricsState -> Tracer -> ObservationHandler -> IO AppState
+initWithPool (sock, adminSock) pool conf loggerState metricsState tracer observer = do
 
   appState <- AppState pool
     <$> newIORef minimumPgVersion -- assume we're in a supported version when starting, this will be corrected on a later step
@@ -159,6 +163,7 @@ initWithPool (sock, adminSock) pool conf loggerState metricsState observer = do
     <*> pure observer
     <*> pure loggerState
     <*> pure metricsState
+    <*> pure tracer
 
   deb <-
     let decisecond = 100000 in
@@ -325,6 +330,9 @@ getSocketREST = stateSocketREST
 getSocketAdmin :: AppState -> Maybe NS.Socket
 getSocketAdmin = stateSocketAdmin
 
+getOTelTracer :: AppState -> Tracer
+getOTelTracer = oTelTracer
+
 getMainThreadId :: AppState -> ThreadId
 getMainThreadId = stateMainThreadId
 

src/PostgREST/CLI.hs

@@ -16,29 +16,29 @@ import qualified Options.Applicative        as O
 
 import Text.Heredoc (str)
 
-import PostgREST.AppState    (AppState)
-import PostgREST.Config      (AppConfig (..))
-import PostgREST.Observation (Observation (..))
-import PostgREST.SchemaCache (querySchemaCache)
-import PostgREST.Version     (prettyVersion)
+import PostgREST.AppState      (AppState)
+import PostgREST.Config        (AppConfig (..))
+import PostgREST.Observation   (Observation (..))
+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
 
-
 main :: CLI -> IO ()
-main CLI{cliCommand, cliPath} = 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

src/PostgREST/OpenTelemetry.hs

@@ -0,0 +1,22 @@
+module PostgREST.OpenTelemetry (withTracer) where
+
+import OpenTelemetry.Attributes (emptyAttributes)
+import OpenTelemetry.Trace      (InstrumentationLibrary (..), Tracer,
+                                 initializeGlobalTracerProvider,
+                                 makeTracer, shutdownTracerProvider,
+                                 tracerOptions)
+import PostgREST.Version        (prettyVersion)
+import Protolude
+
+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
+            , librarySchemaUrl = ""
+            , libraryAttributes = emptyAttributes}

stack.yaml

@@ -10,6 +10,26 @@ nix:
   pure: false
 
 extra-deps:
-  - fuzzyset-0.2.4
-  - hasql-notifications-0.2.2.0
-  - hasql-pool-1.0.1
+  - fuzzyset-0.2.4@sha256:f1b6de8bf33277bf6255207541d65028f1f1ea93af5541b654c86b5674995485,1618
+  - hasql-pool-1.0.1@sha256:3cfb4c7153a6c536ac7e126c17723e6d26ee03794954deed2d72bcc826d05a40,2302
+  - hasql-notifications-0.2.2.0@sha256:a4e591ef3f06647b056567d3b66948c4a85371f05deb5434edb6ce190f7c845d,2021
+
+  - proto-lens-protoc-0.8.0.0@sha256:a146ee8c9af9e445ab05651e688deb0ff849357d320657d6cea5be33cb54b960,2235
+  - proto-lens-setup-0.4.0.7@sha256:acca0b04e033ea0a017f809d91a7dbc942e025ec6bc275fa21647352722c74cc,3122
+  - ghc-source-gen-0.4.4.1@sha256:6029be9ea2fae3ddb3a7a7e132df9b09aa32dc0cf3e31a5cd3d420d220823078,4236
+  - github: develop7/hs-opentelemetry
+    commit: 90c424f5e01cc653cdc7ffe767c05bda55eb048a
+    subdirs:
+      - sdk
+      - api
+      - propagators/b3
+      - propagators/w3c
+      - propagators/datadog
+      - exporters/otlp
+      - utils/exceptions
+      - instrumentation/wai
+      - otlp
+
+  - unix-compat-0.7.1@sha256:bd5bb4e04b2ed707f3e3466470a452354310389506cf0a7a73bf10e4d533f6d1,3619
+  - thread-utils-context-0.3.0.4@sha256:e763da1c6cab3b6d378fb670ca74aa9bf03c9b61b6fcf7628c56363fb0e3e71e,1671
+  - thread-utils-finalizers-0.1.1.0@sha256:24944b71d9f1d01695a5908b4a3b44838fab870883114a323336d537995e0a5b,1381