rework state and ops, store descriptor once (#68)

Author: mjheilmann

lib/grpc_reflection/server.ex

@@ -33,12 +33,12 @@ defmodule GrpcReflection.Server do
       end
 
       @doc """
-      Get the reflection reponse containing the given symbol, if it is exposed by a configured service
+      Get the reflection filename for the given symbol, if it is known
       """
-      @spec get_by_symbol(binary()) ::
-              {:ok, GrpcReflection.Server.descriptor_t()} | {:error, binary}
-      def get_by_symbol(symbol) do
-        Service.get_by_symbol(@cfg, symbol)
+      @spec get_filename_by_symbol(binary()) ::
+              {:ok, String.t()} | {:error, binary}
+      def get_filename_by_symbol(symbol) do
+        Service.get_filename_by_symbol(@cfg, symbol)
       end
 
       @doc """

lib/grpc_reflection/server/process.ex

@@ -11,8 +11,9 @@ defmodule GrpcReflection.Server.Process do
   end
 
   def reflect(state_mod, {:file_containing_symbol, symbol}) do
-    with {:ok, description} <- state_mod.get_by_symbol(symbol) do
-      {:ok, {:file_descriptor_response, description}}
+    with {:ok, filename} <- state_mod.get_filename_by_symbol(symbol),
+         {:ok, descriptor} <- state_mod.get_by_filename(filename) do
+      {:ok, {:file_descriptor_response, descriptor}}
     end
   end
 

lib/grpc_reflection/service.ex

@@ -11,7 +11,7 @@ defmodule GrpcReflection.Service do
   defdelegate put_state(cfg, state), to: Agent
 
   defdelegate list_services(cfg), to: Agent
-  defdelegate get_by_symbol(cfg, symbol), to: Agent
+  defdelegate get_filename_by_symbol(cfg, symbol), to: Agent
   defdelegate get_by_filename(cfg, filename), to: Agent
   defdelegate get_by_extension(cfg, containing_type), to: Agent
   defdelegate get_extension_numbers_by_type(cfg, mod), to: Agent

lib/grpc_reflection/service/agent.ex

@@ -30,8 +30,9 @@ defmodule GrpcReflection.Service.Agent do
     Agent.get(name, &State.lookup_services/1)
   end
 
-  @spec get_by_symbol(cfg_t(), binary()) :: {:ok, State.descriptor_t()} | {:error, binary}
-  def get_by_symbol(cfg, symbol) do
+  @spec get_filename_by_symbol(cfg_t(), binary()) ::
+          {:ok, State.descriptor_t()} | {:error, binary}
+  def get_filename_by_symbol(cfg, symbol) do
     name = start_agent_on_first_call(cfg)
     Agent.get(name, &State.lookup_symbol(symbol, &1))
   end

lib/grpc_reflection/service/builder.ex

@@ -9,10 +9,12 @@ defmodule GrpcReflection.Service.Builder do
   def build_reflection_tree(services) do
     with :ok <- Util.validate_services(services) do
       tree =
-        Enum.reduce(services, State.new(services), fn service, state ->
+        services
+        |> Enum.reduce(State.new(), fn service, state ->
           new_state = process_service(service)
           State.merge(state, new_state)
         end)
+        |> State.group_symbols_by_namespace()
 
       {:ok, tree}
     end
@@ -22,13 +24,14 @@ defmodule GrpcReflection.Service.Builder do
     service_name = service.__meta__(:name)
     service_response = build_response(service_name, service)
 
-    State.new()
-    |> State.add_symbols(%{service_name => service_response})
-    |> State.add_files(%{(service_name <> ".proto") => service_response})
-    |> trace_service_refs(service)
+    [service]
+    |> State.new()
+    |> State.add_file(service_response)
+    |> State.add_symbol(service_name, service_response.name)
+    |> trace_service_refs(service, service_response.name)
   end
 
-  defp trace_service_refs(state, module) do
+  defp trace_service_refs(state, module, module_filename) do
     service_name = module.__meta__(:name)
     methods = get_descriptor(module).method
 
@@ -40,23 +43,21 @@ defmodule GrpcReflection.Service.Builder do
         Enum.find(methods, fn method -> method.name == to_string(function) end)
 
       call_symbol = service_name <> "." <> to_string(function)
-      call_response = build_response(service_name, module)
+
       req_symbol = Util.trim_symbol(req_symbol)
       req_response = build_response(req_symbol, request)
+      req_filename = req_symbol <> ".proto"
       resp_symbol = Util.trim_symbol(resp_symbol)
       resp_response = build_response(resp_symbol, response)
+      resp_filename = resp_symbol <> ".proto"
 
       state
       |> Extensions.add_extensions(service_name, module)
-      |> State.add_symbols(%{
-        call_symbol => call_response,
-        req_symbol => req_response,
-        resp_symbol => resp_response
-      })
-      |> State.add_files(%{
-        (req_symbol <> ".proto") => req_response,
-        (resp_symbol <> ".proto") => resp_response
-      })
+      |> State.add_file(req_response)
+      |> State.add_file(resp_response)
+      |> State.add_symbol(call_symbol, module_filename)
+      |> State.add_symbol(req_symbol, req_filename)
+      |> State.add_symbol(resp_symbol, resp_filename)
       |> Extensions.add_extensions(req_symbol, request)
       |> Extensions.add_extensions(resp_symbol, response)
       |> trace_message_refs(req_symbol, request)
@@ -103,8 +104,8 @@ defmodule GrpcReflection.Service.Builder do
 
       state
       |> Extensions.add_extensions(symbol, mod)
-      |> State.add_symbols(%{symbol => response})
-      |> State.add_files(%{(symbol <> ".proto") => response})
+      |> State.add_file(response)
+      |> State.add_symbol(symbol, response.name)
       |> trace_message_refs(symbol, mod)
     end)
   end

lib/grpc_reflection/service/builder/extensions.ex

@@ -10,7 +10,7 @@ defmodule GrpcReflection.Service.Builder.Extensions do
     case process_extensions(module, symbol, extension_file, module.descriptor()) do
       {:ok, {extension_numbers, extension_payload}} ->
         state
-        |> State.add_files(%{extension_file => extension_payload})
+        |> State.add_file(extension_payload)
         |> State.add_extensions(%{symbol => extension_numbers})
 
       :ignore ->

lib/grpc_reflection/service/state.ex

@@ -3,13 +3,14 @@ defmodule GrpcReflection.Service.State do
 
   defstruct services: [], files: %{}, symbols: %{}, extensions: %{}
 
-  @type descriptor_t :: GrpcReflection.Server.descriptor_t()
+  @type descriptor_t :: Google.Protobuf.FileDescriptorProto.t()
+  @type symbol_t :: %{optional(binary()) => binary()}
   @type entry_t :: %{optional(binary()) => descriptor_t()}
 
   @type t :: %__MODULE__{
           services: list(module()),
           files: entry_t(),
-          symbols: entry_t(),
+          symbols: symbol_t(),
           extensions: %{optional(binary()) => list(integer())}
         }
 
@@ -26,6 +27,7 @@ defmodule GrpcReflection.Service.State do
     }
   end
 
+  # merge two maps, but raise if we have duplicate keys with different values
   defp merge_map!(type, current_map, incoming_map) do
     Map.merge(current_map, incoming_map, fn
       _key, value, value ->
@@ -36,14 +38,27 @@ defmodule GrpcReflection.Service.State do
     end)
   end
 
-  @spec add_files(t(), entry_t()) :: t()
-  def add_files(%__MODULE__{} = state, files) do
-    %{state | files: Map.merge(files, state.files)}
+  # add an item to a map, but raise if the key exists with a different value
+  defp put_map!(type, current_map, key, value) do
+    case current_map[key] do
+      nil -> Map.put(current_map, key, value)
+      ^value -> current_map
+      _ -> raise "#{type} Conflict detected: key #{key} present with multiple values"
+    end
   end
 
-  @spec add_symbols(t(), entry_t()) :: t()
-  def add_symbols(%__MODULE__{} = state, symbols) do
-    %{state | symbols: Map.merge(symbols, state.symbols)}
+  @spec add_file(t(), descriptor_t()) :: t()
+  def add_file(%__MODULE__{} = state, descriptor) do
+    %{state | files: put_map!("Files", state.files, descriptor.name, descriptor)}
+  end
+
+  @spec add_symbol(t(), String.t(), String.t()) :: t()
+  def add_symbol(%__MODULE__{files: files} = state, symbol, filename) do
+    if files[filename] == nil do
+      raise "Symbol #{symbol} added for file that doesn't exist: #{filename}"
+    end
+
+    %{state | symbols: put_map!("Symbols", state.symbols, symbol, filename)}
   end
 
   def add_extensions(%__MODULE__{} = state, extensions) do
@@ -100,4 +115,10 @@ defmodule GrpcReflection.Service.State do
       {:error, "extension numbers not found"}
     end
   end
+
+  def group_symbols_by_namespace(%__MODULE__{} = state) do
+    # group symbols by namespace and combine
+    # IO.inspect(state)
+    state
+  end
 end

test/service/builder_test.exs

@@ -16,8 +16,6 @@ defmodule GrpcReflection.Service.BuilderTest do
              "google.protobuf.Timestamp.proto",
              "testserviceV3.Enum.proto",
              "testserviceV3.TestReply.proto",
-             "testserviceV3.TestRequest.GEntry.proto",
-             "testserviceV3.TestRequest.Payload.Location.proto",
              "testserviceV3.TestRequest.Payload.proto",
              "testserviceV3.TestRequest.Token.proto",
              "testserviceV3.TestRequest.proto",
@@ -39,7 +37,7 @@ defmodule GrpcReflection.Service.BuilderTest do
              "testserviceV3.TestService.CallFunction"
            ]
 
-    Enum.each(Map.values(tree.files) ++ Map.values(tree.symbols), fn payload ->
+    Enum.each(Map.values(tree.files), fn payload ->
       assert payload.syntax == "proto3"
     end)
   end
@@ -53,7 +51,6 @@ defmodule GrpcReflection.Service.BuilderTest do
              "google.protobuf.Timestamp.proto",
              "testserviceV2.Enum.proto",
              "testserviceV2.TestReply.proto",
-             "testserviceV2.TestRequest.GEntry.proto",
              "testserviceV2.TestRequest.proto",
              "testserviceV2.TestRequestExtension.proto",
              "testserviceV2.TestService.proto"
@@ -77,7 +74,7 @@ defmodule GrpcReflection.Service.BuilderTest do
     # this is a bitstring that may contain whitespace characters
     assert extensions |> to_string() |> String.trim() == ""
 
-    Enum.each(Map.values(tree.files) ++ Map.values(tree.symbols), fn
+    Enum.each(Map.values(tree.files), fn
       %{name: "google" <> _, syntax: syntax} -> assert syntax == "proto3"
       %{name: _, syntax: syntax} -> assert syntax == "proto2"
     end)
@@ -87,8 +84,7 @@ defmodule GrpcReflection.Service.BuilderTest do
     assert {:ok, tree} = Builder.build_reflection_tree([EmptyService.Service])
     assert %State{services: [EmptyService.Service]} = tree
 
-    (Map.values(tree.files) ++ Map.values(tree.symbols))
-    |> Enum.each(fn payload ->
+    Enum.each(Map.values(tree.files), fn payload ->
       # empty services default to proto2
       assert payload.syntax == "proto2"
       assert payload.dependency == []
@@ -110,29 +106,17 @@ defmodule GrpcReflection.Service.BuilderTest do
     assert {:ok, tree} = Builder.build_reflection_tree([HLW.TestService.Service])
     assert %State{services: [HLW.TestService.Service]} = tree
 
-    names =
-      (Map.values(tree.files) ++ Map.values(tree.symbols))
-      |> Enum.map(& &1.name)
-
-    assert names ==
-             [
-               "google.protobuf.Any.proto",
-               "google.protobuf.Timestamp.proto",
-               "testserviceV2.Enum.proto",
-               "testserviceV2.TestReply.proto",
-               "testserviceV2.TestRequest.proto",
-               "testserviceV2.TestRequest.proto",
-               "testserviceV2.TestRequestExtension.proto",
-               "testserviceV2.TestService.proto",
-               "google.protobuf.Any.proto",
-               "google.protobuf.Timestamp.proto",
-               "testserviceV2.Enum.proto",
-               "testserviceV2.TestReply.proto",
-               "testserviceV2.TestRequest.proto",
-               "testserviceV2.TestRequest.proto",
-               "testserviceV2.TestService.proto",
-               "testserviceV2.TestService.proto"
-             ]
+    names = Enum.map(Map.values(tree.files), & &1.name)
+
+    assert names == [
+             "google.protobuf.Any.proto",
+             "google.protobuf.Timestamp.proto",
+             "testserviceV2.Enum.proto",
+             "testserviceV2.TestReply.proto",
+             "testserviceV2.TestRequest.proto",
+             "testserviceV2.TestRequestExtension.proto",
+             "testserviceV2.TestService.proto"
+           ]
   end
 
   test "handles a non-service module" do

test/service/server_test.exs

@@ -10,7 +10,7 @@ defmodule GrpcReflection.ServerTest do
   test "adding a service changes responses" do
     assert Service.list_services() == []
 
-    assert Service.get_by_symbol("helloworld.Greeter") ==
+    assert Service.get_filename_by_symbol("helloworld.Greeter") ==
              {:error, "symbol not found"}
 
     assert Service.get_by_filename("helloworld.Greeter.proto") ==
@@ -20,11 +20,11 @@ defmodule GrpcReflection.ServerTest do
 
     assert Service.list_services() == ["helloworld.Greeter"]
 
-    assert {:ok, descriptor} =
-             Service.get_by_symbol("helloworld.Greeter")
+    assert {:ok, filename} =
+             Service.get_filename_by_symbol("helloworld.Greeter")
 
-    assert {:ok, ^descriptor} =
-             Service.get_by_filename("helloworld.Greeter.proto")
+    assert {:ok, %Google.Protobuf.FileDescriptorProto{name: ^filename}} =
+             Service.get_by_filename(filename)
   end
 
   describe "reflection state testing" do
@@ -45,27 +45,27 @@ defmodule GrpcReflection.ServerTest do
     end
 
     test "method files return service descriptors" do
-      assert {:ok, descriptor} =
-               Service.get_by_symbol("helloworld.Greeter")
+      assert {:ok, filename} =
+               Service.get_filename_by_symbol("helloworld.Greeter")
 
-      assert {:ok, ^descriptor} =
-               Service.get_by_symbol("helloworld.Greeter.SayHello")
+      assert {:ok, %Google.Protobuf.FileDescriptorProto{name: ^filename}} =
+               Service.get_by_filename(filename)
     end
 
     test "describing a type returns the type" do
-      assert {:ok, descriptor} =
-               Service.get_by_symbol("helloworld.HelloRequest")
+      assert {:ok, filename} =
+               Service.get_filename_by_symbol("helloworld.HelloRequest")
 
-      assert {:ok, ^descriptor} =
-               Service.get_by_filename("helloworld.HelloRequest.proto")
+      assert {:ok, %Google.Protobuf.FileDescriptorProto{name: ^filename}} =
+               Service.get_by_filename(filename)
     end
 
     test "type with leading period still resolves" do
-      assert {:ok, descriptor} =
-               Service.get_by_symbol(".helloworld.HelloRequest")
+      assert {:ok, filename} =
+               Service.get_filename_by_symbol(".helloworld.HelloRequest")
 
-      assert {:ok, ^descriptor} =
-               Service.get_by_filename("helloworld.HelloRequest.proto")
+      assert {:ok, %Google.Protobuf.FileDescriptorProto{name: ^filename}} =
+               Service.get_by_filename(filename)
     end
   end
 end