saboteur 0.3.0 → 0.5.0

data/lib/saboteur/execution/supervisor.rb

@@ -40,11 +40,33 @@ module Saboteur
 
       # payload: bytes the child wrote to its pipe. stop: nil when the child
       # finished on its own, else :timeout, :memory or :processes.
-      Result = Data.define(:payload, :exit_status, :stopped_by, :detail, :runtime, :leftover_pids)
+      # supervisor_killed: the child had already closed its pipe (delivered
+      # its result) but was still running after POST_EOF_GRACE, so the
+      # supervisor killed it. Its signal exit is then ours, not a failure of
+      # the child. memory_pressure: the process tree was seen growing to at
+      # least half of the memory limit, the evidence that a SIGKILL may be
+      # an out-of-memory kill.
+      Result = Data.define(:payload, :exit_status, :stopped_by, :detail, :runtime, :leftover_pids,
+                           :supervisor_killed, :memory_pressure) do
+        def initialize(supervisor_killed: false, memory_pressure: false, **fields)
+          super
+        end
+
+        # True when the child itself ended unsuccessfully (non-zero exit or
+        # a signal the supervisor did not send).
+        def child_failed?
+          exit_status.is_a?(Process::Status) && !exit_status.success? && !supervisor_killed
+        end
+      end
 
       POLL_INTERVAL = 0.5
       PR_SET_CHILD_SUBREAPER = 36
       TERM_GRACE = 1.0
+      # How long a child that closed its pipe may take to exit by itself
+      # (at_exit handlers, buffered output) before the supervisor kills it.
+      POST_EOF_GRACE = 5.0
+      # Share of the memory limit the tree must reach to count as pressure.
+      MEMORY_PRESSURE_SHARE = 0.5
       REAP_DEADLINE = 5.0
 
       @live = {}
@@ -94,7 +116,10 @@ module Saboteur
       # Runs the block in a forked child with a pipe writer; the child exits
       # when the block returns. Returns a Result. cpu_seconds bounds every
       # descendant's CPU time (default: twice the deadline, at least 10 s).
-      def run(timeout:, cpu_seconds: nil, &block)
+      # on_tick is called from the parent's wait loop after the fork, about
+      # every poll, so callers can report progress without a thread alive
+      # at fork time.
+      def run(timeout:, cpu_seconds: nil, on_tick: nil, &block)
         depth = ENV.fetch(CHILD_ENV, "0").to_i
         raise Error, "refusing to fork: saboteur children nested #{depth} deep (limit #{MAX_DEPTH})" if
           depth >= MAX_DEPTH
@@ -104,10 +129,11 @@ module Saboteur
         $stdout.flush
         $stderr.flush
         started = monotonic
+        @peak_growth_kb = 0
         supervisor = Process.pid
         pid = fork { child(reader, writer, supervisor, cpu_seconds, depth + 1, &block) }
         writer.close
-        contain(pid, reader, started, timeout)
+        contain(pid, reader, started, timeout, on_tick)
       ensure
         reader&.close unless reader&.closed?
       end
@@ -168,24 +194,99 @@ module Saboteur
 
       # --- parent ---------------------------------------------------------
 
-      def contain(pid, reader, started, timeout)
+      def contain(pid, reader, started, timeout, on_tick)
         tree = ProcessTree.new(pid)
         self.class.live[pid] = tree
         own_group(pid)
-        payload, stop, detail = watch(tree, reader, started + timeout)
-        exit_status = stop == :timeout ? terminate(pid) : nil
+        deadline = started + timeout
+        exit_status = nil
+        grace_expired = false
+        begin
+          payload, stop, detail = watch(tree, reader, deadline, on_tick)
+          # EOF: the child has reported. Keep supervising while it finishes
+          # exiting by itself (same deadline, limits and progress ticks).
+          exit_status, stop, detail, grace_expired = await_exit(tree, pid, deadline, on_tick) if stop.nil?
+        ensure
+          # Containment holds on every path out of watch: a raising on_tick
+          # (progress IO failed) or any other error must not skip the
+          # kill/reap, or the child and its tree would run unsupervised
+          # until this process exits.
+          exit_status, supervisor_killed, leftover_pids = release(pid, tree, exit_status, stop, grace_expired)
+        end
+        Result.new(payload: payload, exit_status: exit_status, stopped_by: stop, detail: detail,
+                   runtime: monotonic - started, leftover_pids: leftover_pids,
+                   supervisor_killed: supervisor_killed, memory_pressure: memory_pressure?)
+      end
+
+      # Kills and reaps the child and everything it started. Returns
+      # [exit_status, supervisor_killed, leftover_pids].
+      def release(pid, tree, exit_status, stop, grace_expired)
+        exit_status ||= terminate(pid) if stop == :timeout
+        supervisor_killed = false
+        # The grace ran out with the child seemingly alive. What ends it is
+        # only ours to claim if we see it alive: stop it first.
+        exit_status, supervisor_killed = freeze_child(pid) if grace_expired
         leftover_pids = tree.kill_all - [pid]
         exit_status ||= reap(pid)
         self.class.live.delete(pid)
         reap_adopted(leftover_pids)
         reap_group(pid)
-        Result.new(payload: payload, exit_status: exit_status, stopped_by: stop, detail: detail,
-                   runtime: monotonic - started, leftover_pids: leftover_pids)
+        [exit_status, supervisor_killed, leftover_pids]
+      end
+
+      # After EOF, waits for the child to exit by itself while still checking
+      # the deadline, the memory and process limits and calling on_tick.
+      # Returns [exit_status, stop, detail, grace_expired]. exit_status is set
+      # only when the child exited on its own; grace_expired is true when it
+      # was still running after POST_EOF_GRACE.
+      def await_exit(tree, pid, deadline, on_tick)
+        grace_deadline = [monotonic + POST_EOF_GRACE, deadline].min
+        next_poll = monotonic + POLL_INTERVAL
+        loop do
+          status = waitpid(pid)
+          return [status, nil, nil, false] if status
+
+          now = monotonic
+          return [nil, :timeout, "exceeded deadline", false] if now >= deadline
+          return [nil, nil, nil, true] if now >= grace_deadline
+
+          on_tick&.call
+          if now >= next_poll
+            stop = over_limit(tree)
+            return [nil, *stop, false] if stop
+
+            next_poll = now + POLL_INTERVAL
+          end
+          sleep 0.01
+        end
+      end
+
+      # SIGSTOP cannot be caught, ignored or raced by a signal handler. A
+      # child that reports "stopped" was alive when we acted, so the SIGKILL
+      # that follows is ours. A child that turns out to have ended by itself
+      # keeps its own status. Returns [exit_status or nil, supervisor_killed].
+      def freeze_child(pid)
+        signal("STOP", pid)
+        deadline = monotonic + TERM_GRACE
+        until monotonic > deadline
+          _, status = Process.waitpid2(pid, Process::WNOHANG | Process::WUNTRACED)
+          return [nil, true] if status&.stopped?
+          return [status, false] if status
+
+          sleep 0.01
+        end
+        [nil, false]
+      rescue Errno::ECHILD
+        [:reaped, false]
+      end
+
+      def memory_pressure?
+        @peak_growth_kb >= @limits.max_memory_mb * 1024 * MEMORY_PRESSURE_SHARE
       end
 
       # Reads until EOF, polling the tree's size and memory. Returns
       # [payload, stop, detail].
-      def watch(tree, reader, deadline)
+      def watch(tree, reader, deadline, on_tick)
         payload = +""
         next_poll = monotonic + POLL_INTERVAL
         tree.live_members
@@ -193,6 +294,8 @@ module Saboteur
           now = monotonic
           return [payload, :timeout, "exceeded deadline"] if now >= deadline
 
+          on_tick&.call
+
           if now >= next_poll
             stop = over_limit(tree)
             return [payload, *stop] if stop
@@ -211,6 +314,7 @@ module Saboteur
 
       def over_limit(tree)
         count, growth_kb = tree.usage(baseline_rss_kb)
+        @peak_growth_kb = [@peak_growth_kb, growth_kb].max
         if growth_kb > @limits.max_memory_mb * 1024
           [:memory, "process tree grew by #{growth_kb / 1024} MB (limit #{@limits.max_memory_mb} MB)"]
         elsif count > @limits.max_processes

data/lib/saboteur/execution/worker_pool.rb

@@ -30,9 +30,10 @@ module Saboteur
 
       # on_start: called in each worker with its number before its first
       # item (per-worker database connections, for example).
-      def initialize(jobs:, on_start: nil)
+      def initialize(jobs:, on_start: nil, on_result: nil)
         @jobs = jobs
         @on_start = on_start
+        @on_result = on_result
         @start_failures = 0
       end
 
@@ -78,13 +79,22 @@ module Saboteur
       # next item, retires it when none is left, or replaces it if it died.
       def receive(worker, workers, queue, results, items, &)
         line = worker.results.gets
-        take(worker, line, results) if line
+        take_and_report(worker, line, results) if line
         return replace(worker, workers, queue, results, items, &) if line.nil? || worker.last_words
 
         assign(worker, queue)
         retire(worker, workers) unless worker.current
       end
 
+      # Records one result line, then reports the arrival before the worker
+      # gets more work: a watcher (the canary) can raise here and stop every
+      # queued item from starting (issue #102).
+      def take_and_report(worker, line, results)
+        current = worker.current
+        take(worker, line, results)
+        @on_result&.call(current, results[current]) if current && results[current]
+      end
+
       def parse(line)
         JSON.parse(line)
       rescue JSON::ParserError => e
@@ -112,7 +122,7 @@ module Saboteur
         pending = worker.results.read_nonblock(1 << 20, exception: false)
         return unless pending.is_a?(String)
 
-        pending.each_line { |line| take(worker, line, results) if line.end_with?("\n") }
+        pending.each_line { |line| take_and_report(worker, line, results) if line.end_with?("\n") }
       rescue IOError
         nil
       end
@@ -121,7 +131,13 @@ module Saboteur
         drain(worker, results)
         retire(worker, workers)
         reason = worker.last_words || "worker #{worker.number} died#{how(worker.status)}"
-        results[worker.current] = { "error" => "#{reason}, running this item" } if worker.current
+        if worker.current
+          # The synthetic error row takes the same arrival path as a normal
+          # result, so the canary still trips when the dead worker damaged
+          # the tree before dying (issue #102).
+          results[worker.current] = { "error" => "#{reason}, running this item" }
+          @on_result&.call(worker.current, results[worker.current])
+        end
         give_up(reason) if worker.done.zero?
         return if queue.empty?
 

data/lib/saboteur/execution/wrapper.rb

@@ -0,0 +1,45 @@
+# SPDX-License-Identifier: MIT
+# frozen_string_literal: true
+
+module Saboteur
+  module Execution
+    # Registry of gems that wrap a method at definition time. Their wrapper
+    # is what Ruby runs, so the shadow guard would call every wrapped method
+    # "shadowed" and skip it. A registered wrapper is instead recognized and
+    # re-evaluated together with the def (see Scheduler#eval_prefix). An
+    # unknown wrapper still hits the generic guard, which is the safe side.
+    module Wrapper
+      # :sorbet when the runtime method is sorbet-runtime's sig wrapper:
+      # its source lives inside the gem and a signature is registered for it.
+      def self.around(method, file)
+        return nil unless file&.include?("sorbet-runtime")
+        return nil unless defined?(::T::Utils) && ::T::Utils.respond_to?(:signature_for_method)
+        return nil unless ::T::Utils.signature_for_method(method)
+
+        :sorbet
+      rescue StandardError
+        nil
+      end
+
+      # The def a wrapper dispatches to: sorbet's runtime method is a
+      # forwarding wrapper, while the registered signature holds the real
+      # def. Callers comparing definitions (StaleCaptures) must compare
+      # against this, or the freshly re-registered def reads as stale.
+      # Real sorbet returns a Signature object (its #method is the def);
+      # minimal registries return the method itself. Nil when unknown.
+      def self.original_def(method)
+        return nil unless around(method, method.source_location&.first) == :sorbet
+
+        registered = ::T::Utils.signature_for_method(method)
+        candidate = if registered.is_a?(Method) || registered.is_a?(UnboundMethod)
+                      registered
+                    elsif registered.respond_to?(:method)
+                      registered.method
+                    end
+        candidate.is_a?(Method) || candidate.is_a?(UnboundMethod) ? candidate : nil
+      rescue StandardError
+        nil
+      end
+    end
+  end
+end