rails_observatory 0.1.0

data/lib/rails_observatory/models/error.rb

@@ -0,0 +1,67 @@
+require_relative './redis_model'
+require 'digest'
+module RailsObservatory
+  class Error < RedisModel
+
+    attribute :time, :float
+    attribute :fingerprint, :string
+    attribute :has_causes, :boolean, indexed: false
+    attribute :causes, compressed: true, indexed: false
+    attribute :location, :string
+    attribute :class_name, :string
+    attribute :message, :string
+    attribute :source_extracts, compressed: true, indexed: false
+    attribute :trace, compressed: true, indexed: false
+
+    alias_attribute :id, :fingerprint
+
+    attr_accessor :exception
+
+    def exception=(ex)
+      ex_wrapper = ActionDispatch::ExceptionWrapper.new(Rails.backtrace_cleaner, ex)
+      payload = payload_for_wrapped_exception(ex_wrapper)
+      assign_attributes(payload)
+      self.has_causes = ex_wrapper.has_cause?
+      self.causes = ex_wrapper.wrapped_causes.map { payload_for_wrapped_exception(_1) }
+      self.fingerprint = build_fingerprint(ex_wrapper)
+    end
+
+    private
+
+    def trace_for_ex(wrapped_ex)
+      wrapped_ex.full_trace.each_with_index.map do |trace, idx|
+        {
+          exception_object_id: wrapped_ex.exception.object_id,
+          id: idx,
+          trace: trace,
+          is_application_trace: wrapped_ex.application_trace.include?(trace),
+        }
+      end
+    end
+
+    def payload_for_wrapped_exception(wrapped_ex)
+      {
+        class_name: wrapped_ex.exception_class_name,
+        message: wrapped_ex.message,
+        source_extracts: wrapped_ex.source_extracts,
+        trace: trace_for_ex(wrapped_ex),
+      }
+    end
+
+    def contextual_request_id
+      context = ActiveSupport::ExecutionContext.to_h
+      context[:controller]&.request&.request_id || context[:job]&.request_id
+    end
+
+    def exception_string(wrapped_ex)
+      wrapped_ex.exception_class_name + wrapped_ex.exception.backtrace.map { _1.split(":").slice(0..1).join(":") }.join("\n")
+    end
+
+    def build_fingerprint(wrapped_ex)
+      exceptions = [wrapped_ex]
+      exceptions.push(*wrapped_ex.wrapped_causes) if wrapped_ex.has_cause?
+      Digest::SHA256.hexdigest(exceptions.map { exception_string(_1) }.join("\n"))
+    end
+
+  end
+end

data/lib/rails_observatory/models/event_collection.rb

@@ -0,0 +1,137 @@
+module RailsObservatory
+
+  class EventCollection
+    include Enumerable
+
+    attr_accessor :events
+
+    delegate :push, :<<, :size, to: :events
+    delegate :empty?, to: :to_a
+
+    def initialize(events)
+      @events = events
+    end
+
+    def without(*names)
+      copy = self.clone
+      copy.instance_exec { @without = names }
+      copy
+    end
+
+    def only(*names)
+      copy = self.clone
+      copy.instance_exec { @only = names }
+      copy
+    end
+
+    def each
+      decorate_events unless @decorated
+      iterating_set = @events
+      if @without
+        iterating_set = iterating_set.reject { _1['name'].in?(@without) }
+      end
+      if @only
+        iterating_set = iterating_set.select { _1['name'].in?(@only) }
+      end
+      iterating_set.then(&method(:decorate_with_relative_time)).each { yield _1 }
+    end
+
+    def to_series
+      all_events = to_a
+      min, max = all_events.minmax_by { _1['depth'] }.pluck('depth')
+
+      category_primer = (min..max).map do |depth|
+        {
+          x: depth.to_s,
+          y: nil,
+          event_self_time: 0,
+        }
+      end
+
+      grouped_events = all_events.group_by { _1['name'].split('.').last }.sort_by { _1.first }
+      grouped_events.map do |name, events|
+        {
+          name: name,
+          data: category_primer + events.map do |ev|
+            {
+              x: ev['depth'].to_s,
+              y: [ev['relative_start_at'], ev['relative_end_at']],
+              event_self_time: ev['self_time'],
+              event_name: ev['name'].split('.').first,
+              start_at: ev['start_at'],
+            }
+          end
+        }
+      end
+    end
+
+
+
+    private
+
+    def decorate_events
+      @events = @events.then(&method(:sort_events))
+                       .then(&method(:decorate_with_self_time))
+                       .then(&method(:merge_middleware))
+                       .then(&method(:decorate_with_depth))
+      @decorated = true
+    end
+
+    def sort_events(events)
+      events.sort_by { _1['start_at'] }
+    end
+
+    def merge_middleware(events)
+      middleware_events = events.select { _1['name'] == 'process_middleware.action_dispatch' }
+      return events if middleware_events.empty?
+      merged_middleware = middleware_events.reduce(middleware_events.first.without('self_time')) do |merged, event|
+        merged['self_time'] ||= 0
+        merged['self_time'] += event['self_time']
+        merged['middleware_stack'] ||= []
+        merged['middleware_stack'] << event
+        merged
+      end
+      [merged_middleware] + events.excluding(middleware_events)
+    end
+
+    def decorate_with_depth(events)
+      depth_stack = [] 
+      events.each do |e|
+        event_range = (e['start_at']..e['end_at'])
+        depth_stack.select! { _1.cover?(event_range) }
+        e['depth'] = depth_stack.size
+        depth_stack << event_range
+      end
+    end
+
+    def decorate_with_self_time(events)
+      events.each do |ev|
+        ev_range = (ev['start_at']..ev['end_at'])
+        sub_events = events.excluding(ev).select { ev_range.cover?(_1['start_at'].._1['end_at']) }
+        sub_event_time = non_overlapping_ranges(sub_events).reduce(0) { |sum, range| sum + (range.end - range.begin) }
+        ev['self_time'] = ev['duration'] - sub_event_time
+      end
+    end
+
+    def non_overlapping_ranges(events)
+      events.reduce([]) do |arr, event|
+        event_range = (event['start_at']..event['end_at'])
+        if arr.any? { |r| r.cover?(event_range) }
+          arr
+        else
+          arr << event_range
+        end
+      end
+    end
+
+    def decorate_with_relative_time(events)
+      return events if events.empty?
+      first_event = events.first['start_at']
+      events.each do |ev|
+        ev['relative_start_at'] = ev['start_at'] - first_event
+        ev['relative_end_at'] = ev['end_at'] - first_event
+      end
+    end
+
+  end
+end

data/lib/rails_observatory/models/events.rb

@@ -0,0 +1,22 @@
+module RailsObservatory
+  module Events
+    extend ActiveSupport::Concern
+
+    included do
+      attribute :events, indexed: false, compressed: true
+
+      def events
+        attr_value = super
+        return nil if attr_value.nil?
+        EventCollection.new(attr_value)
+      end
+    end
+
+
+
+    def mail_events
+      events.only('enqueue.action_job', 'deliver.action_mailer')
+            .reject { _1['name'] == 'enqueue.action_job' && _1.dig('payload', 'job', 'class') != 'ActionMailer::MailDeliveryJob' }
+    end
+  end
+end

data/lib/rails_observatory/models/job_trace.rb

@@ -0,0 +1,28 @@
+require_relative './redis_model'
+require_relative './events'
+require_relative './logs'
+
+module RailsObservatory
+  class JobTrace < RedisModel
+    include Events
+    include Logs
+
+    def self.key_prefix
+      "jt"
+    end
+
+    attribute :job_id, :string
+    attribute :queue_name, :string
+    attribute :queue_adapter, :string
+    attribute :job_class, :string
+    attribute :executions, :integer
+    attribute :error, :boolean
+    attribute :time, :float
+    attribute :allocations, :integer, indexed: false
+    attribute :queue_latency, :float, indexed: false
+    attribute :duration, :float, indexed: false
+
+    alias_attribute :id, :job_id
+    alias_attribute :name, :job_class
+  end
+end

data/lib/rails_observatory/models/logs.rb

@@ -0,0 +1,9 @@
+module RailsObservatory
+  module Logs
+    extend ActiveSupport::Concern
+
+    included do
+      attribute :logs, indexed: false, compressed: true
+    end
+  end
+end

data/lib/rails_observatory/models/mail_delivery.rb

@@ -0,0 +1,33 @@
+require_relative './redis_model'
+
+module RailsObservatory
+  class MailDelivery < RedisModel
+    attribute :message_id, :string
+    attribute :time, :float
+    attribute :duration, :float
+    attribute :mailer, :string
+    attribute :to, :string
+    attribute :from, :string
+    attribute :subject, :string
+    attribute :mail, compressed: true, indexed: false
+
+    alias_attribute :id, :message_id
+
+    def to=(val)
+      if val.is_a?(Array)
+        super(val.join(', '))
+      else
+        super
+      end
+    end
+
+    def from=(val)
+      if val.is_a?(Array)
+        super(val.join(', '))
+      else
+        super
+      end
+
+    end
+  end
+end

data/lib/rails_observatory/models/redis_model.rb

@@ -0,0 +1,112 @@
+require_relative '../serializers/serializer'
+require_relative './event_collection'
+require_relative '../redis/time_series'
+require 'zlib'
+module RailsObservatory
+  class RedisModel
+    include ActiveModel::Model
+    include ActiveModel::Attributes
+    include ActiveModel::Serializers::JSON
+
+    class NotFound < StandardError; end
+
+    class << self
+      attr_accessor :indexed_attributes
+      attr_accessor :compressed_attributes
+    end
+
+    def self.attribute(name, *args, indexed: true, compressed: false, **rest)
+      if indexed
+        self.indexed_attributes ||= []
+        indexed_attributes << name
+      end
+      if compressed
+        self.compressed_attributes ||= []
+        compressed_attributes << name.to_s
+      end
+      super(name, *args, **rest)
+    end
+
+    def self.redis
+      Rails.configuration.rails_observatory.redis
+    end
+
+    def redis
+      self.class.redis
+    end
+
+    def self.key_prefix
+      name.demodulize.underscore
+    end
+
+    def self.key_name(id)
+      "#{key_prefix}:#{id}"
+    end
+
+    def self.count
+      total, *results = redis.call("FT.SEARCH", index_name, '*', "SORTBY", "time", "DESC")
+      total
+    end
+
+    def self.all
+      total, *results = redis.call("FT.SEARCH", index_name, '*', "SORTBY", "time", "DESC")
+      Hash[*results].values.map(&:last).map { JSON.parse(_1) }.map { new(_1) }
+    end
+
+    def self.find(id)
+      result = redis.call("JSON.GET", key_name(id), "$") || raise(NotFound, "Could not find #{name} with id #{id}")
+      attrs = JSON.parse(result).first
+
+      compressed_attributes.each do |attr|
+        val = redis.call("GET", [key_prefix, attr].join("_")  + ":#{id}")
+        attrs.merge!(attr => JSON.parse(Zlib.gunzip(val)))
+      end
+
+      self.new(attrs)
+    end
+
+    ATTRIBUTE_TYPE_TO_REDIS_TYPE = {
+      string: "TEXT",
+      integer: "NUMERIC",
+      float: "NUMERIC",
+      boolean: "TAG"
+    }
+
+    def self.index_name
+      "#{key_prefix}-idx"
+    end
+
+    def self.create_redis_index
+      schema = indexed_attributes.flat_map do |attr|
+        ["$.#{attr}", "AS", "#{attr}", ATTRIBUTE_TYPE_TO_REDIS_TYPE[attribute_types[attr.to_s].type]]
+      end
+      redis.call("FT.CREATE", index_name, "ON", "JSON", "PREFIX", "1", key_prefix, "SCHEMA", *schema)
+    end
+
+    def self.index_info
+      info = Hash[*redis.call("FT.INFO", index_name)]
+      info['attributes'] = info['attributes'].map { Hash[*_1] }
+      info['index_definition'] = Hash[*info['index_definition']]
+      info
+    end
+
+    def self.ensure_index
+      redis.call("FT._LIST").include?(index_name) || create_redis_index
+    end
+
+    def attribute_names_for_serialization
+      attributes.keys - self.class.compressed_attributes
+    end
+
+    def save
+      redis.multi do |r|
+        r.call("JSON.SET", self.class.key_name(id), "$", JSON.generate(as_json))
+        self.class.compressed_attributes.each do |attr|
+          compressed_value = Zlib.gzip(JSON.generate(@attributes.fetch_value(attr)), level: Zlib::BEST_COMPRESSION)
+          r.call("SET", [self.class.key_prefix, attr].join("_")  + ":#{id}", compressed_value)
+        end
+      end
+    end
+
+  end
+end

data/lib/rails_observatory/models/request_trace.rb

@@ -0,0 +1,29 @@
+require_relative './redis_model'
+require_relative './events'
+require_relative './logs'
+module RailsObservatory
+  class RequestTrace < RedisModel
+    include Events
+    include Logs
+
+    attribute :request_id, :string
+    attribute :status, :integer
+    attribute :http_method, :string
+    attribute :path, :string
+    attribute :action, :string
+    attribute :format, :string
+    attribute :error, :boolean
+    attribute :route_pattern, :string
+    attribute :time, :float
+    attribute :duration, :float
+    attribute :allocations, :integer, indexed: false
+
+    alias_attribute :id, :request_id
+    alias_attribute :name, :action
+
+    def self.key_prefix
+      "rt"
+    end
+
+  end
+end

data/lib/rails_observatory/railties/active_job_instrumentation.rb

@@ -0,0 +1,48 @@
+require_relative '../event_collector'
+require_relative '../models/job_trace'
+module RailsObservatory
+  module Railties
+    module ActiveJobInstrumentation
+
+      def perform_now
+        TimeSeries.distribution("job.queue_latency", Time.now - enqueued_at, labels: { queue_name: }) unless enqueued_at.nil?
+        labels = { job_class: self.class.name, queue_name: }
+        TimeSeries.increment("job.count", labels:)
+        TimeSeries.increment("job.retry_count", labels:) if executions > 1
+
+        start_at = Time.now
+        start_at_mono = Process.clock_gettime(Process::CLOCK_MONOTONIC, :float_millisecond)
+        result = nil
+        logs = []
+        events = EventCollector.instance.collect_events do
+          logs = LogCollector.collect_logs do
+            result = super
+          end
+        end
+        end_at_mono = Process.clock_gettime(Process::CLOCK_MONOTONIC, :float_millisecond)
+        result
+      rescue Exception => error
+        events = error.instance_variable_get(:@_trace_events)
+        end_at_mono = Process.clock_gettime(Process::CLOCK_MONOTONIC, :float_millisecond)
+        TimeSeries.increment("job.error_count", labels:)
+        raise
+      ensure
+        duration = end_at_mono - start_at_mono
+        TimeSeries.distribution("job.latency", duration, labels:)
+        JobTrace.new(
+          job_id: job_id,
+          time: start_at.to_f,
+          duration:,
+          queue_adapter: ActiveJob.adapter_name(queue_adapter),
+          executions:,
+          job_class: self.class.name,
+          queue_name:,
+          events: events.map { Serializer.serialize(_1) },
+          logs:,
+          error: error.present?
+        ).save
+      end
+
+    end
+  end
+end

data/lib/rails_observatory/railties/redis_runtime.rb

@@ -0,0 +1,11 @@
+module RailsObservatory
+  module Railties
+    module RedisRuntime
+      def append_info_to_payload(payload)
+        super
+
+        payload[:rails_observatory_runtime] = ActiveSupport::ExecutionContext.to_h[:rails_observatory_redis]
+      end
+    end
+  end
+end

data/lib/rails_observatory/redis/logging_middleware.rb

@@ -0,0 +1,22 @@
+module RailsObservatory
+  module LoggingMiddleware
+    def connect(redis_config)
+      puts "[Redis] CONNECT"
+      super
+    end
+
+    def call(command, redis_config)
+      if command.first == "SCRIPT"
+        puts "[Redis] #{command.first}"
+      else
+        puts "[Redis] #{command.first} #{command[1..-1].join(" ")}"
+      end
+      super
+    end
+
+    def call_pipelined(commands, redis_config)
+      puts "[Redis] [Pipelined] #{commands.join("\n")}"
+      super
+    end
+  end
+end

data/lib/rails_observatory/redis/redis_client_instrumentation.rb

@@ -0,0 +1,18 @@
+module RailsObservatory
+  module RedisClientInstrumentation
+    def call(command, redis_config)
+      payload_command = command
+      payload_command = [payload_command.first] if payload_command.first == "SCRIPT"
+      ActiveSupport::Notifications.instrument("call.redis", { command: payload_command.join(' ') }) do
+        super
+      end
+    end
+
+    def call_pipelined(commands, redis_config)
+      res = nil
+      took = Benchmark.realtime { res = super }
+      # puts "Redis call_pipelined took #{took * 1000} ms"
+      res
+    end
+  end
+end

data/lib/rails_observatory/redis/time_series/increment_script.lua

@@ -0,0 +1,67 @@
+-- Helper function to get all combinations of a table
+local function generate_key_combinations(keys)
+  local n = #keys
+  local combs = {}
+  table.insert(combs, {})
+
+  local function helper(curr_comb, start_idx)
+    if start_idx <= n then
+      for i = start_idx, n do
+        local new_comb = {}
+        for _, v in ipairs(curr_comb) do
+          table.insert(new_comb, v)
+        end
+        table.insert(new_comb, keys[i])
+        table.insert(combs, new_comb)
+        helper(new_comb, i + 1)
+      end
+    end
+  end
+
+  helper({}, 1)
+  return combs
+end
+
+-- Main script begins here
+local metric_name = tostring(ARGV[1])  -- Ensure it's a string
+local raw_retention = 10000 -- Hardcoded to 10ms
+local compaction_retention = 31536000000 -- Hardcoded to 1 year in ms (365*24*60*60*1000)
+
+-- Extracting labels
+---@type table
+local labels = {}
+local keys = {}
+for i = 2, #ARGV, 2 do
+  local key = tostring(ARGV[i])
+  local value = tostring(ARGV[i + 1])
+  labels[key] = value
+  redis.call("SADD", metric_name .. ':labels', key)
+  table.insert(keys, key)
+end
+
+local key_combinations = generate_key_combinations(keys)
+
+-- For each combination, upsert and add labels
+for _, comb_keys in ipairs(key_combinations) do
+  local ts_name = metric_name
+  local label_set = {}
+
+  for _, key in ipairs(comb_keys) do
+    ts_name = ts_name .. ":" .. labels[key]
+    table.insert(label_set, key)
+    table.insert(label_set, labels[key])
+  end
+
+  if redis.call("EXISTS", ts_name) == 0 then
+    redis.call("TS.CREATE", ts_name, "RETENTION", raw_retention, "CHUNK_SIZE", 48)
+  end
+
+  local compaction_key = ts_name .. "_" .. "sum"
+  if redis.call("EXISTS", compaction_key) == 0 then
+    redis.call("TS.CREATE", compaction_key, "RETENTION", compaction_retention, "CHUNK_SIZE", 48, "LABELS", "name", metric_name, "compaction", "sum", unpack(label_set))
+    redis.call("TS.CREATERULE", ts_name, compaction_key, "AGGREGATION", "sum", 10000)
+  end
+  redis.call("TS.ADD", ts_name, "*", 1, 'ON_DUPLICATE', 'SUM')
+end
+
+return "OK"

data/lib/rails_observatory/redis/time_series/insertion.rb

@@ -0,0 +1,73 @@
+require 'benchmark'
+
+SCRIPT = File.read(File.join(File.dirname(__FILE__), 'timing_script.lua'))
+INCREMENT_SCRIPT = File.read(File.join(File.dirname(__FILE__), 'increment_script.lua'))
+
+class RedisScript
+
+  def self.redis
+    Rails.configuration.rails_observatory.redis
+  end
+
+  def redis
+    self.class.redis
+  end
+
+  def initialize(lua_string)
+    @script = lua_string
+  end
+
+  def call(*args)
+    @sha1 ||= load_script
+    redis.call("EVALSHA", @sha1, 0, *args)
+  rescue => e
+    if e.message =~ /NOSCRIPT/
+      @sha1 = load_script
+      retry
+    else
+      raise e
+    end
+  end
+
+  def load_script
+    redis.call('SCRIPT', 'LOAD', @script)
+  end
+
+end
+
+TIMING_SCRIPT = RedisScript.new(SCRIPT)
+INCREMENT_CALL = RedisScript.new(INCREMENT_SCRIPT)
+
+module RailsObservatory
+
+  class TimeSeries
+    module Insertion
+
+      # TODO: These need to take in a timestamp
+      def distribution(name, value, labels: {})
+        prefixed_name = begin
+                          if defined?(self::PREFIX)
+                            [self::PREFIX, name].join('.')
+                          else
+                            name
+                          end
+                        end
+        TIMING_SCRIPT.call(prefixed_name, value, labels.to_a.flatten.map(&:to_s))
+      end
+      alias_method :record_timing, :distribution
+
+      def increment(name, labels: {})
+
+        prefixed_name = begin
+                          if defined?(self::PREFIX)
+                            [self::PREFIX, name].join('.')
+                          else
+                            name
+                          end
+                        end
+        INCREMENT_CALL.call(prefixed_name, labels.to_a.flatten.map(&:to_s))
+      end
+      alias_method :record_occurrence, :increment
+    end
+  end
+end