From 53163de836deb795f725995a4e52b81739a1c5eb Mon Sep 17 00:00:00 2001
From: Profpatsch <mail@profpatsch.de>
Date: Mon, 13 May 2024 12:32:56 +0200
Subject: feat(users/Profpatsch/MonadPostgres): time formatting with event

Instead of opening a separate span, we just add events for start and
end of formatting.

Change-Id: I26f6792dfdcd23c01cff415fa0f436d6a22d93fe
Reviewed-on: https://cl.tvl.fyi/c/depot/+/11655
Autosubmit: Profpatsch <mail@profpatsch.de>
Reviewed-by: Profpatsch <mail@profpatsch.de>
Tested-by: BuildkiteCI
---
 .../my-prelude/src/Postgres/MonadPostgres.hs       | 43 ++++++++++++++++++++--
 1 file changed, 40 insertions(+), 3 deletions(-)

(limited to 'users/Profpatsch/my-prelude/src')

diff --git a/users/Profpatsch/my-prelude/src/Postgres/MonadPostgres.hs b/users/Profpatsch/my-prelude/src/Postgres/MonadPostgres.hs
index e5205f2651d4..57daeb3aceef 100644
--- a/users/Profpatsch/my-prelude/src/Postgres/MonadPostgres.hs
+++ b/users/Profpatsch/my-prelude/src/Postgres/MonadPostgres.hs
@@ -26,6 +26,7 @@ import Data.List qualified as List
 import Data.Pool (Pool)
 import Data.Pool qualified as Pool
 import Data.Text qualified as Text
+import Data.Time (getCurrentTime)
 import Data.Typeable (Typeable)
 import Database.PostgreSQL.Simple (Connection, FormatError, FromRow, Query, QueryError, ResultError, SqlError, ToRow)
 import Database.PostgreSQL.Simple qualified as PG
@@ -37,6 +38,7 @@ import Database.PostgreSQL.Simple.Types (Query (..))
 import GHC.IO.Handle (Handle)
 import GHC.Records (getField)
 import Label
+import OpenTelemetry.Trace.Core (NewEvent (newEventName))
 import OpenTelemetry.Trace.Core qualified as Otel hiding (inSpan, inSpan')
 import OpenTelemetry.Trace.Monad qualified as Otel
 import PossehlAnalyticsPrelude
@@ -816,12 +818,16 @@ traceQueryIfEnabled ::
   Transaction m ()
 traceQueryIfEnabled tools span logDatabaseQueries qry params = do
   -- In case we have query logging enabled, we want to do that
-  let formattedQuery =
-        Otel.inSpan "Postgres Query Formatting" Otel.defaultSpanArguments $
-          case params of
+  let formattedQuery = do
+        withEvent
+          span
+          "Query Format start"
+          "Query Format end"
+          $ case params of
             HasNoParams -> pgFormatQueryNoParams' tools qry
             HasSingleParam p -> pgFormatQuery' tools qry p
             HasMultiParams ps -> pgFormatQueryMany' tools qry ps
+
   let doLog errs =
         Otel.addAttributes
           span
@@ -859,6 +865,37 @@ traceQueryIfEnabled tools span logDatabaseQueries qry params = do
       ex <- doExplain
       doLog (T2 (label @"query" q) (label @"explain" (Just ex)))
 
+-- | Add a start and end event to the span, and figure out how long the difference was.
+--
+-- This is more lightweight than starting an extra span for timing things.
+withEvent :: (MonadIO f) => Otel.Span -> Text -> Text -> f b -> f b
+withEvent span start end act = do
+  let mkMs ts = (ts & Otel.timestampNanoseconds & toInteger) `div` 1000_000
+  s <- Otel.getTimestamp
+  Otel.addEvent
+    span
+    ( Otel.NewEvent
+        { newEventName = start,
+          newEventAttributes = mempty,
+          newEventTimestamp = Just s
+        }
+    )
+  res <- act
+  e <- Otel.getTimestamp
+  let tookMs =
+        (mkMs e - mkMs s)
+          -- should be small enough
+          & fromInteger @Int
+  Otel.addEvent
+    span
+    ( Otel.NewEvent
+        { newEventName = end,
+          newEventAttributes = HashMap.fromList [("took ms", Otel.toAttribute tookMs)],
+          newEventTimestamp = Just e
+        }
+    )
+  pure res
+
 instance (ToField t1) => ToRow (Label l1 t1) where
   toRow t2 = toRow $ PG.Only $ getField @l1 t2
 
-- 
cgit 1.4.1