[GRAPH-1102] Customizable doc storage with default implementation

lib/absinthe/subscription.ex

@@ -143,73 +143,56 @@ defmodule Absinthe.Subscription do
   def subscribe(pubsub, field_keys, doc_id, doc) do
     field_keys = List.wrap(field_keys)
 
-    registry = pubsub |> registry_name
-
     doc_value = %{
       initial_phases: PipelineSerializer.pack(doc.initial_phases),
       source: doc.source
     }
 
-    pdict_add_fields(doc_id, field_keys)
-
-    for field_key <- field_keys do
-      {:ok, _} = Registry.register(registry, field_key, doc_id)
-    end
-
-    {:ok, _} = Registry.register(registry, doc_id, doc_value)
-  end
-
-  defp pdict_fields(doc_id) do
-    Process.get({__MODULE__, doc_id}, [])
-  end
-
-  defp pdict_add_fields(doc_id, field_keys) do
-    Process.put({__MODULE__, doc_id}, field_keys ++ pdict_fields(doc_id))
-  end
-
-  defp pdict_delete_fields(doc_id) do
-    Process.delete({__MODULE__, doc_id})
+    storage_module = document_storage(pubsub)
+    storage_process_name = document_storage_name(pubsub)
+    storage_module.put(storage_process_name, doc_id, doc_value, field_keys)
   end
 
   @doc false
   def unsubscribe(pubsub, doc_id) do
-    registry = pubsub |> registry_name
-
-    for field_key <- pdict_fields(doc_id) do
-      Registry.unregister(registry, field_key)
-    end
-
-    Registry.unregister(registry, doc_id)
-
-    pdict_delete_fields(doc_id)
-    :ok
+    storage_module = document_storage(pubsub)
+    storage_process_name = document_storage_name(pubsub)
+    storage_module.delete(storage_process_name, doc_id)
   end
 
   @doc false
   def get(pubsub, key) do
-    name = registry_name(pubsub)
-
-    name
-    |> Registry.lookup(key)
-    |> MapSet.new(fn {_pid, doc_id} -> doc_id end)
-    |> Enum.reduce(%{}, fn doc_id, acc ->
-      case Registry.lookup(name, doc_id) do
-        [] ->
-          acc
-
-        [{_pid, doc} | _rest] ->
-          Map.put_new_lazy(acc, doc_id, fn ->
-            Map.update!(doc, :initial_phases, &PipelineSerializer.unpack/1)
-          end)
-      end
+    storage_module = document_storage(pubsub)
+    storage_process_name = document_storage_name(pubsub)
+
+    storage_process_name
+    |> storage_module.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()
   end
 
   @doc false
   def registry_name(pubsub) do
     Module.concat([pubsub, :Registry])
   end
 
+  @doc false
+  def document_storage_name(pubsub) do
+    Module.concat([pubsub, :Storage])
+  end
+
+  def document_storage(pubsub) do
+    {:ok, document_storage} =
+      pubsub
+      |> registry_name
+      |> Registry.meta(:document_storage)
+
+    document_storage
+  end
+
   @doc false
   def publish_remote(pubsub, mutation_result, subscribed_fields) do
     {:ok, pool_size} =

lib/absinthe/subscription/default_document_storage.ex

@@ -0,0 +1,65 @@
+defmodule Absinthe.Subscription.DefaultDocumentStorage do
+  @behaviour Absinthe.Subscription.DocumentStorage
+
+  @moduledoc """
+  Default document storage for Absinthe. Stores subscription
+  documents and field keys in a Registry process.
+  """
+
+  @impl Absinthe.Subscription.DocumentStorage
+  def child_spec(opts) do
+    Registry.child_spec(opts)
+  end
+
+  @impl Absinthe.Subscription.DocumentStorage
+  def put(storage_process_name, doc_id, doc_value, field_keys) do
+    pdict_add_fields(doc_id, field_keys)
+
+    for field_key <- field_keys do
+      {:ok, _} = Registry.register(storage_process_name, field_key, doc_id)
+    end
+
+    {:ok, _} = Registry.register(storage_process_name, doc_id, doc_value)
+  end
+
+  @impl Absinthe.Subscription.DocumentStorage
+  def delete(storage_process_name, doc_id) do
+    for field_key <- pdict_fields(doc_id) do
+      Registry.unregister(storage_process_name, field_key)
+    end
+
+    pdict_delete_fields(doc_id)
+
+    Registry.unregister(storage_process_name, doc_id)
+
+    :ok
+  end
+
+  @impl Absinthe.Subscription.DocumentStorage
+  def get_docs_by_field_key(storage_process_name, field_key) do
+    storage_process_name
+    |> Registry.lookup(field_key)
+    |> MapSet.new(fn {_pid, doc_id} -> doc_id end)
+    |> Enum.reduce(%{}, fn doc_id, acc ->
+      case Registry.lookup(storage_process_name, doc_id) do
+        [] ->
+          acc
+
+        [{_pid, doc} | _rest] ->
+          Map.put_new(acc, doc_id, doc)
+      end
+    end)
+  end
+
+  defp pdict_fields(doc_id) do
+    Process.get({__MODULE__, doc_id}, [])
+  end
+
+  defp pdict_add_fields(doc_id, field_keys) do
+    Process.put({__MODULE__, doc_id}, field_keys ++ pdict_fields(doc_id))
+  end
+
+  defp pdict_delete_fields(doc_id) do
+    Process.delete({__MODULE__, doc_id})
+  end
+end

lib/absinthe/subscription/document_storage.ex

@@ -0,0 +1,44 @@
+defmodule Absinthe.Subscription.DocumentStorage do
+  @moduledoc """
+  Behaviour for storing subscription documents. Used to tell
+  Absinthe how to store documents and the field keys associated with those
+  documents.
+  """
+
+  @doc """
+  Child spec to determine how to start the
+  Document storage process. This will be supervised. Absinthe will give
+  the process a name and that name will be passed in the other callbacks
+  in order to reference it there.
+  """
+  @callback child_spec(opts :: Keyword.t()) :: Supervisor.child_spec()
+
+  @doc """
+  Adds `doc` to storage with `doc_id` as the key. Associates the given
+  `field_keys` with `doc_id`.
+  """
+  @callback put(
+              storage_process_name :: atom,
+              doc_id :: term,
+              doc :: %{
+                initial_phases: Absinthe.Subscription.PipelineSerializer.packed_pipeline(),
+                source: binary()
+              },
+              field_keys :: [{field :: term, key :: term}]
+            ) ::
+              {:ok, term} | {:error, :reason}
+
+  @doc """
+  Removes the document. Along with any field_keys associated with it
+  """
+  @callback delete(storage_process_name :: atom, doc_id :: term) :: :ok
+
+  @doc """
+  Get all docs associated with `field_key`
+  """
+  @callback get_docs_by_field_key(
+              storage_process_name :: atom,
+              field_key :: {field :: term, key :: term}
+            ) ::
+              map()
+end

lib/absinthe/subscription/supervisor.ex

@@ -24,22 +24,43 @@ defmodule Absinthe.Subscription.Supervisor do
     pool_size = Keyword.get(opts, :pool_size, System.schedulers_online() * 2)
     compress_registry? = Keyword.get(opts, :compress_registry?, true)
 
-    Supervisor.start_link(__MODULE__, {pubsub, pool_size, compress_registry?})
+    document_storage =
+      Keyword.get(opts, :document_storage, Absinthe.Subscription.DefaultDocumentStorage)
+
+    storage_opts =
+      case document_storage do
+        Absinthe.Subscription.DefaultDocumentStorage ->
+          [
+            keys: :duplicate,
+            partitions: System.schedulers_online(),
+            compressed: compress_registry?
+          ]
+
+        _ ->
+          Keyword.get(opts, :storage_opts, Keyword.new())
+      end
+
+    Supervisor.start_link(
+      __MODULE__,
+      {pubsub, pool_size, document_storage, storage_opts}
+    )
   end
 
-  def init({pubsub, pool_size, compress_registry?}) do
+  def init({pubsub, pool_size, document_storage, storage_opts}) do
     registry_name = Absinthe.Subscription.registry_name(pubsub)
-    meta = [pool_size: pool_size]
+    meta = [pool_size: pool_size, document_storage: document_storage]
+
+    storage_opts =
+      Keyword.put(storage_opts, :name, Absinthe.Subscription.document_storage_name(pubsub))
 
     children = [
       {Registry,
        [
-         keys: :duplicate,
+         keys: :unique,
          name: registry_name,
-         partitions: System.schedulers_online(),
-         meta: meta,
-         compressed: compress_registry?
+         meta: meta
        ]},
+      document_storage.child_spec(storage_opts),
       {Absinthe.Subscription.ProxySupervisor, [pubsub, registry_name, pool_size]}
     ]