Check if bundle is valid before restarting (#3066)

### Motivation

The more we restart the server, the higher the risk that something related to Bundler may fail and prevent users from having a smooth experience.

I think we can be a bit smarter about when we restart the server by checking if composing the bundle is successful before we try. If we know that composing fails, it's better to continue running the current version of the server than crashing.

### Implementation

The proposal is to add a custom request that allows the client to ask the server if composing the bundle succeeds. That way, if there's a lockfile update, we can first check if the update produces a valid bundle.

If we couldn't compose the bundle, then we avoid restarting. Otherwise, we trigger the restart, but skip composing the bundle since we just did that ahead of time - saving time for the user and allow them to skip directly to indexing.

### Note

This solution doesn't fully solve #1458, but it does move the needle a little bit at least for VS Code. The big challenge in moving lockfile watching to the server is how to transfer the state of the server currently running to the new instance we launch.

For example, when we receive a notification that the lockfile was changed, we would need to:

1. Compose the new bundle
2. Replace the current process with a new one, using the new bundle (via `exec`)
3. Transfer the state of the previous process to the new one. We need all configurations that were applied during initialize (think the global state), the state of all current documents (the store object) and other pieces of state that exist in the version of the server running before the relaunch

I have not yet figured out how to achieve 3.

### Automated Tests

Added tests.

Author: vinistock

exe/ruby-lsp-launcher

@@ -8,14 +8,24 @@
 
 setup_error = nil
 install_error = nil
+reboot = false
 
-# Read the initialize request before even starting the server. We need to do this to figure out the workspace URI.
-# Editors are not required to spawn the language server process on the same directory as the workspace URI, so we need
-# to ensure that we're setting up the bundle in the right place
-$stdin.binmode
-headers = $stdin.gets("\r\n\r\n")
-content_length = headers[/Content-Length: (\d+)/i, 1].to_i
-raw_initialize = $stdin.read(content_length)
+workspace_uri = ARGV.first
+
+raw_initialize = if workspace_uri && !workspace_uri.start_with?("--")
+  # If there's an argument without `--`, then it's the server asking to compose the bundle and passing to this
+  # executable the workspace URI. We can't require gems at this point, so we built a fake initialize request manually
+  reboot = true
+  "{\"params\":{\"workspaceFolders\":[{\"uri\":\"#{workspace_uri}\"}]}}"
+else
+  # Read the initialize request before even starting the server. We need to do this to figure out the workspace URI.
+  # Editors are not required to spawn the language server process on the same directory as the workspace URI, so we need
+  # to ensure that we're setting up the bundle in the right place
+  $stdin.binmode
+  headers = $stdin.gets("\r\n\r\n")
+  content_length = headers[/Content-Length: (\d+)/i, 1].to_i
+  $stdin.read(content_length)
+end
 
 # Compose the Ruby LSP bundle in a forked process so that we can require gems without polluting the main process
 # `$LOAD_PATH` and `Gem.loaded_specs`. Windows doesn't support forking, so we need a separate path to support it
@@ -91,6 +101,13 @@ rescue StandardError => e
   $LOAD_PATH.unshift(File.expand_path("../lib", __dir__))
 end
 
+# When performing a lockfile re-boot, this executable is invoked to set up the composed bundle ahead of time. In this
+# flow, we are not booting the LSP yet, just checking if the bundle is valid before rebooting
+if reboot
+  # Use the exit status to signal to the server if composing the bundle succeeded
+  exit(install_error || setup_error ? 1 : 0)
+end
+
 # Now that the bundle is set up, we can begin actually launching the server. Note that `Bundler.setup` will have already
 # configured the load path using the version of the Ruby LSP present in the composed bundle. Do not push any Ruby LSP
 # paths into the load path manually or we may end up requiring the wrong version of the gem

lib/ruby_lsp/server.rb

@@ -106,6 +106,8 @@ def process_message(message)
               end,
           ),
         )
+      when "rubyLsp/composeBundle"
+        compose_bundle(message)
       when "$/cancelRequest"
         @global_state.synchronize { @cancelled_requests << message[:params][:id] }
       when nil
@@ -283,6 +285,7 @@ def run_initialize(message)
           document_range_formatting_provider: true,
           experimental: {
             addon_detection: true,
+            compose_bundle: true,
           },
         ),
         serverInfo: {
@@ -1282,5 +1285,30 @@ def window_show_message_request(message)
 
       addon.handle_window_show_message_response(result[:title])
     end
+
+    sig { params(message: T::Hash[Symbol, T.untyped]).void }
+    def compose_bundle(message)
+      already_composed_path = File.join(@global_state.workspace_path, ".ruby-lsp", "bundle_is_composed")
+      command = "#{Gem.ruby} #{File.expand_path("../../exe/ruby-lsp-launcher", __dir__)} #{@global_state.workspace_uri}"
+      id = message[:id]
+
+      # We compose the bundle in a thread so that the LSP continues to work while we're checking for its validity. Once
+      # we return the response back to the editor, then the restart is triggered
+      Thread.new do
+        send_log_message("Recomposing the bundle ahead of restart")
+        pid = Process.spawn(command)
+        _, status = Process.wait2(pid)
+
+        if status&.exitstatus == 0
+          # Create a signal for the restart that it can skip composing the bundle and launch directly
+          FileUtils.touch(already_composed_path)
+          send_message(Result.new(id: id, response: { success: true }))
+        else
+          # This special error code makes the extension avoid restarting in case we already know that the composed
+          # bundle is not valid
+          send_message(Error.new(id: id, code: BUNDLE_COMPOSE_FAILED_CODE, message: "Failed to compose bundle"))
+        end
+      end
+    end
   end
 end

lib/ruby_lsp/setup_bundler.rb

@@ -57,6 +57,7 @@ def initialize(project_path, **options)
       @lockfile_hash_path = T.let(@custom_dir + "main_lockfile_hash", Pathname)
       @last_updated_path = T.let(@custom_dir + "last_updated", Pathname)
       @error_path = T.let(@custom_dir + "install_error", Pathname)
+      @already_composed_path = T.let(@custom_dir + "bundle_is_composed", Pathname)
 
       dependencies, bundler_version = load_dependencies
       @dependencies = T.let(dependencies, T::Hash[String, T.untyped])
@@ -71,6 +72,23 @@ def initialize(project_path, **options)
     def setup!
       raise BundleNotLocked if !@launcher && @gemfile&.exist? && !@lockfile&.exist?
 
+      # If the bundle was composed ahead of time using our custom `rubyLsp/composeBundle` request, then we can skip the
+      # entire process and just return the composed environment
+      if @already_composed_path.exist?
+        $stderr.puts("Ruby LSP> Composed bundle was set up ahead of time. Skipping...")
+        @already_composed_path.delete
+
+        env = bundler_settings_as_env
+        env["BUNDLE_GEMFILE"] = @custom_gemfile.exist? ? @custom_gemfile.to_s : @gemfile.to_s
+
+        if env["BUNDLE_PATH"]
+          env["BUNDLE_PATH"] = File.expand_path(env["BUNDLE_PATH"], @project_path)
+        end
+
+        env["BUNDLER_VERSION"] = @bundler_version.to_s if @bundler_version
+        return env
+      end
+
       # Automatically create and ignore the .ruby-lsp folder for users
       @custom_dir.mkpath unless @custom_dir.exist?
       ignore_file = @custom_dir + ".gitignore"

lib/ruby_lsp/utils.rb

@@ -37,6 +37,8 @@ class DelegateRequestError < StandardError
     CODE = -32900
   end
 
+  BUNDLE_COMPOSE_FAILED_CODE = -33000
+
   # A notification to be sent to the client
   class Message
     extend T::Sig

project-words

@@ -26,6 +26,7 @@ dont
 eglot
 Eglot
 eruby
+exitstatus
 EXTGLOB
 fakehome
 FIXEDENCODING

test/server_test.rb

@@ -1033,6 +1033,29 @@ def test_rubocop_config_changes_trigger_workspace_diagnostic_refresh
     assert_equal("workspace/diagnostic/refresh", request.method)
   end
 
+  def test_compose_bundle_creates_file_to_skip_next_compose
+    Dir.mktmpdir do |dir|
+      Dir.chdir(dir) do
+        @server.process_message({
+          id: 1,
+          method: "initialize",
+          params: {
+            initializationOptions: {},
+            capabilities: { general: { positionEncodings: ["utf-8"] } },
+            workspaceFolders: [{ uri: URI::Generic.from_path(path: dir).to_s }],
+          },
+        })
+
+        capture_subprocess_io do
+          @server.process_message({ id: 2, method: "rubyLsp/composeBundle" })
+        end
+        result = find_message(RubyLsp::Result, id: 2)
+        assert(result.response[:success])
+        assert_path_exists(File.join(dir, ".ruby-lsp", "bundle_is_composed"))
+      end
+    end
+  end
+
   private
 
   def with_uninstalled_rubocop(&block)

test/setup_bundler_test.rb

@@ -858,6 +858,24 @@ def test_update_does_not_fail_if_gems_are_uninstalled
     end
   end
 
+  def test_only_returns_environment_if_bundle_was_composed_ahead_of_time
+    Dir.mktmpdir do |dir|
+      Dir.chdir(dir) do
+        FileUtils.mkdir(".ruby-lsp")
+        FileUtils.touch(File.join(".ruby-lsp", "bundle_is_composed"))
+
+        require "bundler/cli/update"
+        require "bundler/cli/install"
+        Bundler::CLI::Update.expects(:new).never
+        Bundler::CLI::Install.expects(:new).never
+
+        assert_output("", "Ruby LSP> Composed bundle was set up ahead of time. Skipping...\n") do
+          refute_empty(RubyLsp::SetupBundler.new(dir, launcher: true).setup!)
+        end
+      end
+    end
+  end
+
   private
 
   def with_default_external_encoding(encoding, &block)

vscode/src/workspace.ts

@@ -205,6 +205,8 @@ export class Workspace implements WorkspaceInterface {
         return this.start();
       }
 
+      let canRestart = false;
+
       switch (this.lspClient.state) {
         // If the server is still starting, then it may not be ready to handle a shutdown request yet. Trying to send
         // one could lead to a hanging process. Instead we set a flag and only restart once the server finished booting
@@ -214,10 +216,18 @@ export class Workspace implements WorkspaceInterface {
           break;
         // If the server is running, we want to stop it, dispose of the client and start a new one
         case State.Running:
-          await this.stop();
-          await this.lspClient.dispose();
-          this.lspClient = undefined;
-          await this.start();
+          // If the server doesn't support checking the validity of the composed bundle or if composing the bundle was
+          // successful, then we can restart
+          canRestart =
+            !this.lspClient.initializeResult?.capabilities.experimental
+              .compose_bundle || (await this.composingBundleSucceeds());
+
+          if (canRestart) {
+            await this.stop();
+            await this.lspClient.dispose();
+            this.lspClient = undefined;
+            await this.start();
+          }
           break;
         // If the server is already stopped, then we need to dispose it and start a new one
         case State.Stopped:
@@ -441,4 +451,19 @@ export class Workspace implements WorkspaceInterface {
     hash.update(fileContents.toString());
     return hash.digest("hex");
   }
+
+  private async composingBundleSucceeds(): Promise<boolean> {
+    if (!this.lspClient) {
+      return false;
+    }
+
+    try {
+      const response: { success: boolean } = await this.lspClient.sendRequest(
+        "rubyLsp/composeBundle",
+      );
+      return response.success;
+    } catch (error: any) {
+      return false;
+    }
+  }
 }