Merge pull request #124 from NoRedInk/can-modify-redis-timeout-mid-st…

…ream

can modify redis timeout mid stream

nri-redis/CHANGELOG.md

@@ -1,4 +1,10 @@
+# 0.2.0.3 unreleased
+
+- [bugfix] When a query times out, its context is no longer removed from the Stack
+- query timeout setting can be modified at runtime
+
 # 0.2.0.2
+
 - Adds `sismember`
 
 # 0.2.0.1

nri-redis/src/Redis.hs

@@ -18,6 +18,8 @@ module Redis
     Settings.decoder,
     Settings.decoderWithEnvVarPrefix,
     Settings.decoderWithCustomConnectionString,
+    Handler.withQueryTimeoutMilliseconds,
+    Handler.withoutQueryTimeout,
 
     -- * Creating a redis API
     jsonApi,

nri-redis/src/Redis/Handler.hs

@@ -4,6 +4,8 @@
 module Redis.Handler
   ( handler,
     handlerAutoExtendExpire,
+    withQueryTimeoutMilliseconds,
+    withoutQueryTimeout,
   )
 where
 
@@ -57,6 +59,16 @@ handlerAutoExtendExpire namespace settings = do
        )
     |> liftIO
 
+-- | Sets a timeout for the query in milliseconds.
+withQueryTimeoutMilliseconds :: Int -> Internal.Handler' x -> Internal.Handler' x
+withQueryTimeoutMilliseconds timeoutMs handler' =
+  handler' {Internal.queryTimeout = Settings.TimeoutQueryAfterMilliseconds timeoutMs}
+
+-- | Disables timeout for query in milliseconds
+withoutQueryTimeout :: Internal.Handler' x -> Internal.Handler' x
+withoutQueryTimeout handler' =
+  handler' {Internal.queryTimeout = Settings.NoQueryTimeout}
+
 defaultExpiryKeysAfterSeconds :: Int -> Internal.HandlerAutoExtendExpire -> Internal.HandlerAutoExtendExpire
 defaultExpiryKeysAfterSeconds secs handler' =
   let wrapWithExpire :: Internal.Query a -> Internal.Query a
@@ -67,15 +79,15 @@ defaultExpiryKeysAfterSeconds secs handler' =
           |> Internal.sequence
           |> Internal.map2 (\res _ -> res) query'
    in handler'
-        { Internal.doQuery = \query' ->
+        { Internal.doQuery = \queryTimeout query' ->
             wrapWithExpire query'
-              |> Stack.withFrozenCallStack (Internal.doQuery handler'),
-          Internal.doTransaction = \query' ->
+              |> Stack.withFrozenCallStack (Internal.doQuery handler') queryTimeout,
+          Internal.doTransaction = \queryTimeout query' ->
             wrapWithExpire query'
-              |> Stack.withFrozenCallStack (Internal.doTransaction handler'),
-          Internal.doEval = \script' ->
+              |> Stack.withFrozenCallStack (Internal.doTransaction handler') queryTimeout,
+          Internal.doEval = \queryTimeout script' ->
             -- We can't guarantee auto-expire for EVAL, so we just run it as-is
-            Stack.withFrozenCallStack (Internal.doEval handler' script')
+            Stack.withFrozenCallStack (Internal.doEval handler' queryTimeout script')
         }
 
 acquireHandler :: Text -> Settings.Settings -> IO (Internal.Handler' x, Connection)
@@ -95,13 +107,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 ->
+        { Internal.doQuery = \queryTimeout query ->
             let PreparedQuery {redisCtx} = doRawQuery query
              in Stack.withFrozenCallStack platformRedis (Internal.cmds query) connection anything queryTimeout redisCtx,
-          Internal.doTransaction = \query ->
+          Internal.doTransaction = \queryTimeout query ->
             let PreparedQuery {redisCtx} = doRawQuery query
                 redisCmd = Database.Redis.multiExec redisCtx
              in redisCmd
@@ -113,10 +124,11 @@ acquireHandler namespace settings = do
                           Database.Redis.TxError err -> Right (Err (Internal.RedisError (Text.fromList err)))
                     )
                   |> Stack.withFrozenCallStack (platformRedis (Internal.cmds query) connection anything queryTimeout),
-          Internal.doEval = \script' ->
+          Internal.doEval = \queryTimeout script' ->
             Stack.withFrozenCallStack (platformRedisScript script' connection anything queryTimeout),
           Internal.namespace = namespace,
-          Internal.maxKeySize = Settings.maxKeySize settings
+          Internal.maxKeySize = Settings.maxKeySize settings,
+          Internal.queryTimeout = Settings.queryTimeout settings
         },
       connection
     )
@@ -416,7 +428,7 @@ evalsha script connection anything queryTimeout =
     |> Stack.withFrozenCallStack Internal.wrapQuery queryTimeout [Script.evalShaString script] (connectionHost connection) (connectionPort connection)
 
 loadScript ::
-  Stack.HasCallStack =>
+  (Stack.HasCallStack) =>
   Script.Script a ->
   Connection ->
   Platform.DoAnythingHandler ->

nri-redis/src/Redis/Internal.hs

@@ -230,11 +230,12 @@ 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. Settings.QueryTimeout -> Query a -> Task Error a,
+    doTransaction :: (Stack.HasCallStack) => forall a. Settings.QueryTimeout -> Query a -> Task Error a,
+    doEval :: (Stack.HasCallStack) => forall a. (Database.Redis.RedisResult a) => Settings.QueryTimeout -> Script.Script a -> Task Error a,
     namespace :: Text,
-    maxKeySize :: Settings.MaxKeySize
+    maxKeySize :: Settings.MaxKeySize,
+    queryTimeout :: Settings.QueryTimeout
   }
 
 -- | This is a type alias of a handler parametrized by a value that indicates
@@ -257,7 +258,7 @@ query :: (Stack.HasCallStack) => Handler' x -> Query a -> Task Error a
 query handler query' =
   namespaceQuery (namespace handler ++ ":") query'
     |> Task.andThen (ensureMaxKeySize handler)
-    |> Task.andThen (Stack.withFrozenCallStack (doQuery handler))
+    |> Task.andThen (Stack.withFrozenCallStack (doQuery handler) (queryTimeout handler))
 
 -- | Run a redis Query in a transaction. If the query contains several Redis
 -- commands they're all executed together, and Redis will guarantee other
@@ -269,12 +270,12 @@ transaction :: (Stack.HasCallStack) => Handler' x -> Query a -> Task Error a
 transaction handler query' =
   namespaceQuery (namespace handler ++ ":") query'
     |> Task.andThen (ensureMaxKeySize handler)
-    |> Task.andThen (Stack.withFrozenCallStack (doTransaction handler))
+    |> Task.andThen (Stack.withFrozenCallStack (doTransaction handler) (queryTimeout handler))
 
 eval :: (Stack.HasCallStack, Database.Redis.RedisResult a) => Handler' x -> Script.Script a -> Task Error a
 eval handler script =
   Script.mapKeys (\key -> Task.succeed (namespace handler ++ ":" ++ key)) script
-    |> Task.andThen (Stack.withFrozenCallStack (doEval handler))
+    |> Task.andThen (Stack.withFrozenCallStack (doEval handler) (queryTimeout handler))
 
 namespaceQuery :: Text -> Query a -> Task err (Query a)
 namespaceQuery prefix query' =

nri-redis/test/Helpers.hs

@@ -11,17 +11,15 @@ import qualified Prelude
 
 data TestHandlers = TestHandlers
   { autoExtendExpireHandler :: Redis.HandlerAutoExtendExpire,
-    handler :: Redis.Handler,
-    handlerWithMinimalExpire :: Redis.Handler
+    handler :: 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}
-  handlerWithMinimalExpire <- Handler.handler "tests" settings {Settings.queryTimeout = Settings.TimeoutQueryAfterMilliseconds 0}
-  Prelude.pure TestHandlers {autoExtendExpireHandler, handler, handlerWithMinimalExpire}
+  Prelude.pure TestHandlers {autoExtendExpireHandler, handler}
 
 -- | Historical context:
 -- Golden results are slightly different between GHC 9.2.x and 8.10.x due

nri-redis/test/Spec/Redis.hs

@@ -57,12 +57,12 @@ spanForFailingTask task =
         Prelude.fail "Expected task to fail"
 
 tests :: TestHandlers -> Test.Test
-tests TestHandlers {handler, autoExtendExpireHandler, handlerWithMinimalExpire} =
+tests TestHandlers {handler, autoExtendExpireHandler} =
   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 handlerWithMinimalExpire)
+      Test.describe "observability tests" (observabilityTests handler)
     ]
 
 -- We want to test all of our potential makeApi alternatives because it's easy
@@ -73,8 +73,8 @@ tests TestHandlers {handler, autoExtendExpireHandler, handlerWithMinimalExpire}
 -- 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 -> Redis.Handler' x -> List Test.Test
-observabilityTests handler handlerWithMinimalExpire =
+observabilityTests :: Redis.Handler' x -> List Test.Test
+observabilityTests handler =
   [ Test.test "Redis.query reports the span data we expect" <| \() -> do
       span <-
         Redis.query handler (Redis.ping api)
@@ -135,7 +135,7 @@ observabilityTests handler handlerWithMinimalExpire =
       "with 0 ms timeout"
       [ Test.test "Redis.query reports the span data we expect" <| \() -> do
           span <-
-            Redis.query handlerWithMinimalExpire (Redis.ping api)
+            Redis.query (Redis.withQueryTimeoutMilliseconds 0 handler) (Redis.ping api)
               |> spanForFailingTask
           span
             |> Debug.toString