Dataloader integration

README.md

@@ -82,6 +82,14 @@ end
 
 [c:middleware/3]: https://hexdocs.pm/absinthe/Absinthe.Schema.html#c:middleware/3
 
+If you're using [`Dataloader`][dataloader], you will want to use the provided
+`Opencensus.Absinthe.Middleware.Dataloader` Absinthe plugin module in place of
+the default one for tracing batched resolutions. See the [module
+docs][internal_dataloader] for details.
+
+[dataloader]: https://github.com/absinthe-graphql/dataloader
+[internal_dataloader]: https://hexdocs.pm/opencensus_absinthe/Opencensus.Absinthe.Middleware.Dataloader.html
+
 ### Schema
 
 Until Absinthe merge and publish their telemetry support (see below) _and_

lib/opencensus/absinthe.ex

@@ -41,6 +41,14 @@ defmodule Opencensus.Absinthe do
 
   Worst case, you'll need to copy the code from the current `pipeline` target and add a call to
   `Opencensus.Absinthe.add_phases/1` as above.
+
+  If you're using [`Dataloader`][dataloader], you will want to use the provided
+  `Opencensus.Absinthe.Middleware.Dataloader` Absinthe plugin module in place of
+  the default one for tracing batched resolutions. See the [module
+  docs][internal_dataloader] for details.
+
+  [dataloader]: https://github.com/absinthe-graphql/dataloader
+  [internal_dataloader]: https://hexdocs.pm/opencensus_absinthe/Opencensus.Absinthe.Middleware.Dataloader.html
   """
 
   alias Absinthe.Middleware

lib/opencensus/absinthe/acc.ex

@@ -3,6 +3,7 @@ defmodule Opencensus.Absinthe.Acc do
 
   alias Absinthe.Blueprint
   alias Absinthe.Resolution
+  alias Absinthe.Blueprint.Execution
 
   @accumulator_key :opencensus_absinthe
 
@@ -16,12 +17,17 @@ defmodule Opencensus.Absinthe.Acc do
 
   @spec set(Blueprint.t(), any()) :: map()
   def set(%Blueprint{} = bp, our_acc) do
-    acc = bp.execution.acc |> Map.put(@accumulator_key, our_acc)
-    put_in(bp.execution.acc, acc)
+    put_in(bp.execution.acc, Map.put(bp.execution.acc, @accumulator_key, our_acc))
+  end
+
+  @spec set(Execution.t(), any()) :: map()
+  def set(%Execution{} = exec, our_acc) do
+    put_in(exec.acc, Map.put(exec.acc, @accumulator_key, our_acc))
   end
 
   @spec get(Blueprint.t() | Resolution.t()) :: t()
   def get(blueprint_or_resolution)
   def get(%Blueprint{} = bp), do: bp.execution.acc[@accumulator_key]
   def get(%Resolution{} = r), do: r.acc[@accumulator_key]
+  def get(%Execution{} = exec), do: exec.acc[@accumulator_key]
 end

lib/opencensus/absinthe/middleware.ex

@@ -12,15 +12,20 @@ defmodule Opencensus.Absinthe.Middleware do
   @impl true
   @spec call(Resolution.t(), term()) :: Resolution.t()
   def call(%Resolution{state: :unresolved} = resolution, field: field) do
-    acc = Acc.get(resolution)
+    case Acc.get(resolution) do
+      # nil ->
+      #   Logger.error("Handling tracing for a field with no span metadata: #{inspect(field, pretty: true)}")
+      #   resolution
 
-    span_options = %{
-      attributes: field |> extract_metadata() |> Enum.into(%{}, &stringify_keys/1)
-    }
+      acc ->
+        span_options = %{
+          attributes: field |> extract_metadata() |> Enum.into(%{}, &stringify_keys/1)
+        }
 
-    span_ctx = :oc_trace.start_span(field |> repr(), acc.span_ctx, span_options)
-    middleware = resolution.middleware ++ [{{__MODULE__, :on_complete}, span_ctx: span_ctx}]
-    %{resolution | middleware: middleware}
+        span_ctx = :oc_trace.start_span(field |> repr(), acc.span_ctx, span_options)
+        middleware = resolution.middleware ++ [{{__MODULE__, :on_complete}, span_ctx: span_ctx}]
+        %{resolution | middleware: middleware}
+    end
   end
 
   @doc false

lib/opencensus/absinthe/middleware/dataloader.ex

@@ -0,0 +1,72 @@
+if Code.ensure_loaded?(Dataloader) do
+  defmodule Opencensus.Absinthe.Middleware.Dataloader do
+    @moduledoc """
+    This is a small extension on top of `Absinthe.Middleware.Dataloader` that
+    will create spans for each resolution.
+
+    ## Usage
+
+    In your Absinthe schema, simply override the `plugins/0` callback (if you're
+    not already) and prepend this plugin to the list:
+
+        def plugins do
+          [Opencensus.Absinthe.Middleware.Dataloader | Absinthe.Plugin.defaults()]
+        end
+    """
+
+    @behaviour Absinthe.Middleware
+    @behaviour Absinthe.Plugin
+
+    @span_key :dataloader_resolution_span_ctx
+    @counter_key :dataloader_resolution_counter
+
+    alias Opencensus.Absinthe.Acc
+    alias Absinthe.Middleware.Dataloader, as: DefaultDataloader
+
+    @doc """
+    The `Absinthe.Plugin` callback. Starts the OpenCensus span.
+    """
+    def before_resolution(exec) do
+      span_options = %{attributes: %{}}
+      acc = Acc.get(exec)
+
+      {counter, new_acc} =
+        Map.get_and_update(acc, @counter_key, fn cur ->
+          case cur do
+            nil -> {cur, 1}
+            x -> {x, x + 1}
+          end
+        end)
+
+      span_ctx = :oc_trace.start_span("resolution_#{counter || 0}", acc.span_ctx, span_options)
+      new_acc = Map.put(new_acc, @span_key, span_ctx)
+
+      exec
+      |> Acc.set(new_acc)
+      |> DefaultDataloader.before_resolution()
+    end
+
+    @doc """
+    The `Absinthe.Plugin` callback. Finishes the OpenCensus span.
+    """
+    def after_resolution(exec) do
+      acc = Acc.get(exec)
+
+      acc
+      |> Map.get(@span_key)
+      |> :oc_trace.finish_span()
+
+      acc =
+        exec
+        |> Acc.get()
+        |> Map.delete(@span_key)
+
+      exec
+      |> Acc.set(acc)
+      |> DefaultDataloader.after_resolution()
+    end
+
+    def call(resolution, callback), do: DefaultDataloader.call(resolution, callback)
+    def pipeline(pipeline, exec), do: DefaultDataloader.pipeline(pipeline, exec)
+  end
+end

mix.exs

@@ -47,9 +47,10 @@ defmodule Opencensus.Absinthe.MixProject do
 
   defp deps() do
     [
-      {:absinthe_plug, "~> 1.4.0", only: :dev, runtime: false},
+      {:absinthe_plug, "~> 1.4.0", only: [:dev, :docs], runtime: false},
       {:absinthe, "~> 1.4.0"},
       {:credo, "~> 0.10.0", only: [:dev, :test], runtime: false},
+      {:dataloader, "~> 1.0.0", optional: true},
       {:dialyxir, "~> 1.0.0-rc.6", only: :dev, runtime: false},
       {:ex_doc, ">= 0.0.0", only: [:dev, :docs], runtime: false},
       {:excoveralls, "~> 0.10.6", only: [:dev, :test], runtime: false},

mix.lock

@@ -6,6 +6,7 @@
   "counters": {:hex, :counters, "0.2.1", "aa3d97e88f92573488987193d0f48efce0f3b2cd1443bf4ee760bc7f99322f0c", [:mix, :rebar3], [], "hexpm"},
   "credo": {:hex, :credo, "0.10.2", "03ad3a1eff79a16664ed42fc2975b5e5d0ce243d69318060c626c34720a49512", [:mix], [{:bunt, "~> 0.2.0", [hex: :bunt, repo: "hexpm", optional: false]}, {:jason, "~> 1.0", [hex: :jason, repo: "hexpm", optional: false]}], "hexpm"},
   "ctx": {:hex, :ctx, "0.5.0", "78e0f16712e12d707a7f34277381b8e193d7c71eaa24d37330dc02477c09eda5", [:rebar3], [], "hexpm"},
+  "dataloader": {:hex, :dataloader, "1.0.6", "fb724d6d3fb6acb87d27e3b32dea3a307936ad2d245faf9cf5221d1323d6a4ba", [:mix], [{:ecto, ">= 0.0.0", [hex: :ecto, repo: "hexpm", optional: true]}], "hexpm"},
   "dialyxir": {:hex, :dialyxir, "1.0.0-rc.6", "78e97d9c0ff1b5521dd68041193891aebebce52fc3b93463c0a6806874557d7d", [:mix], [{:erlex, "~> 0.2.1", [hex: :erlex, repo: "hexpm", optional: false]}], "hexpm"},
   "earmark": {:hex, :earmark, "1.3.2", "b840562ea3d67795ffbb5bd88940b1bed0ed9fa32834915125ea7d02e35888a5", [:mix], [], "hexpm"},
   "erlex": {:hex, :erlex, "0.2.1", "cee02918660807cbba9a7229cae9b42d1c6143b768c781fa6cee1eaf03ad860b", [:mix], [], "hexpm"},