rage-rb 1.19.2 → 1.20.0

data/lib/rage/configuration.rb

@@ -203,6 +203,14 @@ class Rage::Configuration
   end
   # @!endgroup
 
+  # @!group Telemetry Configuration
+  # Allows configuring telemetry settings.
+  # @return [Rage::Configuration::Telemetry]
+  def telemetry
+    @telemetry ||= Telemetry.new
+  end
+  # @!endgroup
+
   # @!group Session Configuration
   # Allows configuring session settings.
   # @return [Rage::Configuration::Session]
@@ -376,13 +384,13 @@ class Rage::Configuration
     end
   end
 
-  class Middleware
+  class MiddlewareRegistry
     # @private
-    attr_reader :middlewares
+    attr_reader :objects
 
     # @private
     def initialize
-      @middlewares = [[Rage::FiberWrapper]]
+      @objects = []
     end
 
     # Add a new middleware to the end of the stack.
@@ -400,7 +408,8 @@ class Rage::Configuration
     #     end
     #   end
     def use(new_middleware, *args, &block)
-      insert_after(@middlewares.length - 1, new_middleware, *args, &block)
+      validate!(-1, new_middleware)
+      @objects.insert(-1, [new_middleware, args, block])
     end
 
     # Insert a new middleware before an existing middleware in the stack.
@@ -419,11 +428,9 @@ class Rage::Configuration
     #     end
     #   end
     def insert_before(existing_middleware, new_middleware, *args, &block)
-      index = find_middleware_index(existing_middleware)
-      if index == 0 && @middlewares[0][0] == Rage::FiberWrapper
-        puts("Warning: inserting #{new_middleware} before Rage::FiberWrapper may lead to undefined behavior.")
-      end
-      @middlewares = (@middlewares[0...index] + [[new_middleware, args, block]] + @middlewares[index..]).uniq(&:first)
+      index = find_object_index(existing_middleware)
+      validate!(index, new_middleware)
+      @objects.insert(index, [new_middleware, args, block])
     end
 
     # Insert a new middleware after an existing middleware in the stack.
@@ -441,29 +448,65 @@ class Rage::Configuration
     #     end
     #   end
     def insert_after(existing_middleware, new_middleware, *args, &block)
-      index = find_middleware_index(existing_middleware)
-      @middlewares = (@middlewares[0..index] + [[new_middleware, args, block]] + @middlewares[index + 1..]).uniq(&:first)
+      index = find_object_index(existing_middleware) + 1
+      index = 0 if @objects.empty?
+      validate!(index, new_middleware)
+      @objects.insert(index, [new_middleware, args, block])
     end
 
     # Check if a middleware is included in the stack.
-    # @param middleware [Class, Integer] the middleware class or its index in the stack
+    # @param middleware [Class] the middleware class
     # @return [Boolean]
     def include?(middleware)
-      !!find_middleware_index(middleware) rescue false
+      @objects.any? { |o, _, _| o == middleware }
+    end
+
+    # Delete a middleware from the stack.
+    # @param middleware [Class] the middleware class
+    # @example
+    #   Rage.configure do
+    #     config.middleware.delete Rack::Cors
+    #   end
+    def delete(middleware)
+      @objects.reject! { |o, _, _| o == middleware }
     end
 
     private
 
-    def find_middleware_index(middleware)
-      if middleware.is_a?(Integer)
-        if middleware < 0 || middleware >= @middlewares.length
-          raise ArgumentError, "Middleware index should be in the (0...#{@middlewares.length}) range"
+    def find_object_index(object)
+      if object.is_a?(Integer)
+        if @objects[object] || object == 0
+          object
+        else
+          raise ArgumentError, "Could not find middleware at index #{object}"
         end
-        middleware
       else
-        @middlewares.index { |m, _, _| m == middleware }.tap do |i|
-          raise ArgumentError, "Couldn't find #{middleware} in the middleware stack" unless i
-        end
+        index = @objects.index { |o, _, _| o == object }
+        raise ArgumentError, "Could not find `#{object}` in the middleware registry" unless index
+        index
+      end
+    end
+
+    def validate!(_, _)
+    end
+  end
+
+  # See {Rage::Configuration::MiddlewareRegistry Rage::Configuration::MiddlewareRegistry} for details on available methods.
+  class Middleware < MiddlewareRegistry
+    # @private
+    alias_method :middlewares, :objects
+
+    # @private
+    def initialize
+      super
+      @objects = [[Rage::FiberWrapper]]
+    end
+
+    private
+
+    def validate!(index, middleware)
+      if index == 0 && @objects[0][0] == Rage::FiberWrapper
+        puts "WARNING: inserting the `#{middleware}` middleware before `Rage::FiberWrapper` may cause undefined behavior."
       end
     end
   end
@@ -730,6 +773,48 @@ class Rage::Configuration
       @backpressure = Backpressure.new(high_water_mark, low_water_mark, timeout)
     end
 
+    # Allows configuring middleware used by `Rage::Deferred`. See {MiddlewareRegistry} for details on available methods.
+    # @example
+    #   Rage.configure do
+    #     config.deferred.enqueue_middleware.use MyEnqueueMiddleware
+    #     config.deferred.enqueue_middleware.insert_before MyEnqueueMiddleware, MyLoggingMiddleware
+    #   end
+    class Middleware < Rage::Configuration::MiddlewareRegistry
+      private
+
+      def validate!(_, middleware)
+        unless middleware.is_a?(Class)
+          raise ArgumentError, "Deferred middleware has to be a class"
+        end
+
+        unless middleware.method_defined?(:call)
+          raise ArgumentError, "Deferred middleware has to implement the `#call` method"
+        end
+      end
+    end
+
+    # Configure enqueue middleware used by `Rage::Deferred`.
+    # See {EnqueueMiddlewareInterface} for details on the arguments passed to the middleware.
+    # @return [Rage::Configuration::Deferred::Middleware]
+    # @example
+    #   Rage.configure do
+    #     config.deferred.enqueue_middleware.use MyCustomMiddleware
+    #   end
+    def enqueue_middleware
+      @enqueue_middleware ||= Middleware.new
+    end
+
+    # Configure perform middleware used by `Rage::Deferred`.
+    # See {PerformMiddlewareInterface} for details on the arguments passed to the middleware.
+    # @return [Rage::Configuration::Deferred::Middleware]
+    # @example
+    #   Rage.configure do
+    #     config.deferred.perform_middleware.use MyCustomMiddleware
+    #   end
+    def perform_middleware
+      @perform_middleware ||= Middleware.new
+    end
+
     # @private
     def default_disk_storage_path
       Pathname.new("storage")
@@ -781,6 +866,65 @@ class Rage::Configuration
     end
   end
 
+  # The class allows configuring telemetry handlers. See {MiddlewareRegistry} for details on available methods.
+  # @example
+  #   Rage.configure do
+  #     config.telemetry.use MyTelemetryHandler.new
+  #   end
+  # @see Rage::Configuration::MiddlewareRegistry
+  # @see Rage::Telemetry
+  class Telemetry < MiddlewareRegistry
+    # @private
+    # @return [Hash{String => Array<Rage::Telemetry::HandlerRef>}] a map of span IDs to handler references
+    def handlers_map
+      @objects.map(&:first).each_with_object({}) do |handler, memo|
+        handlers_map = handler.is_a?(Class) ? handler.handlers_map : handler.class.handlers_map
+
+        handlers_map.each do |span_id, handler_methods|
+          handler_refs = handler_methods.map do |handler_method|
+            Rage::Telemetry::HandlerRef[handler, handler_method]
+          end
+
+          if memo[span_id]
+            memo[span_id] += handler_refs
+          else
+            memo[span_id] = handler_refs
+          end
+        end
+      end
+    end
+
+    private
+
+    def validate!(_, handler)
+      is_handler = if handler.is_a?(Class)
+        handler.ancestors.include?(Rage::Telemetry::Handler)
+      else
+        handler.is_a?(Rage::Telemetry::Handler)
+      end
+
+      unless is_handler
+        raise ArgumentError, "Cannot add `#{handler}` as a telemetry handler; should inherit `Rage::Telemetry::Handler`"
+      end
+
+      handlers_map = if handler.is_a?(Class)
+        handler.handlers_map
+      else
+        handler.class.handlers_map
+      end
+
+      unless handlers_map&.any?
+        raise ArgumentError, "Telemetry handler `#{handler}` does not define any handlers"
+      end
+
+      handlers_map.values.reduce(&:+).each do |handler_method|
+        unless handler.respond_to?(handler_method)
+          raise ArgumentError, "Telemetry handler `#{handler}` does not implement the `#{handler_method}` handler method"
+        end
+      end
+    end
+  end
+
   class Session
     # @!attribute key
     #   Specify the name of the session cookie.
@@ -847,6 +991,13 @@ class Rage::Configuration
       Rage.__log_processor.add_custom_tags(@log_tags.objects)
       @logger.dynamic_tags = Rage.__log_processor.dynamic_tags
     end
+
+    if defined?(::Rack::Events) && middleware.include?(::Rack::Events)
+      middleware.delete(Rage::BodyFinalizer)
+      middleware.insert_before(::Rack::Events, Rage::BodyFinalizer)
+    end
+
+    Rage::Telemetry.__setup if @telemetry
   end
 end
 
@@ -882,3 +1033,66 @@ end
 #     def call(severity:, tags:, context:, message:, request_info:)
 #     end
 #   end
+
+# @!parse [ruby]
+#   # @note This class does not exist at runtime and is used for documentation purposes only. Do not inherit your middleware classes from it.
+#   class EnqueueMiddlewareInterface
+#     # Called whenever a deferred task is enqueued.
+#     #
+#     # The middleware is expected to call `yield` to pass control to the next middleware in the stack. If the middleware does not call `yield`, the task will not be enqueued.
+#     #
+#     # Rage automatically detects which parameters your middleware's `#call` method accepts, and only passes those parameters. You can omit any of the described parameters in your implementation.
+#     #
+#     # @param task_class [Class] the deferred task class
+#     # @param delay [Integer, nil] the delay in seconds before the task is executed
+#     # @param delay_until [Time, Integer, nil] the time at which the task should be executed
+#     # @param phase [:enqueue] the middleware phase. Useful for middlewares that are shared between enqueue and perform phases
+#     # @param args [Array] the positional arguments passed to the task
+#     # @param kwargs [Hash] the keyword arguments passed to the task
+#     # @param context [Hash] the context is serialized together with the task and allows passing data between middlewares without exposing it to the task itself
+#     # @example
+#     #   class EncryptArgumentsMiddleware
+#     #     def call(args:, kwargs:)
+#     #       args.map! { |arg| MyEncryptionSDK.encrypt(arg) }
+#     #       kwargs.transform_values! { |value| MyEncryptionSDK.encrypt(value) }
+#     #
+#     #       yield
+#     #     end
+#     #   end
+#     def call(task_class:, delay:, delay_until:, phase:, args:, kwargs:, context:)
+#     end
+#   end
+
+# @!parse [ruby]
+#   # @note This class does not exist at runtime and is used for documentation purposes only. Do not inherit your middleware classes from it.
+#   class PerformMiddlewareInterface
+#     # Called whenever a deferred task is performed.
+#     #
+#     # The middleware is expected to call `yield` to pass control to the next middleware in the stack. If the middleware does not call `yield`, the task will not be performed.
+#     #
+#     # Rage automatically detects which parameters your middleware's `#call` method accepts, and only passes those parameters. You can omit any of the described parameters in your implementation.
+#     #
+#     # @param task_class [Class] the deferred task class
+#     # @param task [Rage::Deferred::Task] the deferred task instance
+#     # @param phase [:perform] the middleware phase. Useful for middlewares that are shared between enqueue and perform phases
+#     # @param args [Array] the positional arguments passed to the task
+#     # @param kwargs [Hash] the keyword arguments passed to the task
+#     # @param context [Hash] the context is serialized together with the task and allows passing data between middlewares without exposing it to the task itself
+#     # @example
+#     #   class DecryptArgumentsMiddleware
+#     #     def call(args:, kwargs:)
+#     #       args.map! { |arg| MyEncryptionSDK.decrypt(arg) }
+#     #       kwargs.transform_values! { |value| MyEncryptionSDK.decrypt(value) }
+#     #
+#     #       yield
+#     #
+#     #     rescue
+#     #       # Re-encrypt the arguments in case of an error
+#     #       args.map! { |arg| MyEncryptionSDK.encrypt(arg) }
+#     #       kwargs.transform_values! { |value| MyEncryptionSDK.encrypt(value) }
+#     #       raise
+#     #     end
+#     #   end
+#     def call(task_class:, task:, phase:, args:, kwargs:, context:)
+#     end
+#   end

data/lib/rage/controller/api.rb

@@ -112,56 +112,58 @@ class RageController::API
 
       class_eval <<~RUBY, __FILE__, __LINE__ + 1
         def __run_#{action}
-          #{if query_cache_enabled
-            <<~RUBY
-              ActiveRecord::Base.connection_pool.enable_query_cache!
-            RUBY
-          end}
+          Rage::Telemetry.tracer.span_controller_action_process(controller: self, params: @__params) do
+            #{if query_cache_enabled
+              <<~RUBY
+                ActiveRecord::Base.connection_pool.enable_query_cache!
+              RUBY
+            end}
 
-          #{wrap_parameters_chunk}
-          #{before_actions_chunk}
-          #{action} unless @__before_callback_rendered
-          #{around_actions_end_chunk}
+            #{wrap_parameters_chunk}
+            #{before_actions_chunk}
+            #{action} unless @__before_callback_rendered
+            #{around_actions_end_chunk}
 
-          #{if !after_actions_chunk.empty?
-            <<~RUBY
-              unless @__before_callback_rendered
-                @__rendered = true
-                #{after_actions_chunk}
-              end
-            RUBY
-          end}
+            #{if !after_actions_chunk.empty?
+              <<~RUBY
+                unless @__before_callback_rendered
+                  @__rendered = true
+                  #{after_actions_chunk}
+                end
+              RUBY
+            end}
 
-          [@__status, @__headers, @__body]
+            [@__status, @__headers, @__body]
 
-          #{rescue_handlers_chunk}
+            #{rescue_handlers_chunk}
 
-        ensure
-          #{if query_cache_enabled
-            <<~RUBY
-              ActiveRecord::Base.connection_pool.disable_query_cache!
-            RUBY
-          end}
+          ensure
+            #{if query_cache_enabled
+              <<~RUBY
+                ActiveRecord::Base.connection_pool.disable_query_cache!
+              RUBY
+            end}
 
-          #{if should_release_connections
-            <<~RUBY
-              ActiveRecord::Base.connection_handler.clear_active_connections!(:all)
-            RUBY
-          end}
+            #{if should_release_connections
+              <<~RUBY
+                ActiveRecord::Base.connection_handler.clear_active_connections!(:all)
+              RUBY
+            end}
 
-          #{if method_defined?(:append_info_to_payload) || private_method_defined?(:append_info_to_payload)
-            <<~RUBY
-              context = {}
-              append_info_to_payload(context)
-
-              log_context = Thread.current[:rage_logger][:context]
-              if log_context.empty?
-                Thread.current[:rage_logger][:context] = context
-              else
-                Thread.current[:rage_logger][:context] = log_context.merge(context)
-              end
-            RUBY
-          end}
+            #{if method_defined?(:append_info_to_payload) || private_method_defined?(:append_info_to_payload)
+              <<~RUBY
+                context = {}
+                append_info_to_payload(context)
+
+                log_context = Thread.current[:rage_logger][:context]
+                if log_context.empty?
+                  Thread.current[:rage_logger][:context] = context
+                else
+                  Thread.current[:rage_logger][:context] = log_context.merge(context)
+                end
+              RUBY
+            end}
+          end
         end
       RUBY
     end
@@ -451,6 +453,9 @@ class RageController::API
     @__rendered = false
   end
 
+  # @private
+  attr_reader :__env, :__status, :__headers, :__body
+
   # Get the request object. See {Rage::Request}.
   # @return [Rage::Request]
   def request
@@ -460,7 +465,7 @@ class RageController::API
   # Get the response object. See {Rage::Response}.
   # @return [Rage::Response]
   def response
-    @response ||= Rage::Response.new(@__headers, @__body)
+    @response ||= Rage::Response.new(self)
   end
 
   # Get the cookie object. See {Rage::Cookies}.

data/lib/rage/deferred/context.rb

@@ -5,7 +5,7 @@
 # The class encapsulates the context associated with a deferred task, and allows to store it without modifying the task instance.
 #
 class Rage::Deferred::Context
-  def self.build(task, args, kwargs, storage: nil)
+  def self.build(task, args, kwargs)
     logger = Thread.current[:rage_logger]
 
     [
@@ -14,35 +14,63 @@ class Rage::Deferred::Context
       kwargs.empty? ? nil : kwargs,
       nil,
       logger&.dig(:tags),
-      logger&.dig(:context)
+      logger&.dig(:context),
+      nil
     ]
   end
 
+  # @return [Class] the task class
   def self.get_task(context)
     context[0]
   end
 
+  # @return [Array, nil] arguments the task was enqueued with
   def self.get_args(context)
     context[1]
   end
 
+  # @return [Array] arguments the task was enqueued with, creating it if it does not exist
+  def self.get_or_create_args(context)
+    context[1] ||= []
+  end
+
+  # @return [Hash, nil] keyword arguments the task was enqueued with
   def self.get_kwargs(context)
     context[2]
   end
 
+  # @return [Hash] keyword arguments the task was enqueued with, creating it if it does not exist
+  def self.get_or_create_kwargs(context)
+    context[2] ||= {}
+  end
+
+  # @return [Integer, nil] number of attempts made to process the task
   def self.get_attempts(context)
     context[3]
   end
 
+  # Increments the number of attempts made to process the task
   def self.inc_attempts(context)
     context[3] = context[3].to_i + 1
   end
 
+  # @return [Array, nil] log tags associated with the task
   def self.get_log_tags(context)
     context[4]
   end
 
+  # @return [Hash, nil] log context associated with the task
   def self.get_log_context(context)
     context[5]
   end
+
+  # @return [Hash, nil] user context associated with the task
+  def self.get_user_context(context)
+    context[6]
+  end
+
+  # @return [Hash] user context associated with the task, creating it if it does not exist
+  def self.get_or_create_user_context(context)
+    context[6] ||= {}
+  end
 end

data/lib/rage/deferred/deferred.rb

@@ -74,14 +74,24 @@ module Rage::Deferred
     end
   end
 
-  module Backends
+  # @private
+  def self.__middleware_chain
+    @__middleware_chain ||= MiddlewareChain.new(
+      enqueue_middleware: Rage.config.deferred.enqueue_middleware.objects,
+      perform_middleware: Rage.config.deferred.perform_middleware.objects
+    )
   end
 
-  class PushTimeout < StandardError
+  # @private
+  def self.__initialize
+    __middleware_chain
+    __load_tasks
   end
 
-  # @private
-  class TaskFailed < StandardError
+  module Backends
+  end
+
+  class PushTimeout < StandardError
   end
 end
 
@@ -89,11 +99,13 @@ require_relative "task"
 require_relative "queue"
 require_relative "proxy"
 require_relative "context"
+require_relative "metadata"
+require_relative "middleware_chain"
 require_relative "backends/disk"
 require_relative "backends/nil"
 
 if Iodine.running?
-  Rage::Deferred.__load_tasks
+  Rage::Deferred.__initialize
 else
-  Iodine.on_state(:on_start) { Rage::Deferred.__load_tasks }
+  Iodine.on_state(:on_start) { Rage::Deferred.__initialize }
 end

data/lib/rage/deferred/metadata.rb

@@ -0,0 +1,39 @@
+# frozen_string_literal: true
+
+##
+# Provides metadata about the current deferred task execution.
+#
+class Rage::Deferred::Metadata
+  class << self
+    # Returns the current attempt number.
+    # @return [Integer] the current attempt number (1 for the first run)
+    def attempts
+      Rage::Deferred::Context.get_attempts(context).to_i + 1
+    end
+
+    # Returns the number of retries that have occurred for the current task.
+    # @return [Integer] the number of retries (0 on first run, 1+ on retries)
+    def retries
+      attempts - 1
+    end
+
+    # Checks whether this is a retry execution.
+    # @return [Boolean] `true` if this is a retry, `false` if this is the first run
+    def retrying?
+      attempts > 1
+    end
+
+    # Checks whether the task will be retried if the current execution fails.
+    # @return [Boolean] `true` if a failure will schedule another attempt, `false` otherwise
+    def will_retry?
+      task = Rage::Deferred::Context.get_task(context)
+      task.__should_retry?(attempts)
+    end
+
+    private
+
+    def context
+      Fiber[Rage::Deferred::Task::CONTEXT_KEY]
+    end
+  end
+end

data/lib/rage/deferred/middleware_chain.rb

@@ -0,0 +1,67 @@
+# frozen_string_literal: true
+
+class Rage::Deferred::MiddlewareChain
+  def initialize(enqueue_middleware:, perform_middleware:)
+    @enqueue_middleware = enqueue_middleware
+    @perform_middleware = perform_middleware
+
+    build_enqueue_chain!
+    build_perform_chain!
+  end
+
+  private
+
+  def build_enqueue_chain!
+    raw_arguments = {
+      phase: ":enqueue",
+      args: "Rage::Deferred::Context.get_or_create_args(context)",
+      kwargs: "Rage::Deferred::Context.get_or_create_kwargs(context)",
+      context: "Rage::Deferred::Context.get_or_create_user_context(context)",
+      task_class: "Rage::Deferred::Context.get_task(context)",
+      delay: "delay",
+      delay_until: "delay_until"
+    }
+
+    self.class.class_eval <<~RUBY, __FILE__, __LINE__ + 1
+      def with_enqueue_middleware(context, delay:, delay_until:)
+        #{build_middleware_chain(:@enqueue_middleware, raw_arguments)}
+      end
+    RUBY
+  end
+
+  def build_perform_chain!
+    raw_arguments = {
+      phase: ":perform",
+      args: "Rage::Deferred::Context.get_or_create_args(context)",
+      kwargs: "Rage::Deferred::Context.get_or_create_kwargs(context)",
+      context: "Rage::Deferred::Context.get_or_create_user_context(context)",
+      task_class: "task.class",
+      task: "task"
+    }
+
+    self.class.class_eval <<~RUBY, __FILE__, __LINE__ + 1
+      def with_perform_middleware(context, task:)
+        #{build_middleware_chain(:@perform_middleware, raw_arguments)}
+      end
+    RUBY
+  end
+
+  def build_middleware_chain(middlewares_var, raw_arguments)
+    middlewares = instance_variable_get(middlewares_var)
+    i = middlewares.length
+
+    middlewares.reverse.inject("yield") do |memo, middleware_with_args|
+      middleware, _, _ = middleware_with_args
+      arguments = Rage::Internal.build_arguments(middleware.instance_method(:call), raw_arguments)
+      i -= 1
+
+      <<~RUBY
+        middleware, args, block = #{middlewares_var}[#{i}]
+
+        middleware.new(*args, &block).call(#{arguments}) do
+          #{memo}
+        end
+      RUBY
+    end
+  end
+end