fix: handle document lifecycle notifications synchronously (#78)

Author: doorgan

lib/gen_lsp.ex

@@ -165,6 +165,19 @@ defmodule GenLSP do
   """
   @callback handle_info(message :: any(), state) :: {:noreply, state} when state: GenLSP.LSP.t()
 
+  @default_sync_notifications [
+    GenLSP.Notifications.TextDocumentDidOpen,
+    GenLSP.Notifications.TextDocumentDidChange,
+    GenLSP.Notifications.TextDocumentDidClose,
+    GenLSP.Notifications.TextDocumentDidSave,
+    GenLSP.Notifications.TextDocumentWillSave,
+    GenLSP.Notifications.WorkspaceDidChangeWatchedFiles,
+    GenLSP.Notifications.WorkspaceDidChangeWorkspaceFolders,
+    GenLSP.Notifications.WorkspaceDidChangeConfiguration,
+    GenLSP.Notifications.Initialized,
+    GenLSP.Notifications.Exit
+  ]
+
   @options_schema NimbleOptions.new!(
                     buffer: [
                       type: {:or, [:pid, :atom]},
@@ -182,6 +195,11 @@ defmodule GenLSP do
                       type: :atom,
                       doc:
                         "Used for name registration as described in the \"Name registration\" section in the documentation for `GenServer`."
+                    ],
+                    sync_notifications: [
+                      type: {:list, :atom},
+                      doc:
+                        "List of notification to process synchronously. Defaults to document lifecycle notifications."
                     ]
                   )
 
@@ -196,7 +214,8 @@ defmodule GenLSP do
     opts = NimbleOptions.validate!(opts, @options_schema)
 
     :proc_lib.start_link(__MODULE__, :init, [
-      {module, init_args, Keyword.take(opts, [:name, :buffer, :assigns, :task_supervisor]),
+      {module, init_args,
+       Keyword.take(opts, [:name, :buffer, :assigns, :task_supervisor, :sync_notifications]),
        self()}
     ])
   end
@@ -207,14 +226,16 @@ defmodule GenLSP do
     buffer = opts[:buffer]
     assigns = opts[:assigns]
     task_supervisor = opts[:task_supervisor]
+    sync_notifications = opts[:sync_notifications] || @default_sync_notifications
 
     lsp = %LSP{
       mod: module,
       pid: me,
       buffer: buffer,
       assigns: assigns,
       task_supervisor: task_supervisor,
-      tasks: Map.new()
+      tasks: Map.new(),
+      sync_notifications: MapSet.new(sync_notifications)
     }
 
     case module.init(lsp, init_args) do
@@ -472,76 +493,22 @@ defmodule GenLSP do
         start = System.system_time(:microsecond)
         :telemetry.execute([:gen_lsp, :notification, :client, :start], %{})
 
-        attempt(
-          lsp,
-          "Last message received: handle_notification #{inspect(notification)}",
-          [:gen_lsp, :notification, :client],
-          fn
-            {:error, _} ->
-              Logger.warning("client -> server notification crashed")
-
-            _ ->
-              case GenLSP.Notifications.new(notification) do
-                {:ok, %GenLSP.Notifications.DollarCancelRequest{} = note} ->
-                  result =
-                    :telemetry.span(
-                      [:gen_lsp, :handle_notification],
-                      %{method: note.method},
-                      fn ->
-                        with pid when is_pid(pid) <- lsp.tasks[note.params.id] do
-                          Task.Supervisor.terminate_child(lsp.task_supervisor, pid)
-                        end
-
-                        {{:noreply, lsp}, %{}}
-                      end
-                    )
-
-                  case result do
-                    {:noreply, %LSP{}} ->
-                      duration = System.system_time(:microsecond) - start
-
-                      Logger.debug(
-                        "handled notification client -> server #{note.method} in #{format_time(duration)}",
-                        method: note.method
-                      )
-
-                      :telemetry.execute([:gen_lsp, :notification, :client, :stop], %{
-                        duration: duration
-                      })
-                  end
-
-                {:ok, note} ->
-                  result =
-                    :telemetry.span(
-                      [:gen_lsp, :handle_notification],
-                      %{method: note.method},
-                      fn ->
-                        {lsp.mod.handle_notification(note, lsp), %{}}
-                      end
-                    )
-
-                  case result do
-                    {:noreply, %LSP{}} ->
-                      duration = System.system_time(:microsecond) - start
-
-                      Logger.debug(
-                        "handled notification client -> server #{note.method} in #{format_time(duration)}",
-                        method: note.method
-                      )
-
-                      :telemetry.execute([:gen_lsp, :notification, :client, :stop], %{
-                        duration: duration
-                      })
-                  end
+        case GenLSP.Notifications.new(notification) do
+          {:ok, %GenLSP.Notifications.DollarCancelRequest{} = note} ->
+            handle_cancel_request(lsp, note, start)
 
-                {:error, errors} ->
-                  # the payload is not parseable at all, other than being valid JSON
-                  exception = InvalidNotification.exception({notification, errors})
+          {:ok, note} ->
+            if MapSet.member?(lsp.sync_notifications, note.__struct__) do
+              handle_notification_sync(lsp, note, start)
+            else
+              handle_notification_async(lsp, note, start)
+            end
 
-                  Logger.warning(Exception.format(:error, exception))
-              end
-          end
-        )
+          {:error, errors} ->
+            # the payload is not parseable at all, other than being valid JSON
+            exception = InvalidNotification.exception({notification, errors})
+            Logger.warning(Exception.format(:error, exception))
+        end
 
         loop(lsp, parent, deb)
 
@@ -603,6 +570,89 @@ defmodule GenLSP do
     end)
   end
 
+  defp handle_cancel_request(lsp, note, start) do
+    result =
+      :telemetry.span([:gen_lsp, :handle_notification], %{method: note.method}, fn ->
+        with pid when is_pid(pid) <- lsp.tasks[note.params.id] do
+          Task.Supervisor.terminate_child(lsp.task_supervisor, pid)
+        end
+
+        {{:noreply, lsp}, %{}}
+      end)
+
+    case result do
+      {:noreply, %LSP{}} ->
+        duration = System.system_time(:microsecond) - start
+
+        Logger.debug(
+          "handled notification client -> server #{note.method} in #{format_time(duration)}",
+          method: note.method
+        )
+
+        :telemetry.execute([:gen_lsp, :notification, :client, :stop], %{duration: duration})
+    end
+  end
+
+  defp handle_notification_sync(lsp, note, start) do
+    try do
+      result =
+        :telemetry.span([:gen_lsp, :handle_notification], %{method: note.method}, fn ->
+          {lsp.mod.handle_notification(note, lsp), %{}}
+        end)
+
+      case result do
+        {:noreply, %LSP{}} ->
+          duration = System.system_time(:microsecond) - start
+
+          Logger.debug(
+            "handled notification client -> server #{note.method} in #{format_time(duration)}",
+            method: note.method
+          )
+
+          :telemetry.execute([:gen_lsp, :notification, :client, :stop], %{duration: duration})
+      end
+    rescue
+      e ->
+        :telemetry.execute([:gen_lsp, :notification, :client, :exception], %{
+          message: "Last message received: handle_notification #{inspect(note)}"
+        })
+
+        message = Exception.format(:error, e, __STACKTRACE__)
+        Logger.error(message)
+        error(lsp, message)
+    end
+  end
+
+  defp handle_notification_async(lsp, note, start) do
+    attempt(
+      lsp,
+      "Last message received: handle_notification #{inspect(note)}",
+      [:gen_lsp, :notification, :client],
+      fn
+        {:error, _} ->
+          :ok
+
+        _ ->
+          result =
+            :telemetry.span([:gen_lsp, :handle_notification], %{method: note.method}, fn ->
+              {lsp.mod.handle_notification(note, lsp), %{}}
+            end)
+
+          case result do
+            {:noreply, %LSP{}} ->
+              duration = System.system_time(:microsecond) - start
+
+              Logger.debug(
+                "handled notification client -> server #{note.method} in #{format_time(duration)}",
+                method: note.method
+              )
+
+              :telemetry.execute([:gen_lsp, :notification, :client, :stop], %{duration: duration})
+          end
+      end
+    )
+  end
+
   defp dump!(schematic, structure) do
     {:ok, output} = Schematic.dump(schematic, structure)
     output

lib/gen_lsp/lsp.ex

@@ -11,6 +11,7 @@ defmodule GenLSP.LSP do
     field :pid, pid()
     field :tasks, %{integer() => pid()}
     field :task_supervisor, atom() | pid()
+    field :sync_notifications, MapSet.t(module())
   end
 
   @spec assign(t(), Keyword.t() | (map() -> keyword())) :: t()

test/gen_lsp_test.exs

@@ -6,11 +6,15 @@ defmodule GenLSPTest do
   import GenLSP.Test
   import ExUnit.CaptureLog
 
-  setup do
-    server = server(GenLSPTest.ExampleServer, test_pid: self())
-    client = client(server)
-
-    [server: server, client: client]
+  setup context do
+    if context[:skip_setup] do
+      :ok
+    else
+      server = server(GenLSPTest.ExampleServer, test_pid: self())
+      client = client(server)
+
+      [server: server, client: client]
+    end
   end
 
   test "stores the user state and internal state", %{server: server} do
@@ -396,4 +400,57 @@ defmodule GenLSPTest do
              Process.sleep(100)
            end) =~ "Invalid notification from the client"
   end
+
+  @tag :skip_setup
+  test "processes sync notifications in order" do
+    {:ok, order_agent} = Agent.start_link(fn -> [] end)
+
+    server =
+      server(GenLSPTest.OrderingServer,
+        test_pid: self(),
+        order_agent: order_agent
+      )
+
+    client = client(server)
+
+    notify(client, %{
+      "jsonrpc" => "2.0",
+      "method" => "textDocument/didOpen",
+      "params" => %{
+        "textDocument" => %{
+          "uri" => "file://somefile",
+          "languageId" => "elixir",
+          "version" => 1,
+          "text" => "hello world!"
+        }
+      }
+    })
+
+    notify(client, %{
+      "jsonrpc" => "2.0",
+      "method" => "textDocument/didChange",
+      "params" => %{
+        "textDocument" => %{
+          "uri" => "file://somefile",
+          "version" => 2
+        },
+        "contentChanges" => [%{"text" => "updated content"}]
+      }
+    })
+
+    notify(client, %{
+      "jsonrpc" => "2.0",
+      "method" => "textDocument/didClose",
+      "params" => %{
+        "textDocument" => %{
+          "uri" => "file://somefile"
+        }
+      }
+    })
+
+    assert_receive {:callback, %Notifications.TextDocumentDidOpen{}}
+    assert_receive {:callback, %Notifications.TextDocumentDidChange{}}
+    assert_receive {:callback, %Notifications.TextDocumentDidClose{}}
+    assert Agent.get(order_agent, & &1) == [:did_open, :did_change, :did_close]
+  end
 end

test/support/ordering_server.ex

@@ -0,0 +1,47 @@
+defmodule GenLSPTest.OrderingServer do
+  use GenLSP
+  alias GenLSP.Notifications
+  alias GenLSP.Requests
+  alias GenLSP.Structures
+
+  def start_link(opts) do
+    {test_pid, opts} = Keyword.pop!(opts, :test_pid)
+    {order_agent, opts} = Keyword.pop!(opts, :order_agent)
+    GenLSP.start_link(__MODULE__, {test_pid, order_agent}, opts)
+  end
+
+  @impl true
+  def init(lsp, {test_pid, order_agent}) do
+    {:ok, assign(lsp, test_pid: test_pid, order_agent: order_agent)}
+  end
+
+  @impl true
+  def handle_request(%Requests.Initialize{}, lsp) do
+    {:reply,
+     %Structures.InitializeResult{
+       capabilities: %Structures.ServerCapabilities{},
+       server_info: %{name: "Ordering Test Server"}
+     }, lsp}
+  end
+
+  @impl true
+  def handle_notification(%Notifications.TextDocumentDidOpen{} = note, lsp) do
+    Agent.update(assigns(lsp).order_agent, fn list -> list ++ [:did_open] end)
+    send(assigns(lsp).test_pid, {:callback, note})
+    {:noreply, lsp}
+  end
+
+  @impl true
+  def handle_notification(%Notifications.TextDocumentDidChange{} = note, lsp) do
+    Agent.update(assigns(lsp).order_agent, fn list -> list ++ [:did_change] end)
+    send(assigns(lsp).test_pid, {:callback, note})
+    {:noreply, lsp}
+  end
+
+  @impl true
+  def handle_notification(%Notifications.TextDocumentDidClose{} = note, lsp) do
+    Agent.update(assigns(lsp).order_agent, fn list -> list ++ [:did_close] end)
+    send(assigns(lsp).test_pid, {:callback, note})
+    {:noreply, lsp}
+  end
+end