From 81bb4dc4b3697fcee3772ee052fff72d80a5ba72 Mon Sep 17 00:00:00 2001
From: Robert Mosolgo <rdmosolgo@gmail.com>
Date: Mon, 15 Jan 2024 15:18:19 -0500
Subject: [PATCH] Use trace_with to start and end batches

---
 lib/graphql/batch.rb                 |  8 ++++++--
 lib/graphql/batch/setup_multiplex.rb | 14 ++++++++++++++
 2 files changed, 20 insertions(+), 2 deletions(-)

diff --git a/lib/graphql/batch.rb b/lib/graphql/batch.rb
index 0979cc1..4e6d096 100644
--- a/lib/graphql/batch.rb
+++ b/lib/graphql/batch.rb
@@ -16,8 +16,12 @@ def self.batch(executor_class: GraphQL::Batch::Executor)
     end
 
     def self.use(schema_defn, executor_class: GraphQL::Batch::Executor)
-      instrumentation = GraphQL::Batch::SetupMultiplex.new(schema_defn, executor_class: executor_class)
-      schema_defn.instrument(:multiplex, instrumentation)
+      if schema_defn.respond_to?(:trace_with)
+        schema_defn.trace_with(GraphQL::Batch::SetupMultiplex::Trace, executor_class: executor_class)
+      else
+        instrumentation = GraphQL::Batch::SetupMultiplex.new(schema_defn, executor_class: executor_class)
+        schema_defn.instrument(:multiplex, instrumentation)
+      end
 
       if schema_defn.mutation
         require_relative "batch/mutation_field_extension"
diff --git a/lib/graphql/batch/setup_multiplex.rb b/lib/graphql/batch/setup_multiplex.rb
index 940d2a7..ad1d392 100644
--- a/lib/graphql/batch/setup_multiplex.rb
+++ b/lib/graphql/batch/setup_multiplex.rb
@@ -12,5 +12,19 @@ def before_multiplex(multiplex)
     def after_multiplex(multiplex)
       GraphQL::Batch::Executor.end_batch
     end
+
+    module Trace
+      def initialize(executor_class:, **_rest)
+        @executor_class = executor_class
+        super
+      end
+
+      def execute_multiplex(multiplex:)
+        GraphQL::Batch::Executor.start_batch(@executor_class)
+        super
+      ensure
+        GraphQL::Batch::Executor.end_batch
+      end
+    end
   end
 end