Merge pull request #801 from soutaro/exit

soutaro · web-flow · commit 8b8ed454f846 · 2023-05-24T22:09:48.000+09:00
Improve worker process handling
diff --git a/lib/steep.rb b/lib/steep.rb
@@ -26,6 +26,7 @@
 require "steep/shims/filter_map"
 require "steep/shims/symbol_start_with"
 
+require "steep/thread_waiter"
 require "steep/equatable"
 require "steep/method_name"
 require "steep/node_helper"
diff --git a/lib/steep/drivers/check.rb b/lib/steep/drivers/check.rb
@@ -55,9 +55,9 @@ def run
         master.commandline_args.push(*command_line_patterns)
 
         main_thread = Thread.start do
+          Thread.current.abort_on_exception = true
           master.start()
         end
-        main_thread.abort_on_exception = true
 
         Steep.logger.info { "Initializing server" }
         initialize_id = request_id()
@@ -130,6 +130,9 @@ def run
           stdout.puts Rainbow("Unexpected error reported. 🚨").red.bold
           1
         end
+      rescue Errno::EPIPE => error
+        stdout.puts Rainbow("Steep shutdown with an error: #{error.inspect}").red.bold
+        return 1
       end
 
       def print_expectations(project:, all_files:, expectations_path:, notifications:)
diff --git a/lib/steep/drivers/checkfile.rb b/lib/steep/drivers/checkfile.rb
@@ -119,9 +119,9 @@ def run
         master.typecheck_automatically = false
 
         main_thread = Thread.start do
+          Thread.current.abort_on_exception = true
           master.start()
         end
-        main_thread.abort_on_exception = true
 
         Steep.logger.info { "Initializing server" }
         initialize_id = request_id()
diff --git a/lib/steep/drivers/stats.rb b/lib/steep/drivers/stats.rb
@@ -145,9 +145,9 @@ def run
         master.commandline_args.push(*command_line_patterns)
 
         main_thread = Thread.start do
+          Thread.current.abort_on_exception = true
           master.start()
         end
-        main_thread.abort_on_exception = true
 
         initialize_id = request_id()
         client_writer.write({ method: :initialize, id: initialize_id })
diff --git a/lib/steep/drivers/utils/driver_helper.rb b/lib/steep/drivers/utils/driver_helper.rb
@@ -10,6 +10,16 @@ def load_config(path: steepfile || Pathname("Steepfile"))
           steep_file_path = path.absolute? ? path : Pathname.pwd + path
           Project.new(steepfile_path: steep_file_path).tap do |project|
             Project::DSL.parse(project, path.read, filename: path.to_s)
+
+            project.targets.each do |target|
+              if collection_lock = target.options.collection_lock
+                begin
+                  collection_lock.check_rbs_availability!
+                rescue RBS::Collection::Config::CollectionNotAvailable
+                  raise "Run `rbs collection install` to install type definitions"
+                end
+              end
+            end
           end
         end
 
diff --git a/lib/steep/drivers/watch.rb b/lib/steep/drivers/watch.rb
@@ -57,7 +57,6 @@ def run()
         main_thread = Thread.start do
           master.start()
         end
-        main_thread.abort_on_exception = true
 
         initialize_id = request_id()
         client_writer.write(method: "initialize", id: initialize_id)
@@ -157,7 +156,13 @@ def run()
           main_thread.join
         rescue Interrupt
           master.kill
-          main_thread.join
+          begin
+            main_thread.join
+          rescue
+            master.each_worker do |worker|
+              worker.kill(force: true)
+            end
+          end
         end
 
         0
diff --git a/lib/steep/server/base_worker.rb b/lib/steep/server/base_worker.rb
@@ -40,6 +40,8 @@ def handle_job(job)
       def run
         tags = Steep.logger.formatter.current_tags.dup
         thread = Thread.new do
+          Thread.current.abort_on_exception = true
+
           Steep.logger.formatter.push_tags(*tags)
           Steep.logger.tagged "background" do
             while job = queue.pop
@@ -69,7 +71,6 @@ def run
             end
           end
         end
-        thread.abort_on_exception = true
 
         Steep.logger.tagged "frontend" do
           begin
diff --git a/lib/steep/server/master.rb b/lib/steep/server/master.rb
@@ -450,50 +450,65 @@ def start
             end
           end
 
-          Steep.logger.tagged "main" do
-            while job = job_queue.deq
-              case job
-              when ReceiveMessageJob
-                src = case job.source
+          loop_thread = Thread.new do
+            Steep.logger.formatter.push_tags(*tags)
+            Steep.logger.tagged "main" do
+              while job = job_queue.deq
+                case job
+                when ReceiveMessageJob
+                  src = case job.source
+                        when :client
+                          :client
+                        else
+                          job.source.name
+                        end
+                  Steep.logger.tagged("ReceiveMessageJob(#{src}/#{job.message[:method]}/#{job.message[:id]})") do
+                    if job.response? && result_controller.process_response(job.message)
+                      # nop
+                      Steep.logger.info { "Processed by ResultController" }
+                    else
+                      case job.source
                       when :client
-                        :client
-                      else
-                        job.source.name
-                      end
-                Steep.logger.tagged("ReceiveMessageJob(#{src}/#{job.message[:method]}/#{job.message[:id]})") do
-                  if job.response? && result_controller.process_response(job.message)
-                    # nop
-                    Steep.logger.info { "Processed by ResultController" }
-                  else
-                    case job.source
-                    when :client
-                      process_message_from_client(job.message)
-
-                      if job.message[:method] == "exit"
-                        job_queue.close()
+                        process_message_from_client(job.message)
+
+                        if job.message[:method] == "exit"
+                          job_queue.close()
+                        end
+                      when WorkerProcess
+                        process_message_from_worker(job.message, worker: job.source)
                       end
-                    when WorkerProcess
-                      process_message_from_worker(job.message, worker: job.source)
                     end
                   end
-                end
-              when SendMessageJob
-                case job.dest
-                when :client
-                  Steep.logger.info { "Processing SendMessageJob: dest=client, method=#{job.message[:method] || "-"}, id=#{job.message[:id] || "-"}" }
-                  writer.write job.message
-                when WorkerProcess
-                  Steep.logger.info { "Processing SendMessageJob: dest=#{job.dest.name}, method=#{job.message[:method] || "-"}, id=#{job.message[:id] || "-"}" }
-                  job.dest << job.message
+                when SendMessageJob
+                  case job.dest
+                  when :client
+                    Steep.logger.info { "Processing SendMessageJob: dest=client, method=#{job.message[:method] || "-"}, id=#{job.message[:id] || "-"}" }
+                    writer.write job.message
+                  when WorkerProcess
+                    Steep.logger.info { "Processing SendMessageJob: dest=#{job.dest.name}, method=#{job.message[:method] || "-"}, id=#{job.message[:id] || "-"}" }
+                    job.dest << job.message
+                  end
                 end
               end
             end
+          end
+
+          waiter = ThreadWaiter.new(each_worker.to_a) {|worker| worker.wait_thread }
+          waiter.wait_one()
 
-            read_client_thread.join()
-            worker_threads.each do |thread|
-              thread.join
+          unless job_queue.closed?
+            # Exit by error
+            each_worker do |worker|
+              worker.kill(force: true)
             end
+            raise "Unexpected worker process exit"
+          end
+
+          read_client_thread.join()
+          worker_threads.each do |thread|
+            thread.join
           end
+          loop_thread.join
         end
       end
 
diff --git a/lib/steep/server/worker_process.rb b/lib/steep/server/worker_process.rb
@@ -63,6 +63,9 @@ def self.fork_worker(type, name:, steepfile:, index:, delay_shutdown:, patterns:
         worker.commandline_args = patterns
 
         pid = fork do
+          Process.setpgid(0, 0)
+          stdin_out.close
+          stdout_in.close
           worker.run()
         end
 
@@ -149,8 +152,22 @@ def read(&block)
         reader.read(&block)
       end
 
-      def kill
-        Process.kill(:TERM, @wait_thread.pid)
+      def kill(force: false)
+        Steep.logger.tagged("WorkerProcess#kill@#{name}(#{wait_thread.pid})") do
+          begin
+            signal = force ? :KILL : :TERM
+            Steep.logger.debug("Sending signal SIG#{signal}...")
+            Process.kill(signal, wait_thread.pid)
+            Steep.logger.debug("Successfully sent the signal.")
+          rescue Errno::ESRCH => error
+            Steep.logger.debug("Failed #{error.inspect}")
+          end
+          unless force
+            Steep.logger.debug("Waiting for process exit...")
+            wait_thread.join()
+            Steep.logger.debug("Confirmed process exit.")
+          end
+        end
       end
     end
   end
diff --git a/lib/steep/thread_waiter.rb b/lib/steep/thread_waiter.rb
@@ -0,0 +1,35 @@
+module Steep
+  class ThreadWaiter
+    attr_reader :objects, :queue, :waiter_threads
+
+    def initialize(objects)
+      @objects = objects
+      @queue = Thread::Queue.new()
+      @waiter_threads = Set[].compare_by_identity
+
+      objects.each do |object|
+        thread = yield(object)
+
+        waiter_thread = Thread.new(thread) do |thread|
+          Thread.current.report_on_exception = false
+
+          begin
+            thread.join
+          ensure
+            queue << [object, thread]
+          end
+        end
+
+        waiter_threads << waiter_thread
+      end
+    end
+
+    def wait_one
+      unless waiter_threads.empty?
+        obj, th = queue.pop()
+        waiter_threads.delete(th)
+        obj
+      end
+    end
+  end
+end
diff --git a/sig/steep/drivers/check.rbs b/sig/steep/drivers/check.rbs
@@ -21,7 +21,7 @@ module Steep
 
       def initialize: (stdout: untyped, stderr: untyped) -> void
 
-      def run: () -> untyped
+      def run: () -> Integer
 
       def print_expectations: (project: untyped, all_files: untyped, expectations_path: untyped, notifications: untyped) -> untyped
 
diff --git a/sig/steep/project/options.rbs b/sig/steep/project/options.rbs
@@ -19,7 +19,7 @@ module Steep
 
       attr_accessor paths: PathOptions
 
-      attr_accessor collection_lock: RBS::Collection::Config::Lockfile
+      attr_accessor collection_lock: RBS::Collection::Config::Lockfile?
 
       def initialize: () -> void
     end
diff --git a/sig/steep/server/master.rbs b/sig/steep/server/master.rbs
@@ -227,8 +227,24 @@ module Steep
 
       def initialize: (project: Project, reader: untyped, writer: untyped, interaction_worker: WorkerProcess?, typecheck_workers: Array[WorkerProcess], ?queue: Thread::Queue) -> void
 
+      # Start the Steep language server
+      #
+      # Returns after the language server shutdown.
+      # The shutdown process starts by `exit` message.
+      #
+      # Raises an exception if the worker processes doesn't shutdown properly, like being `#kill`-ed.
+      # Exiting the process is assumed in the case of error.
+      # The only one cleanup it does is stopping the worker processes, by `SIGKILL`.
+      #
       def start: () -> void
 
+      # Stops the language server workers by sending SIGTERM
+      #
+      # This skips the normal shutdown process defined in LSP.
+      # The `#start` call will result in an exception.
+      #
+      def kill: () -> void
+
       def each_worker: () { (WorkerProcess) -> void } -> void
                      | () -> Enumerator[WorkerProcess, void]
 
@@ -253,8 +269,6 @@ module Steep
       def send_request: (method: String, worker: WorkerProcess, ?id: String, ?params: untyped?) ?{ (ResultHandler) -> void } -> ResultHandler
 
       def group_request: () { (GroupHandler) -> void } -> GroupHandler
-
-      def kill: () -> void
     end
   end
 end
diff --git a/sig/steep/server/worker_process.rbs b/sig/steep/server/worker_process.rbs
@@ -87,7 +87,11 @@ module Steep
 
       def read: () { (untyped) -> void } -> void
 
-      def kill: () -> void
+      # Stop the process by sending SIGTERM and wait for the process exit
+      #
+      # Send `SIGKILL` and returns immediately if `force:` is `true`.
+      #
+      def kill: (?force: bool) -> void
     end
   end
 end
diff --git a/sig/steep/thread_waiter.rbs b/sig/steep/thread_waiter.rbs
@@ -0,0 +1,13 @@
+module Steep
+  class ThreadWaiter[A]
+    attr_reader objects: Array[A]
+
+    attr_reader queue: Thread::Queue
+
+    attr_reader waiter_threads: Set[Thread]
+
+    def initialize: (Array[A]) { (A) -> Thread } -> void
+
+    def wait_one: () -> Thread?
+  end
+end
diff --git a/test/master_test.rb b/test/master_test.rb
