Check if bundle is valid before restarting

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;
+    }
+  }
 }