rails_mind 0.1.0 → 0.3.0

data/lib/rails_mind/integrations/action_mailer.rb

@@ -0,0 +1,195 @@
+require "active_support/notifications"
+require "securerandom"
+
+module RailsMind
+  module Integrations
+    # Wrap the public delivery entry points without touching arguments or mail
+    # contents. The stored class/action names are the only MessageDelivery state
+    # read here; reading .message would prematurely render the lazy message.
+    module MailerDeliveryContext
+      def deliver_now
+        MailerNotifications.observe_delivery(self, forced: false) { super }
+      end
+
+      def deliver_now!
+        MailerNotifications.observe_delivery(self, forced: true) { super }
+      end
+    end
+
+    module MailerCallbacks
+      def self.included(base)
+        base.around_deliver do |mailer, block|
+          MailerNotifications.record_delivery_options(mailer)
+          block.call
+        end
+      end
+    end
+
+    # Mail#deliver! bypasses ActionMailer's notification handler as well as its
+    # delivery flags. Observe that public method only inside a known forced
+    # ActionMailer invocation, preserving its return value and exceptions.
+    module ForcedMailerTransport
+      def deliver!(...)
+        MailerNotifications.observe_forced_transport(self) { super }
+      end
+    end
+
+    module MailerNotifications
+      class << self
+        def install!
+          return if @subscriptions
+          @subscriptions = []
+          subscribe("process.action_mailer") do |event|
+            payload = event.payload
+            properties = metadata(payload[:mailer], payload[:action]).merge(operation: "process")
+            properties[:attempt_id] = attempt[:attempt_id] if attempt
+            emit(event, properties, status: event_error(event) ? "error" : "observed")
+          end
+          subscribe("deliver.action_mailer") do |event|
+            current = attempt
+            current[:observed] = true if current
+            properties = metadata(event.payload[:mailer], current&.dig(:action)).merge(
+              operation: "delivery_attempt_completed", attempt_id: current&.dig(:attempt_id) || SecureRandom.uuid)
+            properties.merge!(delivery_options(current[:message])) if current
+            properties[:forced] = current[:forced] if current
+            properties[:perform_deliveries] = event.payload[:perform_deliveries] if [ true, false ].include?(event.payload[:perform_deliveries])
+            status = if event_error(event) then "error"
+            elsif properties[:perform_deliveries] == false && properties[:forced] == false then "disabled"
+            else "observed"
+            end
+            emit(event, properties, status: status)
+          end
+        end
+
+        def uninstall!
+          @subscriptions&.each { |subscription| ActiveSupport::Notifications.unsubscribe(subscription) }
+          @subscriptions = nil
+        end
+
+        def observe_delivery(delivery, forced:)
+          current = new_attempt(delivery, forced: forced)
+          return yield unless current
+
+          previous = attempt
+          Thread.current[:rails_mind_mailer_attempt] = current
+          Context.with_trace do
+            result = nil
+            failure = nil
+            begin
+              result = yield
+            rescue StandardError => error
+              failure = error
+              raise
+            ensure
+              finish_unobserved(current, result, failure)
+              RailsMind.capture_exception(failure, source: "action_mailer") if failure
+            end
+          end
+        ensure
+          Thread.current[:rails_mind_mailer_attempt] = previous if current
+        end
+
+        # Configuration describes the attempted transport, never its recipients,
+        # credentials, provider responses, or eventual delivery outcome.
+        def record_delivery_options(mailer)
+          return unless attempt
+          attempt[:message] = mailer.message
+        rescue StandardError
+          nil
+        end
+
+        def observe_forced_transport(message)
+          current = attempt
+          # Callbacks/interceptors can send an unrelated raw Mail::Message while
+          # this attempt is active. Only attribute the ActionMailer message that
+          # our delivery callback recorded, never another message on the stack.
+          return yield unless current && current[:forced] && current[:message].equal?(message)
+          current[:observed] = true
+          started_at = Time.now
+          clock = Process.clock_gettime(Process::CLOCK_MONOTONIC)
+          failure = nil
+          begin
+            yield
+          rescue StandardError => error
+            failure = error
+            raise
+          ensure
+            record_forced_transport(current, message, started_at, clock, failure)
+          end
+        end
+
+        private
+          def attempt = Thread.current[:rails_mind_mailer_attempt]
+
+          def new_attempt(delivery, forced:)
+            { mailer: delivery.instance_variable_get(:@mailer_class).name,
+              action: delivery.instance_variable_get(:@action).to_s,
+              forced: forced, attempt_id: SecureRandom.uuid, observed: false,
+              started_at: Time.now, started_clock: Process.clock_gettime(Process::CLOCK_MONOTONIC) }
+          rescue StandardError
+            nil
+          end
+
+          def metadata(mailer, action)
+            { instrumentation: "action_mailer", schema_version: 1, mailer: mailer, action: action }.reject { |_, value| value.nil? }
+          end
+
+          def delivery_options(message)
+            return {} unless message.is_a?(::Mail::Message)
+            { perform_deliveries: message.perform_deliveries, raise_delivery_errors: message.raise_delivery_errors,
+              delivery_method: message.delivery_method.class.name }.reject { |_, value| value.nil? }
+          rescue StandardError
+            {}
+          end
+
+          def record_forced_transport(current, message, started_at, clock, failure)
+            properties = metadata(current[:mailer], current[:action]).merge(delivery_options(message)).merge(
+              operation: "delivery_attempt_completed", attempt_id: current[:attempt_id], forced: true)
+            properties[:exception_class] = failure.class.name if failure
+            RailsMind.client.emit(:span, "#{current[:mailer]}##{current[:action]} delivery_attempt_completed",
+              occurred_at: started_at, duration_ms: (Process.clock_gettime(Process::CLOCK_MONOTONIC) - clock) * 1000,
+              status: failure ? "error" : "observed", properties: properties)
+            RailsMind.capture_exception(failure, source: "action_mailer") if failure
+          rescue StandardError
+            nil
+          end
+
+          def finish_unobserved(current, result, failure)
+            return if current[:observed]
+            properties = metadata(current[:mailer], current[:action]).merge(delivery_options(current[:message])).merge(
+              operation: "delivery_not_observed", attempt_id: current[:attempt_id], forced: current[:forced])
+            properties[:exception_class] = failure.class.name if failure
+            status = failure ? "error" : (result == false ? "aborted" : "unknown")
+            RailsMind.client.emit(:span, "#{current[:mailer]}##{current[:action]} delivery_not_observed",
+              occurred_at: current[:started_at], duration_ms: (Process.clock_gettime(Process::CLOCK_MONOTONIC) - current[:started_clock]) * 1000,
+              status: status, properties: properties)
+          rescue StandardError
+            nil
+          end
+
+          def event_error(event)
+            event.payload[:exception_object] || event.payload[:error]
+          end
+
+          def emit(event, properties, status:)
+            error = event_error(event)
+            properties[:exception_class] = error.class.name if error
+            Context.with_trace do
+              label = [ properties[:mailer], properties[:action] ].compact.join("#")
+              RailsMind.client.emit(:span, "#{label} #{properties[:operation]}",
+                occurred_at: Time.now - event.duration / 1000, duration_ms: event.duration, status: status, properties: properties)
+              RailsMind.capture_exception(error, source: "action_mailer") if error
+            end
+          end
+
+          def subscribe(name, &callback)
+            @subscriptions << ActiveSupport::Notifications.subscribe(name) do |event|
+              callback.call(event) unless Context.suppressed?
+            rescue StandardError
+              nil
+            end
+          end
+      end
+    end
+  end
+end

data/lib/rails_mind/integrations/rails.rb

@@ -14,17 +14,18 @@ module RailsMind
 
       def call(env)
         Context.with(request_id: env["action_dispatch.request_id"] || SecureRandom.uuid) do
-          previous = Thread.current[:rails_mind_queries]
-          Thread.current[:rails_mind_queries] = { count: 0, fingerprints: Hash.new(0) }
-          @app.call(env)
-        rescue StandardError => error
-          # Preserve identity before this request boundary unwinds. The Rails
-          # reporter may observe the same exception later; object dedup prevents
-          # a second RailsMind event while other reporters remain unaffected.
-          RailsMind.capture_exception(error, source: "rack")
-          raise
-        ensure
-          Thread.current[:rails_mind_queries] = previous
+          Context.with_trace do
+            previous = Thread.current[:rails_mind_queries]
+            Thread.current[:rails_mind_queries] = { count: 0, fingerprints: Hash.new(0) }
+            @app.call(env)
+          rescue StandardError => error
+            # Capture identity before this request boundary unwinds. Object
+            # dedup still leaves other error reporters unaffected.
+            RailsMind.capture_exception(error, source: "rack")
+            raise
+          ensure
+            Thread.current[:rails_mind_queries] = previous
+          end
         end
       end
     end
@@ -32,16 +33,37 @@ module RailsMind
     module JobContext
       def self.included(base)
         base.around_enqueue do |job, block|
-          job.rails_mind_context = Context.snapshot.transform_keys(&:to_s)
-          block.call
+          Context.with_trace do
+            job.rails_mind_context = Context.effective_snapshot.transform_keys(&:to_s)
+            block.call
+          end
         end
         base.around_perform do |job, block|
-          job.rails_mind_started_at = Time.now
-          Context.with((job.rails_mind_context || {}).merge("job_id" => job.job_id)) { block.call }
+          job.rails_mind_started_at ||= Time.now
+          Context.with((job.rails_mind_context || {}).merge("job_id" => job.job_id)) do
+            Context.with_trace do
+              job.rails_mind_context = Context.effective_snapshot.transform_keys(&:to_s)
+              block.call
+            end
+          end
         end
       end
 
-      attr_accessor :rails_mind_context, :rails_mind_started_at
+      attr_accessor :rails_mind_context, :rails_mind_started_at, :rails_mind_execution_timing
+
+      # ActiveJob runs retry/discard handlers after around_perform has unwound.
+      # Keep the restored boundary around the entire execution so those handlers
+      # and their newly serialized attempts retain the originating trace.
+      def perform_now(...)
+        self.rails_mind_started_at = nil
+        self.rails_mind_execution_timing = nil
+        Context.with((rails_mind_context || {}).merge("job_id" => job_id)) do
+          Context.with_trace do
+            self.rails_mind_context = Context.effective_snapshot.transform_keys(&:to_s)
+            super
+          end
+        end
+      end
 
       def serialize = super.merge("rails_mind_context" => rails_mind_context || {})
 
@@ -76,17 +98,38 @@ module RailsMind
             fingerprint = Digest::SHA256.hexdigest(normalized)[0, 24]
             queries[:fingerprints][fingerprint] += 1 if queries[:fingerprints].key?(fingerprint) || queries[:fingerprints].size < 100
           end
-          %w[perform.active_job enqueue_retry.active_job retry_stopped.active_job discard.active_job].each do |name|
+          %w[enqueue.active_job enqueue_at.active_job perform_start.active_job perform.active_job enqueue_retry.active_job retry_stopped.active_job discard.active_job].each do |name|
             subscribe(name) do |event|
               job = event.payload[:job]
               next unless job
-              error = event.payload[:exception_object] || event.payload[:error]
-              Context.with((job.rails_mind_context || {}).merge("job_id" => job.job_id)) do
-                RailsMind.client.emit(:job, job.class.name, duration_ms: event.duration,
-                  status: error ? "error" : (name == "perform.active_job" ? "ok" : name.split(".").first),
-                  properties: { operation: name.split(".").first, queue: job.queue_name, executions: job.executions,
-                    queue_delay_ms: job.enqueued_at && job.rails_mind_started_at ? [ (job.rails_mind_started_at.to_f - job.enqueued_at.to_f) * 1000, 0 ].max : nil })
-                RailsMind.capture_exception(error, source: "active_job") if error
+              operation = name.split(".").first
+              enqueue = %w[enqueue enqueue_at].include?(operation)
+              error = event.payload[:exception_object] || event.payload[:error] || (job.enqueue_error if enqueue)
+              if operation == "perform_start"
+                job.rails_mind_started_at = Time.now
+                job.rails_mind_execution_timing = { enqueued_at: job.enqueued_at, scheduled_at: job.scheduled_at }
+              end
+              # A callback may abort before the SDK's around_enqueue is entered.
+              values = job.rails_mind_context || Context.effective_snapshot
+              Context.with(values.merge("job_id" => job.job_id)) do
+                Context.with_trace do
+                  job.rails_mind_context = Context.effective_snapshot.transform_keys(&:to_s)
+                  status = if error
+                    "error"
+                  elsif event.payload[:aborted]
+                    "aborted"
+                  elsif enqueue
+                    "enqueued"
+                  elsif operation == "perform_start"
+                    "started"
+                  else
+                    operation == "perform" ? "ok" : operation
+                  end
+                  kind = enqueue || operation == "perform_start" ? :span : :job
+                  RailsMind.client.emit(kind, job.class.name, duration_ms: event.duration, status: status,
+                    properties: job_properties(job, operation))
+                  RailsMind.capture_exception(error, source: "active_job") if error
+                end
               end
             end
           end
@@ -107,6 +150,29 @@ module RailsMind
 
         private
 
+        def job_properties(job, operation)
+          timing = if %w[perform_start perform].include?(operation)
+            job.rails_mind_execution_timing || {}
+          else
+            { enqueued_at: job.enqueued_at, scheduled_at: job.scheduled_at }
+          end
+          enqueued_at, scheduled_at = timing.values_at(:enqueued_at, :scheduled_at)
+          execution_timing = job.rails_mind_execution_timing || {}
+          queued = execution_timing[:enqueued_at]
+          scheduled = execution_timing[:scheduled_at]
+          started = job.rails_mind_started_at
+          properties = { instrumentation: "active_job", schema_version: 1, operation: operation,
+            job_class: job.class.name, queue: job.queue_name, executions: job.executions,
+            enqueued_at: enqueued_at&.iso8601(6), scheduled_at: scheduled_at&.iso8601(6) }
+          unless %w[enqueue enqueue_at].include?(operation)
+            properties[:queue_delay_ms] = [ (started.to_f - queued.to_f) * 1000, 0 ].max if started && queued
+            if started && queued && scheduled
+              properties[:eligible_queue_delay_ms] = [ (started.to_f - [ queued.to_f, scheduled.to_f ].max) * 1000, 0 ].max
+            end
+          end
+          properties.compact
+        end
+
         def subscribe(name, &callback)
           @subscriptions << ActiveSupport::Notifications.subscribe(name) do |event|
             callback.call(event) unless Context.suppressed?

data/lib/rails_mind/railtie.rb

@@ -1,4 +1,5 @@
 require_relative "integrations/rails"
+require_relative "integrations/action_mailer"
 
 module RailsMind
   class Railtie < Rails::Railtie
@@ -8,10 +9,18 @@ module RailsMind
 
     config.after_initialize do
       Integrations::RailsNotifications.install!
+      Integrations::MailerNotifications.install!
       Rails.error.subscribe(Integrations::ErrorSubscriber.new)
       ActiveSupport.on_load(:active_job) do
         include Integrations::JobContext unless included_modules.include?(Integrations::JobContext)
       end
+      ActiveSupport.on_load(:action_mailer) do
+        include Integrations::MailerCallbacks unless included_modules.include?(Integrations::MailerCallbacks)
+        unless ActionMailer::MessageDelivery.ancestors.include?(Integrations::MailerDeliveryContext)
+          ActionMailer::MessageDelivery.prepend(Integrations::MailerDeliveryContext)
+        end
+        Mail::Message.prepend(Integrations::ForcedMailerTransport) unless Mail::Message.ancestors.include?(Integrations::ForcedMailerTransport)
+      end
       at_exit { RailsMind.shutdown }
     end
 

data/lib/rails_mind/sampling.rb

@@ -0,0 +1,211 @@
+require "json"
+require "securerandom"
+require "time"
+require "digest"
+
+module RailsMind
+  # Called under Collector's mutex. Decisions and queue-loss accounting share
+  # the same critical section as report rotation; reports are immutable once cut.
+  class Sampling
+    REPORT_INTERVAL = 60
+    POLICY_INTERVAL = 60
+    MAX_POLICY_TTL = 600
+    MAX_ENDPOINTS = 100
+    MAX_PENDING = 5
+    MAX_COUNT = 1_000_000_000
+    MAX_REPORT_AGE = 7 * 86_400
+    ENDPOINT = /\A[A-Z][A-Za-z0-9_]*(?:::[A-Z][A-Za-z0-9_]*)*#[A-Za-z_][A-Za-z0-9_!?]*\z/
+    UUID = /\A[0-9a-f]{8}-(?:[0-9a-f]{4}-){3}[0-9a-f]{12}\z/i
+
+    def initialize(config, clock: -> { Time.now.utc }, monotonic: -> { Process.clock_gettime(Process::CLOCK_MONOTONIC) })
+      @config, @clock, @monotonic = config, clock, monotonic
+      @counts, @reports = {}, []
+      @from = @clock.call
+      @last_rotation = @monotonic.call
+      @next_policy_at = @next_report_at = 0
+      @policy = nil
+      @activated = false
+      @limited = false
+      @policy_ids = []
+      @stats = { observed: 0, retained: 0, sampled_out: 0, queue_dropped: 0, unreported: 0, reports_lost: 0 }
+    end
+
+    def consider(event)
+      return [ true, nil ] unless event["kind"] == "request"
+      @activated = true
+      rate, state, policy_id = effective_policy
+      name = event["name"]
+      unless eligible?(event) && name.is_a?(String) && name.bytesize <= 200 && name.match?(ENDPOINT)
+        annotate(event, 1.0, "protected", policy_id)
+        return [ true, nil ]
+      end
+      if !@counts.key?(name) && @counts.size >= MAX_ENDPOINTS
+        @stats[:unreported] += 1
+        @limited = true
+        annotate(event, 1.0, "unavailable", policy_id)
+        return [ true, nil ]
+      end
+      count = @counts[name] ||= { "name" => name, "observed" => 0, "retained" => 0, "sampled_out" => 0, "queue_dropped" => 0 }
+      if count["observed"] >= MAX_COUNT
+        @stats[:unreported] += 1
+        @limited = true
+        annotate(event, 1.0, "unavailable", policy_id)
+        return [ true, nil ]
+      end
+      rate = @policy[:endpoints].fetch(name, rate) if state == "active"
+      selected = rate >= 1 || (rate.positive? && fraction(event["id"]) < rate)
+      count["observed"] += 1
+      count[selected ? "retained" : "sampled_out"] += 1
+      @stats[:observed] += 1
+      @stats[selected ? :retained : :sampled_out] += 1
+      @policy_ids << (policy_id || state) unless @policy_ids.include?(policy_id || state) || @policy_ids.size >= 2
+      annotate(event, rate, state, policy_id) if selected
+      [ selected, count ]
+    end
+
+    def queue_dropped(count)
+      return unless count
+      count["queue_dropped"] += 1
+      @stats[:queue_dropped] += 1
+    end
+
+    def apply_policy(value)
+      return false unless value.is_a?(Hash) && value["version"] == 1 && boolean?(value["enabled"])
+      id = value["id"]
+      return false unless id.is_a?(String) && id.bytesize.between?(1, 200) && !id.match?(/\p{Cc}/)
+      issued = parse_time(value["issued_at"])
+      expires = parse_time(value["expires_at"])
+      now = @clock.call
+      return false unless issued && expires && issued <= now + 300 && expires > now && expires > issued && expires - issued <= MAX_POLICY_TTL
+      return false if @policy && issued < @policy[:issued_at]
+      rate = value["default_rate"]
+      rows = value["endpoints"]
+      return false unless valid_rate?(rate) && rows.is_a?(Array) && rows.size <= MAX_ENDPOINTS
+      endpoints = {}
+      rows.each do |row|
+        return false unless row.is_a?(Hash) && row["name"].is_a?(String) && row["name"].bytesize.between?(1, 800) && row["name"].length <= 200 && !row["name"].match?(/\p{Cc}/) && valid_rate?(row["rate"])
+        return false if endpoints.key?(row["name"])
+        endpoints[row["name"]] = row["rate"].to_f
+      end
+      @policy = { id: id, issued_at: issued, expires_at: expires, deadline: @monotonic.call + [ expires - now, MAX_POLICY_TTL ].min,
+        enabled: value["enabled"], default_rate: rate.to_f, endpoints: endpoints }
+      true
+    rescue ArgumentError, EncodingError
+      false
+    end
+
+    def policy_due?
+      @activated && @monotonic.call >= @next_policy_at
+    end
+
+    def policy_requested!
+      @next_policy_at = @monotonic.call + POLICY_INTERVAL
+    end
+
+    def rotate(force: false, byte_budget:)
+      return unless supported?
+      now = @clock.call
+      while @reports.first && Time.iso8601(@reports.first.fetch("from")) < now - MAX_REPORT_AGE
+        @reports.shift
+        @stats[:reports_lost] += 1
+        @limited = true
+      end
+      return if @counts.empty? || (!force && @monotonic.call - @last_rotation < REPORT_INTERVAL)
+      if now < @from || now - @from > 86_400
+        @stats[:unreported] += @counts.values.sum { |entry| entry["observed"] }
+        @limited = true
+      else
+        rows = []
+        template = { "version" => 1, "limited" => @limited, "id" => SecureRandom.uuid, "from" => @from.iso8601(6), "to" => now.iso8601(6),
+          "enabled" => @config.automatic_sampling, "policy_id" => @policy_ids.size == 1 ? @policy_ids.first : "mixed", "endpoints" => rows }
+        @counts.each_value do |entry|
+          candidate = template.merge("endpoints" => rows + [ entry ])
+          if JSON.generate(candidate).bytesize > byte_budget
+            store_report(template.merge("endpoints" => rows)) if rows.any?
+            rows = []
+            template = template.merge("id" => SecureRandom.uuid)
+          end
+          if JSON.generate(template.merge("endpoints" => [ entry ])).bytesize <= byte_budget
+            rows << entry.dup
+          else
+            @stats[:unreported] += entry["observed"]
+            @limited = true
+          end
+        end
+        store_report(template.merge("endpoints" => rows)) if rows.any?
+      end
+      @counts, @policy_ids = {}, []
+      @from, @last_rotation = now, @monotonic.call
+    end
+
+    def report(force: false)
+      @reports.first if force || @monotonic.call >= @next_report_at
+    end
+
+    def report_sent!
+      @next_report_at = @monotonic.call + REPORT_INTERVAL
+    end
+
+    def acknowledge(id)
+      return unless id.is_a?(String) && id.match?(UUID)
+      removed = @reports.reject! { |report| report["id"] == id }
+      @next_report_at = 0 if removed && @reports.any?
+    end
+
+    def discard!
+      @stats[:reports_lost] += @reports.size
+      @stats[:unreported] += @counts.values.sum { |entry| entry["observed"] }
+      @reports, @counts = [], {}
+    end
+
+    def supported? = !@policy.nil?
+    def pending? = supported? && (@reports.any? || @counts.any?)
+    def stats = @stats.merge(limited: @limited, pending_reports: @reports.size, policy_state: effective_policy[1], policy_id: @policy&.fetch(:id))
+
+    private
+
+    def eligible?(event)
+      status, duration = event["status"], event["duration_ms"]
+      repeated = event.fetch("properties", {})["repeated_query_count"]
+      status.is_a?(String) && status.match?(/\A[23][0-9]{2}\z/) && duration.is_a?(Numeric) && duration.finite? && duration >= 0 && duration < 1000 && !(repeated.is_a?(Numeric) && repeated >= 3)
+    end
+
+    def effective_policy
+      return [ 1.0, "disabled", @policy&.fetch(:id) ] unless @config.automatic_sampling
+      return [ 1.0, "unavailable", nil ] unless @policy
+      return [ 1.0, "expired", @policy[:id] ] if @monotonic.call >= @policy[:deadline] || @clock.call >= @policy[:expires_at]
+      return [ 1.0, "disabled", @policy[:id] ] unless @policy[:enabled]
+      [ @policy[:default_rate], "active", @policy[:id] ]
+    end
+
+    def annotate(event, rate, state, policy_id)
+      event["properties"] = event.fetch("properties", {}).merge("rails_mind_sampling" => {
+        "version" => 1, "policy_id" => policy_id, "rate" => rate, "state" => state })
+    end
+
+    def store_report(report)
+      if @reports.size >= MAX_PENDING
+        @reports.shift
+        @stats[:reports_lost] += 1
+        @limited = true
+      end
+      report["limited"] ||= @limited
+      @reports << deep_freeze(JSON.parse(JSON.generate(report)))
+    end
+
+    def deep_freeze(value)
+      case value
+      when Hash then value.each { |key, child| deep_freeze(key); deep_freeze(child) }
+      when Array then value.each { |child| deep_freeze(child) }
+      end
+      value.freeze
+    end
+
+    def fraction(id) = Digest::SHA256.digest("rails_mind:sampling:v1:#{id.to_s.downcase}").unpack1("Q>") / (2.0**64)
+    def boolean?(value) = value == true || value == false
+    def valid_rate?(value) = (value.is_a?(Float) || value.is_a?(Integer)) && value.finite? && value.between?(0, 1)
+    def parse_time(value)
+      Time.iso8601(value) if value.is_a?(String) && value.bytesize <= 40
+    end
+  end
+end

data/lib/rails_mind/transport.rb

@@ -3,7 +3,8 @@ require "json"
 
 module RailsMind
   class Transport
-    Result = Struct.new(:success, :retryable, :retry_after, keyword_init: true)
+    MAX_RESPONSE_BYTES = 64 * 1024
+    Result = Struct.new(:success, :retryable, :retry_after, :sampling_policy, :sampling_report_ack, keyword_init: true)
     TRANSIENT_ERRORS = [ Timeout::Error, IOError, EOFError, SocketError, Errno::ECONNRESET,
       Errno::ECONNREFUSED, Errno::ETIMEDOUT, Errno::EPIPE, OpenSSL::SSL::SSLError ].freeze
 
@@ -15,28 +16,65 @@ module RailsMind
     end
 
     def send_batch(json)
+      request = Net::HTTP::Post.new(@uri)
+      request.body = json
+      perform(request)
+    end
+
+    def fetch_sampling_policy
+      uri = @uri.dup
+      uri.path = uri.path.sub(%r{/events$}, "/sampling_policy")
+      perform(Net::HTTP::Get.new(uri))
+    end
+
+    private
+
+    def perform(request)
       Context.suppress do
-        request = Net::HTTP::Post.new(@uri)
         request["Authorization"] = "Bearer #{@config.token}"
         request["Content-Type"] = "application/json"
         request["User-Agent"] = "rails_mind/#{VERSION}"
-        request.body = json
+        request["X-RailsMind-Automatic-Sampling"] = @config.automatic_sampling ? "enabled" : "disabled"
         http = @http_factory.call(@uri)
         http.use_ssl = @uri.scheme == "https"
         http.open_timeout = @config.open_timeout
         http.read_timeout = @config.read_timeout
         http.write_timeout = @config.read_timeout
         http.max_retries = 0
-        response = http.start { |connection| connection.request(request) }
+        body, streamed, oversized = +"", false, false
+        response = http.start do |connection|
+          connection.request(request) do |reply|
+            streamed = true
+            reply.read_body do |chunk|
+              if body.bytesize + chunk.bytesize > MAX_RESPONSE_BYTES
+                oversized = true
+              elsif !oversized
+                body << chunk
+              end
+            end
+          end
+        end
+        body = response.body.to_s if !streamed && response.respond_to?(:body)
         code = response.code.to_i
+        control = if !oversized && body.bytesize <= MAX_RESPONSE_BYTES && ((200..299).cover?(code) || code == 429)
+          parse_control(body)
+        else
+          {}
+        end
         Result.new(success: (200..299).cover?(code), retryable: code == 429 || code >= 500,
-          retry_after: retry_after(response["Retry-After"]))
+          retry_after: retry_after(response["Retry-After"]), sampling_policy: control["sampling_policy"],
+          sampling_report_ack: control["sampling_report_ack"])
       end
     rescue *TRANSIENT_ERRORS
       Result.new(success: false, retryable: true)
     end
 
-    private
+    def parse_control(body)
+      value = JSON.parse(body)
+      value.is_a?(Hash) ? value : {}
+    rescue JSON::ParserError, EncodingError
+      {}
+    end
 
     def retry_after(value)
       return nil unless value

data/lib/rails_mind/version.rb

@@ -1,3 +1,3 @@
 module RailsMind
-  VERSION = "0.1.0"
+  VERSION = "0.3.0"
 end

data/lib/rails_mind.rb

@@ -3,6 +3,7 @@ require_relative "rails_mind/configuration"
 require_relative "rails_mind/context"
 require_relative "rails_mind/redactor"
 require_relative "rails_mind/transport"
+require_relative "rails_mind/sampling"
 require_relative "rails_mind/collector"
 require_relative "rails_mind/client"