Dataloader integration

README.md

@@ -58,6 +58,25 @@ end
 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 are using `DocumentProvider` modules, you will need to integrate into
+their `pipeline/1` callback as well. If your `DocumentProvider` modules do not
+yet override this callback, then this is fairly straightforward:
+
+```elixir
+def pipeline(%{pipeline: as_configured}) do
+  as_configured
+  |> Absinthe.Pipeline.from(__absinthe_plug_doc__(:remaining_pipeline))
+  |> Opencensus.Absinthe.add_schema_phases()
+end
+```
+
+If you already override the `pipeline/1` callback, just append this to the end:
+
+```elixir
+# ... result
+|> Opencensus.Absinthe.add_schema_phases()
+```
+
 ### Middleware
 
 Your [middleware callback][c:middleware/3] needs to run its output through
@@ -82,6 +101,14 @@ end
 
 [c:middleware/3]: https://hexdocs.pm/absinthe/Absinthe.Schema.html#c:middleware/3
 
+If you're using [`Dataloader`][dataloader], you should 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,29 @@ 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`, you should use the provided
+  `Opencensus.Absinthe.Middleware.Dataloader` Absinthe plugin module in place of
+  the default one for tracing batched resolutions.
+
+  If you are using `DocumentProvider` modules, you will need to integrate into
+  their `pipeline/1` callback as well. If your `DocumentProvider` modules do not
+  yet override this callback, then this is fairly straightforward:
+
+  ```elixir
+  def pipeline(%{pipeline: as_configured}) do
+    as_configured
+    |> Absinthe.Pipeline.from(__absinthe_plug_doc__(:remaining_pipeline))
+    |> Opencensus.Absinthe.add_schema_phases()
+  end
+  ```
+
+  If you already override the `pipeline/1` callback, just append this to the end:
+
+  ```elixir
+  # ... result
+  |> Opencensus.Absinthe.add_schema_phases()
+  ```
   """
 
   alias Absinthe.Middleware
@@ -68,6 +91,28 @@ defmodule Opencensus.Absinthe do
     )
   end
 
+  @doc """
+  Add tracing phases to an existing pipeline for schema.
+
+  ```elixir
+  pipeline =
+    Absinthe.Pipeline.for_document(schema, pipeline_opts)
+    |> Opencensus.Absinthe.add_schema_phases()
+  ```
+  """
+  @spec add_schema_phases(Absinthe.Pipeline.t()) :: Absinthe.Pipeline.t()
+  def add_schema_phases(pipeline) do
+    pipeline
+    |> Absinthe.Pipeline.insert_after(
+      Absinthe.Phase.Schema,
+      Opencensus.Absinthe.Phase.SchemaPush
+    )
+    |> Absinthe.Pipeline.insert_after(
+      Absinthe.Phase.Document.Result,
+      Opencensus.Absinthe.Phase.Pop
+    )
+  end
+
   @doc """
   Add tracing middleware for field resolution.
 

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

lib/opencensus/absinthe/phase/push.ex

@@ -9,8 +9,9 @@ defmodule Opencensus.Absinthe.Phase.Push do
 
   @impl true
   @spec run(Blueprint.t(), keyword()) :: Phase.result_t()
-  def run(blueprint, _opts \\ []) do
-    parent_span_ctx = :ocp.with_child_span("Blueprint")
+  def run(blueprint, opts \\ []) do
+    child_span = Keyword.get(opts, :child_span, "Blueprint")
+    parent_span_ctx = :ocp.with_child_span(child_span)
     span_ctx = :ocp.current_span_ctx()
 
     acc = %Acc{

lib/opencensus/absinthe/phase/schema_push.ex

@@ -0,0 +1,15 @@
+defmodule Opencensus.Absinthe.Phase.SchemaPush do
+  @moduledoc false
+
+  use Absinthe.Phase
+
+  alias Absinthe.Blueprint
+  alias Absinthe.Phase
+
+  @impl true
+  @spec run(Blueprint.t(), keyword()) :: Phase.result_t()
+  def run(blueprint, opts \\ []) do
+    opts = Keyword.put_new(opts, :child_span, "Blueprint")
+    Opencensus.Absinthe.Phase.Push.run(blueprint, opts)
+  end
+end

mix.exs

@@ -52,6 +52,7 @@ defmodule Opencensus.Absinthe.MixProject do
     [
       {:absinthe, "~> 1.4.0"},
       {:absinthe_plug, "~> 1.4.0", optional: true},
+      {:dataloader, "~> 1.0.0", optional: true},
       {:credo, "~> 1.1.0", only: :test},
       {:dialyxir, "~> 1.0.0-rc.6", only: :dev, runtime: false},
       {:ex_doc, ">= 0.0.0", only: :docs},

mix.lock

@@ -4,8 +4,9 @@
   "bunt": {:hex, :bunt, "0.2.0", "951c6e801e8b1d2cbe58ebbd3e616a869061ddadcc4863d0a2182541acae9a38", [:mix], [], "hexpm"},
   "certifi": {:hex, :certifi, "2.5.1", "867ce347f7c7d78563450a18a6a28a8090331e77fa02380b4a21962a65d36ee5", [:rebar3], [{:parse_trans, "~>3.3", [hex: :parse_trans, repo: "hexpm", optional: false]}], "hexpm"},
   "counters": {:hex, :counters, "0.2.1", "aa3d97e88f92573488987193d0f48efce0f3b2cd1443bf4ee760bc7f99322f0c", [:mix, :rebar3], [], "hexpm"},
-  "credo": {:hex, :credo, "1.1.0", "e0c07b2fd7e2109495f582430a1bc96b2c71b7d94c59dfad120529f65f19872f", [:mix], [{:bunt, "~> 0.2.0", [hex: :bunt, repo: "hexpm", optional: false]}, {:jason, "~> 1.0", [hex: :jason, repo: "hexpm", optional: false]}], "hexpm"},
+  "credo": {:hex, :credo, "1.1.2", "02b6422f3e659eb74b05aca3c20c1d8da0119a05ee82577a82e6c2938bf29f81", [: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.2", "cb0e6878fdf86dc63509eaf2233a71fa73fc383c8362c8ff8e8b6f0c2bb7017c", [:mix], [], "hexpm"},