Use Task.yield/2 instead of Task.await/2 in batching fn, so we can log error

lib/absinthe/middleware/batch.ex

@@ -60,6 +60,8 @@ defmodule Absinthe.Middleware.Batch do
   @behaviour Absinthe.Middleware
   @behaviour Absinthe.Plugin
 
+  require Logger
+
   @typedoc """
   The function to be called with the aggregate batch information.
 
@@ -143,18 +145,29 @@ defmodule Absinthe.Middleware.Batch do
 
       metadata = emit_start_event(system_time, batch_fun, batch_opts, batch_data)
 
-      {batch_opts, task, start_time_mono, metadata}
+      {batch_opts, task, start_time_mono, metadata, batch_fun}
     end)
-    |> Map.new(fn {batch_opts, task, start_time_mono, metadata} ->
-      timeout = Keyword.get(batch_opts, :timeout, 5_000)
-      result = Task.await(task, timeout)
+    |> Map.new(&yield_batching_result/1)
+  end
 
-      end_time_mono = System.monotonic_time()
-      duration = end_time_mono - start_time_mono
-      emit_stop_event(duration, end_time_mono, metadata, result)
+  defp yield_batching_result({batch_opts, task, start_time_mono, metadata, batch_fun}) do
+    timeout = Keyword.get(batch_opts, :timeout, 5_000)
 
-      result
-    end)
+    case Task.yield(task, timeout) || Task.shutdown(task) do
+      {:ok, result} ->
+        end_time_mono = System.monotonic_time()
+        duration = end_time_mono - start_time_mono
+        emit_stop_event(duration, end_time_mono, metadata, result)
+
+        result
+
+      _ ->
+        Logger.error(
+          "Failed to get batching result in #{timeout}ms for\nfn: #{inspect(batch_fun)}"
+        )
+
+        Process.exit(self(), :timeout)
+    end
   end
 
   @batch_start [:absinthe, :middleware, :batch, :start]

test/absinthe/middleware/batch_test.exs

@@ -1,6 +1,14 @@
 defmodule Absinthe.Middleware.BatchTest do
   use Absinthe.Case, async: true
 
+  import ExUnit.CaptureLog
+
+  defmodule TimeoutModule do
+    def arbitrary_fn_name(_, _) do
+      :timer.sleep(2000)
+    end
+  end
+
   defmodule Schema do
     use Absinthe.Schema
 
@@ -60,6 +68,17 @@ defmodule Absinthe.Middleware.BatchTest do
           end)
         end
       end
+
+      field :timeout, :string do
+        resolve fn _, _, _ ->
+          batch(
+            {TimeoutModule, :arbitrary_fn_name, %{arbitrary: :data}},
+            nil,
+            fn batch -> {:ok, batch} end,
+            timeout: 1
+          )
+        end
+      end
     end
 
     def by_id(_, ids) do
@@ -150,4 +169,24 @@ defmodule Absinthe.Middleware.BatchTest do
 
     assert {:ok, %{data: %{"ctx" => "some_value"}}} == Absinthe.run(doc, Schema)
   end
+
+  test "when batch task timeouts it logs batching options" do
+    doc = """
+    {timeout}
+    """
+
+    assert capture_log(fn ->
+             pid =
+               spawn(fn ->
+                 Absinthe.run(doc, Schema)
+               end)
+
+             wait_for_process_to_exit(pid)
+           end) =~ "fn: {Absinthe.Middleware.BatchTest.TimeoutModule, :arbitrary_fn_name, %{arbitrary: :data}}"
+  end
+
+  defp wait_for_process_to_exit(pid) do
+    :timer.sleep(1)
+    if Process.alive?(pid), do: wait_for_process_to_exit(pid)
+  end
 end