Memory middleware for Sidekiq

memory_sampling_middleware.rb

# lib/sidekiq/memory_sampling_middleware.rb
# frozen_string_literal: true
class Sidekiq::MemorySamplingMiddleware
  attr_accessor :logger

  def initialize(options = { logger: Sidekiq::Logging.logger })
    self.logger = options.fetch(:logger) { raise 'Missing logger parameter in options' }
  end

  CLASS_KEY = 'class'
  JID_KEY = 'jid'
  ARGS_KEY = 'args'

  # Process a sidekiq job and send to the configured logger:
  #   * the queue
  #   * the worker class
  #   * the job id
  #   * the job parameters
  #   * the job execution start time
  #   * the job execution end time
  #   * the job execution duration in milliseconds
  #   * the memory usage estimation of the process BEFORE the job was executed (initial)
  #   * the memory usage estimation of the process AFTER the job was executed (final)
  #   * the memory usage increase/decrease between before and after the job was executed (delta)
  #
  # Expected parameters are:
  #   * worker: a sidekiq worker instance. Unused for the moment.
  #   * job: the job hash. Must contains the following keys: class, jid, args
  #   * queue: the queue name
  def call(_worker, job, queue)
    initial = current_memory
    start_time = Time.now
    yield
  ensure
    end_time = Time.now
    final = current_memory
    logger.info('SidekiqMemorySampling ' + {
      queue:        queue,
      worker:       job.fetch(CLASS_KEY),
      jid:          job.fetch(JID_KEY),
      args:         job.fetch(ARGS_KEY),
      started_at:   start_time.iso8601,
      ended_at:     end_time.iso8601,
      elapsed_ms:   (end_time - start_time) * 1000,
      initial_mem:  initial,
      final_mem:    final,
      delta_mem:    delta(initial, final)
    }.map { |k, v| "#{k}=#{truncate(v)}" }.join(' '))
  end

  private

  NONE = :no_memory_sampling

  def delta(initial, final)
    return NONE if initial == NONE || final == NONE
    final - initial
  end

  def current_memory
    # From NewRelic Agent
    MemorySampler.get_memory
  rescue StandardError => e
    logger.error "Memory sampling failed: #{e.message}"
    Airbrake.notify e
    NONE
  end

  def truncate(value, length = 40, ellipsis = '... (continued)')
    if value.is_a?(Array)
      value.map { |v| truncate(v) }
    elsif value.is_a?(String)
      value.truncate(length, omission: ellipsis)
    else
      value
    end
  end
end