potato-mesh/web/lib/potato_mesh/application/worker_pool.rb

# Copyright © 2025-26 l5yth & contributors
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# frozen_string_literal: true

require "timeout"

module PotatoMesh
  module App
    # WorkerPool executes submitted blocks using a bounded set of Ruby threads.
    #
    # The pool enforces an upper bound on queued tasks, surfaces errors raised
    # by jobs, and supports graceful shutdown during application teardown.
    class WorkerPool
      # Raised when the worker pool queue has reached its configured capacity.
      class QueueFullError < StandardError; end

      # Raised when a task fails to complete before the requested timeout.
      class TaskTimeoutError < StandardError; end

      # Raised when scheduling occurs after the pool has been shut down.
      class ShutdownError < StandardError; end

      # Internal structure responsible for coordinating task completion.
      class Task
        # @return [Object, nil] value produced by the task block when available.
        attr_reader :value

        # @return [StandardError, nil] error raised by the task block when set.
        attr_reader :error

        def initialize
          @mutex = Mutex.new
          @condition = ConditionVariable.new
          @complete = false
          @value = nil
          @error = nil
        end

        # Mark the task as completed successfully.
        #
        # @param result [Object] value produced by the job.
        # @return [void]
        def fulfill(result)
          @mutex.synchronize do
            return if @complete

            @complete = true
            @value = result
            @condition.broadcast
          end
        end

        # Mark the task as failed with the provided error.
        #
        # @param failure [StandardError] exception raised while executing the job.
        # @return [void]
        def reject(failure)
          @mutex.synchronize do
            return if @complete

            @complete = true
            @error = failure
            @condition.broadcast
          end
        end

        # Wait for the task to complete, raising any stored failure.
        #
        # @param timeout [Numeric, nil] optional timeout in seconds.
        # @return [Object] the value produced by the job when successful.
        # @raise [TaskTimeoutError] when the timeout elapses prior to completion.
        # @raise [StandardError] when the job raised an exception.
        def wait(timeout: nil)
          deadline = timeout && monotonic_now + timeout

          @mutex.synchronize do
            until @complete
              if deadline
                remaining = deadline - monotonic_now
                raise TaskTimeoutError, "task deadline exceeded" if remaining <= 0

                @condition.wait(@mutex, remaining)
              else
                @condition.wait(@mutex)
              end
            end

            raise @error if @error

            @value
          end
        end

        # Check whether the task has finished executing.
        #
        # @return [Boolean] true when the task is complete.
        def complete?
          @mutex.synchronize { @complete }
        end

        private

        def monotonic_now
          Process.clock_gettime(Process::CLOCK_MONOTONIC)
        end
      end

      STOP_SIGNAL = Object.new

      # @return [Array<Thread>] threads created to service the pool.
      attr_reader :threads

      # Initialize a worker pool using the supplied configuration.
      #
      # @param size [Integer] number of worker threads to spawn.
      # @param max_queue [Integer, nil] optional upper bound on queued jobs.
      # @param task_timeout [Numeric, nil] optional per-task execution timeout.
      # @param name [String] prefix assigned to worker thread names.
      def initialize(size:, max_queue: nil, task_timeout: nil, name: "worker-pool")
        raise ArgumentError, "size must be positive" unless size.is_a?(Integer) && size.positive?

        @name = name
        @queue = max_queue ? SizedQueue.new(max_queue) : Queue.new
        @threads = []
        @stopped = false
        @mutex = Mutex.new
        @task_timeout = normalize_task_timeout(task_timeout)
        spawn_workers(size)
      end

      # Determine whether the worker pool is still accepting work.
      #
      # @return [Boolean] true when the pool remains active.
      def alive?
        @mutex.synchronize { !@stopped }
      end

      # Submit a block of work for asynchronous execution.
      #
      # @yieldreturn [Object] result produced by the job block.
      # @return [Task] task tracking the asynchronous execution.
      # @raise [QueueFullError] when the queue cannot accept additional work.
      # @raise [ShutdownError] when the pool is no longer active.
      def schedule(&block)
        raise ArgumentError, "block required" unless block

        task = Task.new

        @mutex.synchronize do
          raise ShutdownError, "worker pool has been shut down" if @stopped

          begin
            @queue.push([task, block], true)
          rescue ThreadError => e
            raise QueueFullError, e.message
          end
        end

        task
      end

      # Stop accepting work and wait for the worker threads to finish.
      #
      # Pending tasks that have not yet started executing are rejected with a
      # {ShutdownError}; workers that fail to honour STOP_SIGNAL within
      # +timeout+ are hard-killed if +force_kill_after+ is supplied.  This
      # bounds CTRL+C / process-exit time even when a worker is blocked deep
      # in a network call that ignores its task-level Timeout watcher.
      #
      # @param timeout [Numeric, nil] seconds to wait for each worker to exit
      #   cooperatively.
      # @param force_kill_after [Numeric, nil] additional seconds to wait
      #   after a hard {Thread#kill}; nil disables hard-kill entirely
      #   (legacy behavior).
      # @return [void]
      def shutdown(timeout: nil, force_kill_after: nil)
        threads = nil

        @mutex.synchronize do
          return if @stopped

          @stopped = true
          threads = @threads.dup
        end

        drain_pending_tasks!
        threads.each { @queue << STOP_SIGNAL }
        threads.each { |thread| thread.join(timeout) }
        return unless force_kill_after

        threads.each do |thread|
          next unless thread.alive?

          thread.kill
          thread.join(force_kill_after)
        end
      end

      private

      # Pop and reject every task currently queued so workers see STOP_SIGNAL
      # immediately instead of after draining a long backlog of in-flight
      # crawl/announce jobs.  Callers awaiting those tasks observe a
      # {ShutdownError} rather than silent timeouts.
      #
      # @return [void]
      def drain_pending_tasks!
        loop do
          entry = begin
              @queue.pop(true)
            rescue ThreadError
              nil
            end
          break unless entry

          task, _block = entry
          next unless task.respond_to?(:reject)

          task.reject(ShutdownError.new("worker pool shut down"))
        end
      end

      public

      private

      def spawn_workers(size)
        size.times do |index|
          worker = Thread.new do
            Thread.current.name = "#{@name}-#{index}" if Thread.current.respond_to?(:name=)
            Thread.current.report_on_exception = false if Thread.current.respond_to?(:report_on_exception=)
            # Daemon threads allow the process to exit even if a job is stuck.
            Thread.current.daemon = true if Thread.current.respond_to?(:daemon=)

            loop do
              task, block = @queue.pop
              break if task.equal?(STOP_SIGNAL)

              begin
                result = if @task_timeout
                    Timeout.timeout(@task_timeout, TaskTimeoutError, "task exceeded timeout") do
                      block.call
                    end
                  else
                    block.call
                  end
                task.fulfill(result)
              rescue StandardError => e
                task.reject(e)
              end
            end
          end
          @threads << worker
        end
      end

      # Normalize the per-task timeout into a positive float value.
      #
      # @param task_timeout [Numeric, nil] candidate timeout value.
      # @return [Float, nil] positive timeout in seconds or nil when disabled.
      def normalize_task_timeout(task_timeout)
        return nil if task_timeout.nil?

        value = Float(task_timeout)
        return nil unless value.positive?

        value
      rescue ArgumentError, TypeError
        nil
      end
    end
  end
end