Ensure runtime server exits after shutdown

lib/ruby_lsp/ruby_lsp_rails/addon.rb

@@ -39,7 +39,7 @@ def initialize
         @client_mutex = Mutex.new #: Mutex
         @client_mutex.lock
 
-        Thread.new do
+        @boot_thread = Thread.new do
           @addon_mutex.synchronize do
             # We need to ensure the Rails client is fully loaded before we activate the server addons
             @client_mutex.synchronize do
@@ -50,7 +50,7 @@ def initialize
             end
             offer_to_run_pending_migrations
           end
-        end
+        end #: Thread
       end
 
       #: -> RunnerClient
@@ -77,6 +77,7 @@ def activate(global_state, outgoing_queue)
       # @override
       #: -> void
       def deactivate
+        @boot_thread.join
         @rails_runner_client.shutdown
       end
 

lib/ruby_lsp/ruby_lsp_rails/runner_client.rb

@@ -51,18 +51,6 @@ class EmptyMessageError < MessageError; end
       def initialize(outgoing_queue, global_state)
         @outgoing_queue = outgoing_queue #: Thread::Queue
         @mutex = Mutex.new #: Mutex
-        # Spring needs a Process session ID. It uses this ID to "attach" itself to the parent process, so that when the
-        # parent ends, the spring process ends as well. If this is not set, Spring will throw an error while trying to
-        # set its own session ID
-        begin
-          Process.setpgrp
-          Process.setsid
-        rescue Errno::EPERM
-          # If we can't set the session ID, continue
-        rescue NotImplementedError
-          # setpgrp() may be unimplemented on some platform
-          # https://github.com/Shopify/ruby-lsp-rails/issues/348
-        end
 
         log_message("Ruby LSP Rails booting server")
 
@@ -96,27 +84,14 @@ def initialize(outgoing_queue, global_state)
         @rails_root = initialize_response[:root] #: String
         log_message("Finished booting Ruby LSP Rails server")
 
-        unless ENV["RAILS_ENV"] == "test"
-          at_exit do
-            if @wait_thread.alive?
-              sleep(0.5) # give the server a bit of time if we already issued a shutdown notification
-              force_kill
-            end
-          end
-        end
-
         # Responsible for transmitting notifications coming from the server to the outgoing queue, so that we can do
-        # things such as showing progress notifications initiated by the server
+        # things such as showing progress notifications initiated by the server. The loop exits naturally when the
+        # server closes its stderr write end (i.e., when the server process exits), at which point `read_notification`
+        # returns nil.
         @notifier_thread = Thread.new do
-          until @stderr.closed?
-            notification = read_notification
-
-            unless @outgoing_queue.closed? || !notification
-              @outgoing_queue << notification
-            end
+          while (notification = read_notification)
+            @outgoing_queue << notification unless @outgoing_queue.closed?
           end
-        rescue IOError
-          # The server was shutdown and stderr is already closed
         end #: Thread
       rescue StandardError
         raise InitializationError, @stderr.read
@@ -265,18 +240,25 @@ def trigger_i18n_reload
 
       #: -> void
       def shutdown
+        return if stopped?
+
         log_message("Ruby LSP Rails shutting down server")
         send_message("shutdown")
-        sleep(0.5) # give the server a bit of time to shutdown
-        [@stdin, @stdout, @stderr].each(&:close)
-      rescue IOError
-        # The server connection may have died
-        force_kill
+
+        @stdin.close unless @stdin.closed?
+
+        # Wait for the server to exit. Once it does, all handles it inherited (including its stderr write end) are
+        # released, which lets the notifier thread drain remaining bytes and observe EOF.
+        @wait_thread.join
+        @notifier_thread.join
+
+        @stdout.close unless @stdout.closed?
+        @stderr.close unless @stderr.closed?
       end
 
       #: -> bool
       def stopped?
-        [@stdin, @stdout, @stderr].all?(&:closed?) && !@wait_thread.alive?
+        [@stdin, @stdout, @stderr].all?(&:closed?) && !@wait_thread.alive? && !@notifier_thread.alive?
       end
 
       #: -> bool
@@ -337,15 +319,6 @@ def read_response
         nil
       end
 
-      #: -> void
-      def force_kill
-        # Windows does not support the `TERM` signal, so we're forced to use `KILL` here
-        Process.kill(
-          Signal.list["KILL"], #: as !nil
-          @wait_thread.pid,
-        )
-      end
-
       #: (::String message, ?type: ::Integer) -> void
       def log_message(message, type: RubyLsp::Constant::MessageType::LOG)
         return if @outgoing_queue.closed?

lib/ruby_lsp/ruby_lsp_rails/server.rb

@@ -289,8 +289,14 @@ def start
         send_result({ message: "ok", root: ::Rails.root.to_s })
 
         while @running
-          headers = @stdin.gets("\r\n\r\n") #: as String
-          json = @stdin.read(headers[/Content-Length: (\d+)/i, 1].to_i) #: as String
+          headers = @stdin.gets("\r\n\r\n")
+          break unless headers
+
+          length = headers[/Content-Length: (\d+)/i, 1]
+          break unless length
+
+          json = @stdin.read(length.to_i)
+          break unless json
 
           request = JSON.parse(json, symbolize_names: true)
           execute(request.fetch(:method), request[:params])
@@ -562,4 +568,8 @@ def resolve_i18n_key(key)
 
 if ARGV.first == "start"
   RubyLsp::Rails::Server.new(capabilities: JSON.parse(ARGV[1], symbolize_names: true)).start
+
+  # Ensure that we exit the process after finishing the server loop. This prevents child processes that may have been
+  # spawned by the user's Rails application from keeping this process alive
+  exit!(0)
 end

test/ruby_lsp_rails/runner_client_test.rb

@@ -22,10 +22,7 @@ class RunnerClientTest < ActiveSupport::TestCase
 
       teardown do
         @client.shutdown
-
-        # On Windows, the server process sometimes takes a lot longer to shutdown and may end up getting force killed,
-        # which makes this assertion flaky
-        assert_predicate(@client, :stopped?) unless Gem.win_platform?
+        assert_predicate(@client, :stopped?)
         @outgoing_queue.close
       end
 
@@ -145,6 +142,7 @@ class RunnerClientTest < ActiveSupport::TestCase
         begin
           assert(response.key?(:columns))
         ensure
+          client.shutdown
           outgoing_queue.close
           FileUtils.mv("test/dummy/config/application.rb.bak", "test/dummy/config/application.rb")
         end