Support disabling timeout mid-way through

lib/graphql/schema/timeout.rb

@@ -71,15 +71,23 @@ def execute_multiplex(multiplex:)
         def execute_field(query:, field:, **_rest)
           timeout_state = query.context.namespace(@timeout).fetch(:state)
           # If the `:state` is `false`, then `max_seconds(query)` opted out of timeout for this query.
-          if timeout_state != false && Process.clock_gettime(Process::CLOCK_MONOTONIC, :millisecond) > timeout_state.fetch(:timeout_at)
+          if timeout_state == false
+            super
+          elsif Process.clock_gettime(Process::CLOCK_MONOTONIC, :millisecond) > timeout_state.fetch(:timeout_at)
             error = GraphQL::Schema::Timeout::TimeoutError.new(field)
             # Only invoke the timeout callback for the first timeout
             if !timeout_state[:timed_out]
               timeout_state[:timed_out] = true
               @timeout.handle_timeout(error, query)
+              timeout_state = query.context.namespace(@timeout).fetch(:state)
             end
 
-            error
+            # `handle_timeout` may have set this to be `false`
+            if timeout_state != false
+              error
+            else
+              super
+            end
           else
             super
           end
@@ -102,6 +110,15 @@ def handle_timeout(error, query)
         # override to do something interesting
       end
 
+      # Call this method (eg, from {#handle_timeout}) to disable timeout tracking
+      # for the given query.
+      # @param query [GraphQL::Query]
+      # @return [void]
+      def disable_timeout(query)
+        query.context.namespace(self)[:state] = false
+        nil
+      end
+
       # This error is raised when a query exceeds `max_seconds`.
       # Since it's a child of {GraphQL::ExecutionError},
       # its message will be added to the response's `errors` key.

spec/graphql/schema/timeout_spec.rb

@@ -9,8 +9,17 @@ def execute_field(query:, **opts)
     end
   end
 
+  class CustomTimeout < GraphQL::Schema::Timeout
+    def handle_timeout(error, query)
+      if query.context[:disable_timeout]
+        disable_timeout(query)
+      end
+      super
+    end
+  end
+
   let(:max_seconds) { 1 }
-  let(:timeout_class) { GraphQL::Schema::Timeout }
+  let(:timeout_class) { CustomTimeout }
   let(:timeout_schema) {
     nested_sleep_type = Class.new(GraphQL::Schema::Object) do
       graphql_name "NestedSleep"
@@ -36,9 +45,13 @@ def nested_sleep(seconds:)
 
       field :sleep_for, Float do
         argument :seconds, Float
+        argument :disable_timeout, "Boolean", default_value: false
       end
 
-      def sleep_for(seconds:)
+      def sleep_for(seconds:, disable_timeout:)
+        if disable_timeout
+          context[:disable_timeout] = true
+        end
         sleep(seconds)
         seconds
       end
@@ -254,4 +267,39 @@ def handle_timeout(err, query)
       end
     end
   end
+
+
+  describe "disabling the timeout" do
+    let(:query_string) {%|
+    query GetTimeouts($disable: Boolean) {
+      a: sleepFor(seconds: 0.4)
+      b: sleepFor(seconds: 0.4)
+      c: sleepFor(seconds: 0.4, disableTimeout: $disable)
+      d: sleepFor(seconds: 0.4)
+      e: sleepFor(seconds: 0.4)
+    }
+  |}
+
+    it "can be disabled" do
+      expected_data = {
+        "a"=>0.4,
+        "b"=>0.4,
+        "c"=>0.4,
+        "d"=>nil,
+        "e"=>nil,
+      }
+
+      assert_graphql_equal expected_data, result["data"]
+
+      disabled_timeout_result = timeout_schema.execute(query_string, context: query_context, variables: { disable: true })
+      expected_disabled_data = {
+        "a"=>0.4,
+        "b"=>0.4,
+        "c"=>0.4,
+        "d"=>0.4,
+        "e"=>0.4,
+      }
+      assert_graphql_equal expected_disabled_data, disabled_timeout_result["data"]
+    end
+  end
 end