From f4db5d9762069d29a533f942d2dc090afcce0bde Mon Sep 17 00:00:00 2001
From: Bryan Joseph <bryanj@adobe.com>
Date: Thu, 20 Jun 2024 15:30:37 -0500
Subject: [PATCH] telemetry

---
 lib/absinthe/subscription.ex                  | 70 ++++++++++++++++---
 lib/absinthe/subscription/document_storage.ex |  1 -
 2 files changed, 61 insertions(+), 10 deletions(-)

diff --git a/lib/absinthe/subscription.ex b/lib/absinthe/subscription.ex
index 105f78af..3a31c810 100644
--- a/lib/absinthe/subscription.ex
+++ b/lib/absinthe/subscription.ex
@@ -149,26 +149,78 @@ defmodule Absinthe.Subscription do
     }
 
     storage_implementation = storage_implementation(pubsub)
-    storage_implementation.subscribe(pubsub, doc_id, doc_value, field_keys)
+
+    :telemetry.span(
+      [:absinthe, :subscription, :storage, :subscribe],
+      %{
+        doc_id: doc_id,
+        doc: doc,
+        field_keys: field_keys,
+        storage_implementation: storage_implementation
+      },
+      fn ->
+        result = storage_implementation.subscribe(pubsub, doc_id, doc_value, field_keys)
+
+        {result,
+         %{
+           doc_id: doc_id,
+           doc: doc,
+           field_keys: field_keys,
+           storage_implementation: storage_implementation
+         }}
+      end
+    )
   end
 
   @doc false
   def unsubscribe(pubsub, doc_id) do
     storage_implementation = storage_implementation(pubsub)
-    storage_implementation.unsubscribe(pubsub, doc_id)
+
+    :telemetry.span(
+      [:absinthe, :subscription, :storage, :unsubscribe],
+      %{
+        doc_id: doc_id,
+        storage_implementation: storage_implementation
+      },
+      fn ->
+        result = storage_implementation.unsubscribe(pubsub, doc_id)
+
+        {result,
+         %{
+           doc_id: doc_id,
+           storage_implementation: storage_implementation
+         }}
+      end
+    )
   end
 
   @doc false
   def get(pubsub, key) do
     storage_implementation = storage_implementation(pubsub)
 
-    pubsub
-    |> storage_implementation.get_docs_by_field_key(key)
-    |> Enum.map(fn {doc_id, %{initial_phases: initial_phases} = doc} ->
-      initial_phases = PipelineSerializer.unpack(initial_phases)
-      {doc_id, Map.put(doc, :initial_phases, initial_phases)}
-    end)
-    |> Map.new()
+    :telemetry.span(
+      [:absinthe, :subscription, :storage, :get],
+      %{
+        key: key,
+        storage_implementation: storage_implementation
+      },
+      fn ->
+        result =
+          pubsub
+          |> storage_implementation.get_docs_by_field_key(key)
+          |> Enum.map(fn {doc_id, %{initial_phases: initial_phases} = doc} ->
+            initial_phases = PipelineSerializer.unpack(initial_phases)
+            {doc_id, Map.put(doc, :initial_phases, initial_phases)}
+          end)
+          |> Map.new()
+
+        {result,
+         %{
+           key: key,
+           storage_implementation: storage_implementation
+         }}
+      end
+    )
   end
 
   @doc false
diff --git a/lib/absinthe/subscription/document_storage.ex b/lib/absinthe/subscription/document_storage.ex
index 4dcb1234..f565816c 100644
--- a/lib/absinthe/subscription/document_storage.ex
+++ b/lib/absinthe/subscription/document_storage.ex
@@ -4,7 +4,6 @@ defmodule Absinthe.Subscription.DocumentStorage do
   Absinthe how to store documents and the field keys subcribed to those
   documents.
   """
-  alias Module.Behaviour
 
   @doc """
   Child spec to determine how to start the