rage-rb 1.15.1 → 1.17.0

data/lib/rage/configuration.rb

@@ -35,6 +35,15 @@ require "erb"
 #
 # > Defines one or several old secrets that need to be rotated. Can accept a single key or an array of keys. Rage will fall back to the `FALLBACK_SECRET_KEY_BASE` environment variable if this is not set.
 #
+# • _config.after_initialize_
+#
+# > Schedule a block of code to run after Rage has finished loading the application code. Use this to reference application-level constants during the initialization process.
+# > ```
+# Rage.config.after_initialize do
+#   SUPER_USER = User.find_by!(super: true)
+# end
+# > ```
+#
 # # Middleware Configuration
 #
 # • _config.middleware.use_
@@ -115,7 +124,7 @@ require "erb"
 #
 # • _config.cable.protocol_
 #
-# > Specifies the protocol the server will use. The only value currently supported is `Rage::Cable::Protocol::ActioncableV1Json`. The client application will need to use [@rails/actioncable](https://www.npmjs.com/package/@rails/actioncable) to talk to the server.
+# > Specifies the protocol the server will use. Supported values include {Rage::Cable::Protocols::ActioncableV1Json :actioncable_v1_json} and {Rage::Cable::Protocols::RawWebSocketJson :raw_websocket_json}. Defaults to {Rage::Cable::Protocols::ActioncableV1Json :actioncable_v1_json}.
 #
 # • _config.cable.allowed_request_origins_
 #
@@ -136,6 +145,38 @@ require "erb"
 # end
 # > ```
 #
+# # Deferred Configuration
+# • _config.deferred.backend_
+#
+# > Specifies the backend for deferred tasks. Supported values are `:disk`, which uses disk storage, or `nil`, which disables persistence of deferred tasks.
+# > The `:disk` backend accepts the following options:
+# >
+# > - `:path` - the path to the directory where deferred tasks will be stored. Defaults to `storage`.
+# > - `:prefix` - the prefix for the deferred task files. Defaults to `deferred-`.
+# > - `:fsync_frequency` - the frequency of `fsync` calls in seconds. Defaults to `0.5`.
+#
+# > ```ruby
+# config.deferred.backend = :disk, { path: "storage" }
+# > ```
+#
+# • _config.deferred.backpressure_
+#
+# > Enables the backpressure for deferred tasks. The backpressure is used to limit the number of pending tasks in the queue. It accepts a hash with the following options:
+# >
+# > - `:high_water_mark` - the maximum number of pending tasks in the queue. Defaults to `1000`.
+# > - `:low_water_mark` - the minimum number of pending tasks in the queue before the backpressure is released. Defaults to `800`.
+# > - `:timeout` - the timeout for the backpressure in seconds. Defaults to `2`.
+#
+# > ```ruby
+# config.deferred.backpressure = { high_water_mark: 1000, low_water_mark: 800, timeout: 2 }
+# > ```
+#
+# > Additionally, you can set the backpressure value to `true` to use the default values:
+#
+# > ```ruby
+# config.deferred.backpressure = true
+# ```
+#
 # # Transient Settings
 #
 # The settings described in this section should be configured using **environment variables** and are either temporary or will become the default in the future.
@@ -153,6 +194,8 @@ require "erb"
 # > Instructs Rage to not reuse Active Record connections between different fibers.
 #
 class Rage::Configuration
+  include Hooks
+
   attr_accessor :logger
   attr_reader :log_formatter, :log_level
   attr_writer :secret_key_base, :fallback_secret_key_base
@@ -197,10 +240,22 @@ class Rage::Configuration
     @openapi ||= OpenAPI.new
   end
 
+  def deferred
+    @deferred ||= Deferred.new
+  end
+
   def internal
     @internal ||= Internal.new
   end
 
+  def after_initialize(&block)
+    push_hook(block, :after_initialize)
+  end
+
+  def run_after_initialize!
+    run_hooks_for!(:after_initialize, self)
+  end
+
   class Server
     attr_accessor :port, :workers_count, :timeout, :max_clients
     attr_reader :threads_count
@@ -257,15 +312,29 @@ class Rage::Configuration
   end
 
   class Cable
-    attr_accessor :protocol, :allowed_request_origins, :disable_request_forgery_protection
+    attr_accessor :allowed_request_origins, :disable_request_forgery_protection
+    attr_reader :protocol
 
     def initialize
-      @protocol = Rage::Cable::Protocol::ActioncableV1Json
+      @protocol = Rage::Cable::Protocols::ActioncableV1Json
       @allowed_request_origins = if Rage.env.development? || Rage.env.test?
         /localhost/
       end
     end
 
+    def protocol=(protocol)
+      @protocol = case protocol
+      when Class
+        protocol
+      when :actioncable_v1_json
+        Rage::Cable::Protocols::ActioncableV1Json
+      when :raw_websocket_json
+        Rage::Cable::Protocols::RawWebSocketJson
+      else
+        raise ArgumentError, "Unknown protocol. Supported values are `:actioncable_v1_json` and `:raw_websocket_json`."
+      end
+    end
+
     # @private
     def middlewares
       @middlewares ||= begin
@@ -314,6 +383,121 @@ class Rage::Configuration
     attr_accessor :tag_resolver
   end
 
+  class Deferred
+    attr_reader :backpressure
+
+    def initialize
+      @configured = false
+    end
+
+    def backend
+      unless @backend_class
+        @backend_class = Rage::Deferred::Backends::Disk
+        @backend_options = parse_disk_backend_options({})
+      end
+
+      @backend_class.new(**@backend_options)
+    end
+
+    def backend=(config)
+      @configured = true
+
+      backend_id, opts = if config.is_a?(Array)
+        [config[0], config[1]]
+      else
+        [config, {}]
+      end
+
+      @backend_class = case backend_id
+      when :disk
+        @backend_options = parse_disk_backend_options(opts)
+        Rage::Deferred::Backends::Disk
+      when nil
+        Rage::Deferred::Backends::Nil
+      else
+        raise ArgumentError, "unsupported backend value; supported keys are `:disk` and `nil`"
+      end
+    end
+
+    class Backpressure
+      attr_reader :high_water_mark, :low_water_mark, :timeout, :sleep_interval, :timeout_iterations
+
+      def initialize(high_water_mark = nil, low_water_mark = nil, timeout = nil)
+        @high_water_mark = high_water_mark || 1_000
+        @low_water_mark = low_water_mark || (@high_water_mark * 0.8).round
+
+        @timeout = timeout || 2
+        @sleep_interval = 0.05
+        @timeout_iterations = (@timeout / @sleep_interval).round
+      end
+    end
+
+    def backpressure=(config)
+      @configured = true
+
+      if config == true
+        @backpressure = Backpressure.new
+        return
+      elsif config == false
+        @backpressure = nil
+        return
+      end
+
+      if opts.except(:high_water_mark, :low_water_mark, :timeout).any?
+        raise ArgumentError, "unsupported backpressure options; supported keys are `:high_water_mark`, `:low_water_mark`, `:timeout`"
+      end
+
+      high_water_mark, low_water_mark, timeout = config.values_at(:high_water_mark, :low_water_mark, :timeout)
+      @backpressure = Backpressure.new(high_water_mark, low_water_mark, timeout)
+    end
+
+    def default_disk_storage_path
+      Pathname.new("storage")
+    end
+
+    def default_disk_storage_prefix
+      "deferred-"
+    end
+
+    def has_default_disk_storage?
+      default_disk_storage_path.glob("#{default_disk_storage_prefix}*").any?
+    end
+
+    def configured?
+      @configured
+    end
+
+    private
+
+    def parse_disk_backend_options(opts)
+      if opts.except(:path, :prefix, :fsync_frequency).any?
+        raise ArgumentError, "unsupported backend options; supported values are `:path`, `:prefix`, `:fsync_frequency`"
+      end
+
+      parsed_options = {}
+
+      parsed_options[:path] = if opts[:path]
+        opts[:path].is_a?(Pathname) ? opts[:path] : Pathname.new(opts[:path])
+      else
+        default_disk_storage_path
+      end
+
+      parsed_options[:prefix] = if opts[:prefix]
+        opts[:prefix].end_with?("-") ? opts[:prefix] : "#{opts[:prefix]}-"
+      else
+        default_disk_storage_prefix
+      end
+
+      parsed_options[:fsync_frequency] = if opts[:fsync_frequency]
+        (opts[:fsync_frequency].to_i * 1_000).round
+      else
+        500
+      end
+
+      parsed_options
+    end
+  end
+
   # @private
   class Internal
     attr_accessor :rails_mode

data/lib/rage/controller/api.rb

@@ -618,6 +618,7 @@ class RageController::API
   #  stale?(etag: "123", last_modified: Time.utc(2023, 12, 15))
   #  stale?(last_modified: Time.utc(2023, 12, 15))
   #  stale?(etag: "123")
+  # @note `stale?` will set ETag and Last-Modified response headers made of passed arguments in the method. Value for ETag will be additionally hashified using SHA1 algorithm, whereas value for Last-Modified will be converted to the string which represents time as RFC 1123 date of HTTP-date defined by RFC 2616.
   # @note `stale?` will set the response status to 304 if the request is fresh. This side effect will cause a double render error, if `render` gets called after this method. Make sure to implement a proper conditional in your action to prevent this from happening:
   #  ```ruby
   #  if stale?(etag: "123")
@@ -625,7 +626,10 @@ class RageController::API
   #  end
   #  ```
   def stale?(etag: nil, last_modified: nil)
-    still_fresh = request.fresh?(etag:, last_modified:)
+    response.etag = etag
+    response.last_modified = last_modified
+
+    still_fresh = request.fresh?(etag: response.etag, last_modified: last_modified)
 
     head :not_modified if still_fresh
     !still_fresh

data/lib/rage/deferred/backends/disk.rb

@@ -0,0 +1,258 @@
+# frozen_string_literal: true
+
+require "zlib"
+
+##
+# `Rage::Deferred::Backends` implements a storage layer to persist deferred tasks.
+# A storage should implement the following instance methods:
+#
+# * `add` - called when a task has to be added to the storage;
+# * `remove` - called when a task has to be removed from the storage;
+# * `pending_tasks` - the method should iterate over the underlying storage and return a list of tasks to replay;
+#
+class Rage::Deferred::Backends::Disk
+  STORAGE_VERSION = "0"
+  STORAGE_SIZE_INCREASE_RATIO = 1.5
+
+  DEFAULT_PUBLISH_AT = "0"
+  DEFAULT_STORAGE_SIZE_LIMIT = 2_000_000
+
+  def initialize(path:, prefix:, fsync_frequency:)
+    @storage_path = path
+    @storage_prefix = "#{prefix}#{STORAGE_VERSION}"
+    @fsync_frequency = fsync_frequency
+
+    @storage_path.mkpath
+
+    # try to open and take ownership of all storage files in the storage directory
+    storage_files = @storage_path.glob("#{@storage_prefix}-*").filter_map do |file_path|
+      file = file_path.open("a+b")
+      if file.flock(File::LOCK_EX | File::LOCK_NB)
+        sleep 0.01 # reduce contention between workers
+        file
+      else
+        file.close
+      end
+    end
+
+    # if there are no storage files - create one;
+    # otherwise the first one is used as the main storage; the rest will be merged into the main storage
+    if storage_files.empty?
+      @storage = create_storage
+    else
+      @storage = storage_files[0]
+      @recovered_storages = storage_files[1..] if storage_files.length > 1
+    end
+
+    # create seed value for the task IDs
+    task_id_seed = Time.now.to_i # TODO: ensure timestamps in the file are not higher
+    @task_id_base, @task_id_i = "#{task_id_seed}-#{Process.pid}", 0
+    Iodine.run_every(1_000) do
+      task_id_seed += 1
+      @task_id_base, @task_id_i = "#{task_id_seed}-#{Process.pid}", 0
+    end
+
+    @storage_size_limit = DEFAULT_STORAGE_SIZE_LIMIT
+    @storage_size = @storage.size
+    @fsync_scheduled = false
+    @should_rotate = false
+
+    # we use different counters for different tasks:
+    # delayed tasks are stored in the hash; for regular tasks we only maintain a counter;
+    # this information is only used during storage rotation
+    @immediate_tasks_in_queue = 0
+    @delayed_tasks = {}
+
+    # ensure data is written to disk
+    @storage_has_changes = false
+    Iodine.run_every(@fsync_frequency) do
+      if @storage_has_changes
+        @storage_has_changes = false
+        @storage.fsync
+      end
+    end
+  end
+
+  # Add a record to the log representing a new task.
+  # @param task [Rage::Deferred::Task]
+  # @param publish_at [Integer, nil]
+  # @param task_id [String, nil]
+  # @return [String]
+  def add(task, publish_at: nil, task_id: nil)
+    serialized_task = Marshal.dump(task).dump
+
+    persisted_task_id = task_id || generate_task_id
+
+    entry = build_add_entry(persisted_task_id, serialized_task, publish_at)
+    write_to_storage(entry)
+
+    if publish_at
+      @delayed_tasks[persisted_task_id] = [serialized_task, publish_at]
+    else
+      @immediate_tasks_in_queue += 1
+    end
+
+    persisted_task_id
+  end
+
+  # Add a record to the log representing a task removal.
+  # @param task_id [String]
+  def remove(task_id)
+    write_to_storage(build_remove_entry(task_id))
+
+    if @delayed_tasks.has_key?(task_id)
+      @delayed_tasks.delete(task_id)
+    else
+      @immediate_tasks_in_queue -= 1
+    end
+
+    # rotate the storage once the size is over the limit and all non-delayed tasks are processed
+    rotate_storage if @should_rotate && @immediate_tasks_in_queue == 0
+  end
+
+  # Return a list of pending tasks in the storage.
+  # @return [Array<(String, Rage::Deferred::Task, Integer, Integer)>
+  def pending_tasks
+    if @recovered_storages
+      # `@recovered_storages` will only be present if the server has previously crashed and left
+      # some storage files behind, or if the new cluster is started with fewer workers than before;
+      # TLDR: this code is expected to execute very rarely
+      @recovered_storages.each { |storage| recover_tasks(storage) }
+    end
+
+    tasks = {}
+    corrupted_tasks_count = 0
+
+    # find pending tasks in the storage
+    @storage.tap(&:rewind).each_line(chomp: true) do |entry|
+      signature, op, payload = entry[0...8], entry[9...12], entry[9..]
+      next if signature&.empty? || payload&.empty? || op&.empty?
+
+      unless signature == Zlib.crc32(payload).to_s(16).rjust(8, "0")
+        corrupted_tasks_count += 1
+        next
+      end
+
+      if op == "add"
+        task_id = entry[13...entry.index(":", 13).to_i]
+        tasks[task_id] = entry
+      elsif op == "rem"
+        task_id = entry[13..]
+        tasks.delete(task_id)
+      end
+    end
+
+    if corrupted_tasks_count != 0
+      puts "WARNING: Detected #{corrupted_tasks_count} corrupted deferred task(s)"
+    end
+
+    tasks.filter_map do |task_id, entry|
+      _, _, _, serialized_publish_at, serialized_task = entry.split(":", 5)
+
+      task = Marshal.load(serialized_task.undump)
+
+      publish_at = (serialized_publish_at == DEFAULT_PUBLISH_AT ? nil : serialized_publish_at.to_i)
+
+      if publish_at
+        @delayed_tasks[task_id] = [serialized_task, publish_at]
+      else
+        @immediate_tasks_in_queue += 1
+      end
+
+      [task_id, task, publish_at]
+
+    rescue ArgumentError, NameError => e
+      puts "ERROR: Can't deserialize the task with id #{task_id}: (#{e.class}) #{e.message}"
+      nil
+    end
+  end
+
+  private
+
+  def generate_task_id
+    @task_id_i += 1
+    "#{@task_id_base}-#{@task_id_i}"
+  end
+
+  def create_storage
+    file = @storage_path.join("#{@storage_prefix}-#{Time.now.strftime("%Y%m%d")}-#{Process.pid}-#{rand(0x100000000).to_s(36)}")
+
+    file.open("a+b").tap { |f| f.flock(File::LOCK_EX) }
+  end
+
+  def write_to_storage(content, adjust_size_limit: false)
+    @storage.write(content)
+    @storage_has_changes = true
+
+    @storage_size += content.bytesize
+    @should_rotate = true if @storage_size >= @storage_size_limit
+
+    if adjust_size_limit
+      # if the data copied from recovered storages or during the rotation takes up most of the storage, we might
+      # end up in an infinite rotation loop; instead, we dynamically increase the storage size limit
+      if @storage_size * STORAGE_SIZE_INCREASE_RATIO >= @storage_size_limit
+        @storage_size_limit *= STORAGE_SIZE_INCREASE_RATIO
+        @should_rotate = false
+      end
+    end
+  end
+
+  def rotate_storage
+    old_storage = @storage
+    @storage = nil # in case `create_storage` ends up blocking the fiber
+
+    # create a new storage and update internal state;
+    # after this point all new tasks will be written to the new storage
+    @should_rotate = false
+    @storage_size = 0
+    @storage_size_limit = DEFAULT_STORAGE_SIZE_LIMIT
+    @storage = create_storage
+
+    # copy delayed tasks to the new storage in batches
+    @delayed_tasks.keys.each_slice(100) do |task_ids|
+      entries = task_ids.filter_map do |task_id|
+        # don't copy the task if it has already been processed during the rotation
+        next unless @delayed_tasks.has_key?(task_id)
+
+        serialized_task, publish_at = @delayed_tasks[task_id]
+        build_add_entry(task_id, serialized_task, publish_at)
+      end
+
+      write_to_storage(entries.join, adjust_size_limit: true)
+
+      Fiber.pause
+    end
+
+    # delete the old storage ensuring the copied data has already been written to disk
+    Iodine.run_after(@fsync_frequency) do
+      old_storage.close
+      File.unlink(old_storage.path)
+    end
+  end
+
+  def build_add_entry(task_id, serialized_task, publish_at)
+    entry = "add:#{task_id}:#{publish_at || DEFAULT_PUBLISH_AT}:#{serialized_task}"
+    crc = Zlib.crc32(entry).to_s(16).rjust(8, "0")
+
+    "#{crc}:#{entry}\n"
+  end
+
+  def build_remove_entry(task_id)
+    entry = "rem:#{task_id}"
+    crc = Zlib.crc32(entry).to_s(16).rjust(8, "0")
+
+    "#{crc}:#{entry}\n"
+  end
+
+  def recover_tasks(storage)
+    # copy records to the main storage
+    while (content = storage.read(262_144))
+      write_to_storage(content, adjust_size_limit: true)
+    end
+
+    Iodine.run_after(@fsync_frequency) do
+      storage.close
+      File.unlink(storage.path)
+    end
+  end
+end

data/lib/rage/deferred/backends/nil.rb

@@ -0,0 +1,16 @@
+# frozen_string_literal: true
+
+class Rage::Deferred::Backends::Nil
+  def initialize(**)
+  end
+
+  def add(_, **)
+  end
+
+  def remove(_)
+  end
+
+  def pending_tasks
+    []
+  end
+end

data/lib/rage/deferred/deferred.rb

@@ -0,0 +1,95 @@
+# frozen_string_literal: true
+
+##
+# `Rage::Deferred` is an in-process background task queue with at-least-once delivery guarantee that allows you to schedule tasks to be executed later.
+# It can be used to offload long-running operations, such as sending emails or communicating with external APIs.
+#
+# To schedule a task, first define a task class that includes `Rage::Deferred::Task` and implements the `#perform` method.
+#
+# ```ruby
+# class SendWelcomeEmail
+#   include Rage::Deferred::Task
+#
+#   def perform(email)
+#     # logic to send the welcome email
+#   end
+# end
+# ```
+#
+# Then, push the task to the deferred queue:
+#
+# ```ruby
+# SendWelcomeEmail.enqueue(email: user.email)
+# ```
+#
+# You can also specify a delay for the task execution using the `delay` option:
+#
+# ```ruby
+# SendWelcomeEmail.enqueue(email: user.email, delay: 10) # execute after 10 seconds
+# ```
+#
+# Or you can specify a specific time in the future when the task should be executed:
+#
+# ```ruby
+# SendWelcomeEmail.enqueue(email: user.email, delay_until: Time.now + 3600) # execute in 1 hour
+# ```
+#
+module Rage::Deferred
+  # Push an instance to the deferred queue without including the `Rage::Deferred::Task` module.
+  # @param instance [Object] the instance to wrap
+  # @param delay [Integer, nil] the delay in seconds before the task is executed
+  # @param delay_until [Time, nil] the specific time when the task should be executed
+  # @example Schedule an arbitrary method to be called in the background
+  #   class SendWelcomeEmail < Struct.new(:email)
+  #     def call
+  #     end
+  #   end
+  #
+  #   email_service = SendWelcomeEmail.new(email: user.email)
+  #   Rage::Deferred.wrap(email_service).call
+  def self.wrap(instance, delay: nil, delay_until: nil)
+    Rage::Deferred::Proxy.new(instance, delay:, delay_until:)
+  end
+
+  # @private
+  def self.__backend
+    @__backend ||= Rage.config.deferred.backend
+  end
+
+  # @private
+  def self.__queue
+    @__queue ||= Rage::Deferred::Queue.new(__backend)
+  end
+
+  # @private
+  def self.__load_tasks
+    current_time = Time.now.to_i
+
+    __backend.pending_tasks.each do |task_id, task_wrapper, publish_at|
+      publish_in = publish_at - current_time if publish_at
+      __queue.schedule(task_id, task_wrapper, publish_in:)
+    rescue => e
+      puts "ERROR: Failed to load deferred task #{task_id}: #{e.class} (#{e.message}). Removing task from the queue."
+      __backend.remove(task_id)
+    end
+  end
+
+  module Backends
+  end
+
+  class PushTimeout < StandardError
+  end
+end
+
+require_relative "task"
+require_relative "queue"
+require_relative "proxy"
+require_relative "metadata"
+require_relative "backends/disk"
+require_relative "backends/nil"
+
+if Iodine.running?
+  Rage::Deferred.__load_tasks
+else
+  Iodine.on_state(:on_start) { Rage::Deferred.__load_tasks }
+end

data/lib/rage/deferred/metadata.rb

@@ -0,0 +1,43 @@
+# frozen_string_literal: true
+
+##
+# Metadata for deferred tasks.
+# The class encapsulates the metadata associated with a deferred task, and allows to store it without modifying the task instance.
+#
+class Rage::Deferred::Metadata
+  def self.build(task, args, kwargs)
+    request_id = Thread.current[:rage_logger][:tags][0] if Thread.current[:rage_logger]
+
+    [
+      task,
+      args.empty? ? nil : args,
+      kwargs.empty? ? nil : kwargs,
+      nil,
+      request_id
+    ]
+  end
+
+  def self.get_task(metadata)
+    metadata[0]
+  end
+
+  def self.get_args(metadata)
+    metadata[1]
+  end
+
+  def self.get_kwargs(metadata)
+    metadata[2]
+  end
+
+  def self.get_attempts(metadata)
+    metadata[3]
+  end
+
+  def self.get_request_id(metadata)
+    metadata[4]
+  end
+
+  def self.inc_attempts(metadata)
+    metadata[3] = metadata[3].to_i + 1
+  end
+end