better observability for redis timeouts

nri-redis/src/Redis/Handler.hs

@@ -30,12 +30,6 @@ handler :: Text -> Settings.Settings -> Data.Acquire.Acquire Internal.Handler
 handler namespace settings = do
   (namespacedHandler, _) <- Data.Acquire.mkAcquire (acquireHandler namespace settings) releaseHandler
   namespacedHandler
-    |> ( \handler' ->
-           case Settings.queryTimeout settings of
-             Settings.NoQueryTimeout -> handler'
-             Settings.TimeoutQueryAfterMilliseconds milliseconds ->
-               timeoutAfterMilliseconds (toFloat milliseconds) handler'
-       )
     |> Prelude.pure
 
 -- | Produce a namespaced handler for Redis access.
@@ -44,12 +38,6 @@ handlerAutoExtendExpire :: Text -> Settings.Settings -> Data.Acquire.Acquire Int
 handlerAutoExtendExpire namespace settings = do
   (namespacedHandler, _) <- Data.Acquire.mkAcquire (acquireHandler namespace settings) releaseHandler
   namespacedHandler
-    |> ( \handler' ->
-           case Settings.queryTimeout settings of
-             Settings.NoQueryTimeout -> handler'
-             Settings.TimeoutQueryAfterMilliseconds milliseconds ->
-               timeoutAfterMilliseconds (toFloat milliseconds) handler'
-       )
     |> ( \handler' -> case Settings.defaultExpiry settings of
            Settings.NoDefaultExpiry ->
              -- We create the handler as part of starting the application. Throwing
@@ -69,20 +57,6 @@ handlerAutoExtendExpire namespace settings = do
        )
     |> liftIO
 
-timeoutAfterMilliseconds :: Float -> Internal.Handler' x -> Internal.Handler' x
-timeoutAfterMilliseconds milliseconds handler' =
-  handler'
-    { Internal.doQuery =
-        Stack.withFrozenCallStack (Internal.doQuery handler')
-          >> Task.timeout milliseconds Internal.TimeoutError,
-      Internal.doTransaction =
-        Stack.withFrozenCallStack (Internal.doTransaction handler')
-          >> Task.timeout milliseconds Internal.TimeoutError,
-      Internal.doEval =
-        Stack.withFrozenCallStack (Internal.doEval handler')
-          >> Task.timeout milliseconds Internal.TimeoutError
-    }
-
 defaultExpiryKeysAfterSeconds :: Int -> Internal.HandlerAutoExtendExpire -> Internal.HandlerAutoExtendExpire
 defaultExpiryKeysAfterSeconds secs handler' =
   let wrapWithExpire :: Internal.Query a -> Internal.Query a
@@ -121,11 +95,12 @@ acquireHandler namespace settings = do
             Database.Redis.UnixSocket _ -> Nothing
     pure Connection {connectionHedis, connectionHost, connectionPort}
   anything <- Platform.doAnythingHandler
+  let queryTimeout = (Settings.queryTimeout settings)
   pure
     ( Internal.Handler'
         { Internal.doQuery = \query ->
             let PreparedQuery {redisCtx} = doRawQuery query
-             in Stack.withFrozenCallStack platformRedis (Internal.cmds query) connection anything redisCtx,
+             in Stack.withFrozenCallStack platformRedis (Internal.cmds query) connection anything queryTimeout redisCtx,
           Internal.doTransaction = \query ->
             let PreparedQuery {redisCtx} = doRawQuery query
                 redisCmd = Database.Redis.multiExec redisCtx
@@ -137,9 +112,9 @@ acquireHandler namespace settings = do
                           Database.Redis.TxAborted -> Right (Err Internal.TransactionAborted)
                           Database.Redis.TxError err -> Right (Err (Internal.RedisError (Text.fromList err)))
                     )
-                  |> Stack.withFrozenCallStack (platformRedis (Internal.cmds query) connection anything),
+                  |> Stack.withFrozenCallStack (platformRedis (Internal.cmds query) connection anything queryTimeout),
           Internal.doEval = \script' ->
-            Stack.withFrozenCallStack (platformRedisScript script' connection anything),
+            Stack.withFrozenCallStack (platformRedisScript script' connection anything queryTimeout),
           Internal.namespace = namespace,
           Internal.maxKeySize = Settings.maxKeySize settings
         },
@@ -306,9 +281,9 @@ doRawQuery query =
         |> PreparedQuery
         |> map (Ok << Prelude.fromIntegral)
     Internal.Sismember key val ->
-            Database.Redis.sismember (toB key) val
-            |> PreparedQuery
-            |> map Ok
+      Database.Redis.sismember (toB key) val
+        |> PreparedQuery
+        |> map Ok
     Internal.Smembers key ->
       Database.Redis.smembers (toB key)
         |> PreparedQuery
@@ -362,13 +337,14 @@ data Connection = Connection
   }
 
 platformRedis ::
-  Stack.HasCallStack =>
+  (Stack.HasCallStack) =>
   [Text] ->
   Connection ->
   Platform.DoAnythingHandler ->
+  Settings.QueryTimeout ->
   Database.Redis.Redis (Either Database.Redis.Reply (Result Internal.Error a)) ->
   Task Internal.Error a
-platformRedis cmds connection anything action =
+platformRedis cmds connection anything queryTimeout action =
   Database.Redis.runRedis (connectionHedis connection) action
     |> map toResult
     |> map
@@ -379,7 +355,7 @@ platformRedis cmds connection anything action =
       )
     |> handleExceptions
     |> Platform.doAnything anything
-    |> Stack.withFrozenCallStack Internal.traceQuery cmds (connectionHost connection) (connectionPort connection)
+    |> Stack.withFrozenCallStack Internal.wrapQuery queryTimeout cmds (connectionHost connection) (connectionPort connection)
 
 toResult :: Either Database.Redis.Reply a -> Result Internal.Error a
 toResult reply =
@@ -406,17 +382,18 @@ platformRedisScript ::
   Script.Script a ->
   Connection ->
   Platform.DoAnythingHandler ->
+  Settings.QueryTimeout ->
   Task Internal.Error a
-platformRedisScript script connection anything = do
+platformRedisScript script connection anything queryTimeout = do
   -- Try EVALSHA
-  evalsha script connection anything
+  evalsha script connection anything queryTimeout
     |> Task.onError
       ( \err ->
           case err of
             Internal.RedisError "NOSCRIPT No matching script. Please use EVAL." -> do
               -- If it fails with NOSCRIPT, load the script and try again
-              loadScript script connection anything
-              evalsha script connection anything
+              loadScript script connection anything queryTimeout
+              evalsha script connection anything queryTimeout
             _ -> Task.fail err
       )
 
@@ -425,8 +402,9 @@ evalsha ::
   Script.Script a ->
   Connection ->
   Platform.DoAnythingHandler ->
+  Settings.QueryTimeout ->
   Task Internal.Error a
-evalsha script connection anything =
+evalsha script connection anything queryTimeout =
   Database.Redis.evalsha
     (toB (Script.luaScriptHash script))
     (map toB (Script.keys script))
@@ -435,23 +413,24 @@ evalsha script connection anything =
     |> map toResult
     |> handleExceptions
     |> Platform.doAnything anything
-    |> Stack.withFrozenCallStack Internal.traceQuery [Script.evalShaString script] (connectionHost connection) (connectionPort connection)
+    |> Stack.withFrozenCallStack Internal.wrapQuery queryTimeout [Script.evalShaString script] (connectionHost connection) (connectionPort connection)
 
 loadScript ::
   Stack.HasCallStack =>
   Script.Script a ->
   Connection ->
   Platform.DoAnythingHandler ->
+  Settings.QueryTimeout ->
   Task Internal.Error ()
-loadScript script connection anything = do
+loadScript script connection anything queryTimeout = do
   Database.Redis.scriptLoad (toB (Script.luaScript script))
     |> Database.Redis.runRedis (connectionHedis connection)
     |> map toResult
     |> handleExceptions
     -- The result is the hash, which we already have. No sense in decoding it.
     |> map (map (\_ -> ()))
     |> Platform.doAnything anything
-    |> Stack.withFrozenCallStack Internal.traceQuery [Script.scriptLoadString script] (connectionHost connection) (connectionPort connection)
+    |> Stack.withFrozenCallStack Internal.wrapQuery queryTimeout [Script.scriptLoadString script] (connectionHost connection) (connectionPort connection)
 
 toB :: Text -> Data.ByteString.ByteString
 toB = Data.Text.Encoding.encodeUtf8

nri-redis/src/Redis/Internal.hs

@@ -23,7 +23,7 @@ module Redis.Internal
     eval,
     foldWithScan,
     -- internal tools
-    traceQuery,
+    wrapQuery,
     maybesToDict,
     keysTouchedByQuery,
   )
@@ -111,7 +111,7 @@ cmds query'' =
     Sadd key vals -> [unwords ("SADD" : key : List.map (\_ -> "*****") (NonEmpty.toList vals))]
     Scard key -> [unwords ["SCARD", key]]
     Srem key vals -> [unwords ("SREM" : key : List.map (\_ -> "*****") (NonEmpty.toList vals))]
-    Sismember key _ -> [unwords ["SISMEMBER", key , "*****"]]
+    Sismember key _ -> [unwords ["SISMEMBER", key, "*****"]]
     Smembers key -> [unwords ["SMEMBERS", key]]
     Zadd key vals -> [unwords ("ZADD" : key : List.concatMap (\(_, val) -> ["*****", Text.fromFloat val]) (Dict.toList vals))]
     Zrange key start stop -> [unwords ["ZRANGE", key, Text.fromInt start, Text.fromInt stop]]
@@ -230,9 +230,9 @@ data HasAutoExtendExpire = NoAutoExtendExpire | AutoExtendExpire
 -- A handler that can only be parametrized by a value of this kind.
 -- Meaning that we use the values of the type parameter at a type level.
 data Handler' (x :: HasAutoExtendExpire) = Handler'
-  { doQuery :: Stack.HasCallStack => forall a. Query a -> Task Error a,
-    doTransaction :: Stack.HasCallStack => forall a. Query a -> Task Error a,
-    doEval :: Stack.HasCallStack => forall a. Database.Redis.RedisResult a => Script.Script a -> Task Error a,
+  { doQuery :: (Stack.HasCallStack) => forall a. Query a -> Task Error a,
+    doTransaction :: (Stack.HasCallStack) => forall a. Query a -> Task Error a,
+    doEval :: (Stack.HasCallStack) => forall a. (Database.Redis.RedisResult a) => Script.Script a -> Task Error a,
     namespace :: Text,
     maxKeySize :: Settings.MaxKeySize
   }
@@ -253,7 +253,7 @@ type HandlerAutoExtendExpire = Handler' 'AutoExtendExpire
 -- Note: A 'Query' in this library can consist of one or more queries in sequence.
 -- if a 'Query' contains multiple queries, it may make more sense, if possible
 -- to run them using 'transaction'
-query :: Stack.HasCallStack => Handler' x -> Query a -> Task Error a
+query :: (Stack.HasCallStack) => Handler' x -> Query a -> Task Error a
 query handler query' =
   namespaceQuery (namespace handler ++ ":") query'
     |> Task.andThen (ensureMaxKeySize handler)
@@ -265,7 +265,7 @@ query handler query' =
 --
 -- In redis terms, this is wrappping the 'Query' in `MULTI` and `EXEC
 -- see redis transaction semantics here: https://redis.io/topics/transactions
-transaction :: Stack.HasCallStack => Handler' x -> Query a -> Task Error a
+transaction :: (Stack.HasCallStack) => Handler' x -> Query a -> Task Error a
 transaction handler query' =
   namespaceQuery (namespace handler ++ ":") query'
     |> Task.andThen (ensureMaxKeySize handler)
@@ -424,7 +424,7 @@ keysTouchedByQuery query' =
     Zrevrank key _ -> Set.singleton key
     WithResult _ q -> keysTouchedByQuery q
 
-maybesToDict :: Ord key => List key -> List (Maybe a) -> Dict.Dict key a
+maybesToDict :: (Ord key) => List key -> List (Maybe a) -> Dict.Dict key a
 maybesToDict keys values =
   List.map2 (,) keys values
     |> List.filterMap
@@ -435,7 +435,15 @@ maybesToDict keys values =
       )
     |> Dict.fromList
 
-traceQuery :: Stack.HasCallStack => [Text] -> Text -> Maybe Int -> Task e a -> Task e a
+wrapQuery :: (Stack.HasCallStack) => Settings.QueryTimeout -> [Text] -> Text -> Maybe Int -> Task Error a -> Task Error a
+wrapQuery queryTimeout commands host port task =
+  traceQuery commands host port <| case queryTimeout of
+    Settings.NoQueryTimeout ->
+      task
+    Settings.TimeoutQueryAfterMilliseconds timeoutMs ->
+      Task.timeout (toFloat timeoutMs) TimeoutError task
+
+traceQuery :: (Stack.HasCallStack) => [Text] -> Text -> Maybe Int -> Task Error a -> Task Error a
 traceQuery commands host port task =
   let info =
         RedisCommands.emptyDetails

nri-redis/test/Helpers.hs

@@ -11,15 +11,17 @@ import qualified Prelude
 
 data TestHandlers = TestHandlers
   { autoExtendExpireHandler :: Redis.HandlerAutoExtendExpire,
-    handler :: Redis.Handler
+    handler :: Redis.Handler,
+    handlerWithMinimalExpire :: Redis.Handler
   }
 
 getHandlers :: Conduit.Acquire TestHandlers
 getHandlers = do
   settings <- Conduit.liftIO (Environment.decode Settings.decoder)
   autoExtendExpireHandler <- Handler.handlerAutoExtendExpire "tests-auto-extend-expire" settings {Settings.defaultExpiry = Settings.ExpireKeysAfterSeconds 1}
   handler <- Handler.handler "tests" settings {Settings.defaultExpiry = Settings.NoDefaultExpiry}
-  Prelude.pure TestHandlers {autoExtendExpireHandler, handler}
+  handlerWithMinimalExpire <- Handler.handler "tests" settings {Settings.queryTimeout = Settings.TimeoutQueryAfterMilliseconds 0}
+  Prelude.pure TestHandlers {autoExtendExpireHandler, handler, handlerWithMinimalExpire}
 
 -- | Historical context:
 -- Golden results are slightly different between GHC 9.2.x and 8.10.x due

nri-redis/test/Spec/Redis.hs

@@ -23,7 +23,7 @@ import qualified Prelude
 -- put this at the top of the file so that adding tests doesn't push
 -- the line number of the source location of this file down, which would
 -- change golden test results
-spanForTask :: Show e => Task e () -> Expect.Expectation' Platform.TracingSpan
+spanForTask :: (Show e) => Task e () -> Expect.Expectation' Platform.TracingSpan
 spanForTask task =
   Expect.fromIO <| do
     spanVar <- MVar.newEmptyMVar
@@ -39,13 +39,30 @@ spanForTask task =
         MVar.takeMVar spanVar
           |> map constantValuesForVariableFields
 
+spanForFailingTask :: Task e () -> Expect.Expectation' Platform.TracingSpan
+spanForFailingTask task =
+  Expect.fromIO <| do
+    spanVar <- MVar.newEmptyMVar
+    res <-
+      Platform.rootTracingSpanIO
+        "test-request"
+        (MVar.putMVar spanVar)
+        "test-root"
+        (\log -> Task.attempt log task)
+    case res of
+      Err _ ->
+        MVar.takeMVar spanVar
+          |> map constantValuesForVariableFields
+      Ok _ ->
+        Prelude.fail "Expected task to fail"
+
 tests :: TestHandlers -> Test.Test
-tests TestHandlers {handler, autoExtendExpireHandler} =
+tests TestHandlers {handler, autoExtendExpireHandler, handlerWithMinimalExpire} =
   Test.describe
     "Redis Library"
     [ Test.describe "query tests using handler" (queryTests handler),
       Test.describe "query tests using auto extend expire handler" (queryTests autoExtendExpireHandler),
-      Test.describe "observability tests" (observabilityTests handler)
+      Test.describe "observability tests" (observabilityTests handler handlerWithMinimalExpire)
     ]
 
 -- We want to test all of our potential makeApi alternatives because it's easy
@@ -56,8 +73,8 @@ tests TestHandlers {handler, autoExtendExpireHandler} =
 -- value "test/Main.hs". If it points to one of the src files of the redis
 -- library it means stack frames for redis query in bugsnag, newrelic, etc will
 -- not point to the application code making the query!
-observabilityTests :: Redis.Handler' x -> List Test.Test
-observabilityTests handler =
+observabilityTests :: Redis.Handler' x -> Redis.Handler' x -> List Test.Test
+observabilityTests handler handlerWithMinimalExpire =
   [ Test.test "Redis.query reports the span data we expect" <| \() -> do
       span <-
         Redis.query handler (Redis.ping api)
@@ -113,7 +130,20 @@ observabilityTests handler =
           |> spanForTask
       span
         |> Debug.toString
-        |> Expect.equalToContentsOf (goldenResultsDir ++ "/observability-spec-reporting-redis-counter-transaction")
+        |> Expect.equalToContentsOf (goldenResultsDir ++ "/observability-spec-reporting-redis-counter-transaction"),
+    Test.describe
+      "with 0 ms timeout"
+      [ Test.test "Redis.query reports the span data we expect" <| \() -> do
+          span <-
+            Redis.query handlerWithMinimalExpire (Redis.ping api)
+              |> spanForFailingTask
+          span
+            |> Debug.toString
+            |> Expect.all
+              [ Expect.equalToContentsOf (goldenResultsDir ++ "/observability-spec-timeout-reporting-redis-query"),
+                \spanText -> Expect.true (Text.contains "Redis Query" spanText)
+              ]
+      ]
   ]
 
 queryTests :: Redis.Handler' x -> List Test.Test

nri-redis/test/golden-results-9.2/observability-spec-reporting-redis-counter-query

@@ -32,9 +32,9 @@ TracingSpan
                     { srcLocPackage = "main"
                     , srcLocModule = "Spec.Redis"
                     , srcLocFile = "test/Spec/Redis.hs"
-                    , srcLocStartLine = 105
+                    , srcLocStartLine = 122
                     , srcLocStartCol = 9
-                    , srcLocEndLine = 105
+                    , srcLocEndLine = 122
                     , srcLocEndCol = 28
                     }
                 )

nri-redis/test/golden-results-9.2/observability-spec-reporting-redis-counter-transaction

@@ -32,9 +32,9 @@ TracingSpan
                     { srcLocPackage = "main"
                     , srcLocModule = "Spec.Redis"
                     , srcLocFile = "test/Spec/Redis.hs"
-                    , srcLocStartLine = 112
+                    , srcLocStartLine = 129
                     , srcLocStartCol = 9
-                    , srcLocEndLine = 112
+                    , srcLocEndLine = 129
                     , srcLocEndCol = 34
                     }
                 )

nri-redis/test/golden-results-9.2/observability-spec-reporting-redis-hash-query

@@ -32,9 +32,9 @@ TracingSpan
                     { srcLocPackage = "main"
                     , srcLocModule = "Spec.Redis"
                     , srcLocFile = "test/Spec/Redis.hs"
-                    , srcLocStartLine = 77
+                    , srcLocStartLine = 94
                     , srcLocStartCol = 9
-                    , srcLocEndLine = 77
+                    , srcLocEndLine = 94
                     , srcLocEndCol = 25
                     }
                 )

nri-redis/test/golden-results-9.2/observability-spec-reporting-redis-hash-transaction

@@ -32,9 +32,9 @@ TracingSpan
                     { srcLocPackage = "main"
                     , srcLocModule = "Spec.Redis"
                     , srcLocFile = "test/Spec/Redis.hs"
-                    , srcLocStartLine = 84
+                    , srcLocStartLine = 101
                     , srcLocStartCol = 9
-                    , srcLocEndLine = 84
+                    , srcLocEndLine = 101
                     , srcLocEndCol = 31
                     }
                 )