memory_profiler.rb in single file

memory_profiler.rb

# frozen_string_literal: true

# Single file copy of https://github.com/SamSaffron/memory_profiler  for loading in Sonic Pi.
# Licence: https://github.com/SamSaffron/memory_profiler/blob/master/LICENSE.txt
# Copyright (c) 2013 Sam Saffron
# MIT License

module MemoryProfiler
  VERSION = "0.9.14"
end
module MemoryProfiler
  class Helpers

    def initialize
      @gem_guess_cache = Hash.new
      @location_cache = Hash.new { |h, k| h[k] = Hash.new.compare_by_identity }
      @class_name_cache = Hash.new.compare_by_identity
      @string_cache = Hash.new
    end

    def guess_gem(path)
      @gem_guess_cache[path] ||=
        if /(\/gems\/.*)*\/gems\/(?<gemname>[^\/]+)/ =~ path
          gemname
        elsif /\/rubygems[\.\/]/ =~ path
          "rubygems"
        elsif /ruby\/2\.[^\/]+\/(?<stdlib>[^\/\.]+)/ =~ path
          stdlib
        elsif /(?<app>[^\/]+\/(bin|app|lib))/ =~ path
          app
        else
          "other"
        end
    end

    def lookup_location(file, line)
      @location_cache[file][line] ||= "#{file}:#{line}"
    end

    def lookup_class_name(klass)
      @class_name_cache[klass] ||= ((klass.is_a?(Class) && klass.name) || '<<Unknown>>').to_s
    end

    def lookup_string(obj)
      # This string is shortened to 200 characters which is what the string report shows
      # The string report can still list unique strings longer than 200 characters
      #   separately because the object_id of the shortened string will be different
      @string_cache[obj] ||= String.new << obj[0, 200]
    end
  end
end

module MemoryProfiler

  class Polychrome

    def path(text)
      blue(text)
    end

    def string(text)
      green(text)
    end

    def line(text)
      cyan(text)
    end

    private

    def black(str)
      "\033[30m#{str}\033[0m"
    end

    def red(str)
      "\033[31m#{str}\033[0m"
    end

    def green(str)
      "\033[32m#{str}\033[0m"
    end

    def brown(str)
      "\033[33m#{str}\033[0m"
    end

    def blue(str)
      "\033[34m#{str}\033[0m"
    end

    def magenta(str)
      "\033[35m#{str}\033[0m"
    end

    def cyan(str)
      "\033[36m#{str}\033[0m"
    end

    def gray(str)
      "\033[37m#{str}\033[0m"
    end

    def bg_black(str)
      "\033[40m#{str}\033[0m"
    end

    def bg_red(str)
      "\033[41m#{str}\033[0m"
    end

    def bg_green(str)
      "\033[42m#{str}\033[0m"
    end

    def bg_brown(str)
      "\033[43m#{str}\033[0m"
    end

    def bg_blue(str)
      "\033[44m#{str}\033[0m"
    end

    def bg_magenta(str)
      "\033[45m#{str}\033[0m"
    end

    def bg_cyan(str)
      "\033[46m#{str}\033[0m"
    end

    def bg_gray(str)
      "\033[47m#{str}\033[0m"
    end

    def bold(str)
      "\033[1m#{str}\033[22m"
    end

    def reverse_color(str)
      "\033[7m#{str}\033[27m"
    end

  end

end

module MemoryProfiler

  class Monochrome

    def path(text)
      text
    end

    def string(text)
      text
    end

    def line(text)
      text
    end
  end

end

module MemoryProfiler
  module TopN
    # Fast approach for determining the top_n entries in a Hash of Stat objects.
    # Returns results for both memory (memsize summed) and objects allocated (count) as a tuple.
    def top_n(max, metric_method)

      stat_totals =
        self.values
          .group_by(&metric_method)
          .map do |metric, stats|
            [metric, stats.reduce(0) { |sum, stat| sum + stat.memsize }, stats.size]
          end

      stats_by_memsize =
        stat_totals
          .sort_by! { |metric, memsize, _count| [-memsize, metric] }
          .take(max)
          .map! { |metric, memsize, _count| { data: metric, count: memsize } }

      stats_by_count =
        stat_totals
          .sort_by! { |metric, _memsize, count| [-count, metric] }
          .take(max)
          .map! { |metric, _memsize, count| { data: metric, count: count } }

      [stats_by_memsize, stats_by_count]
    end
  end
end

module MemoryProfiler
  class Stat

    attr_reader :class_name, :gem, :file, :location, :memsize, :string_value

    def initialize(class_name, gem, file, location, memsize, string_value)
      @class_name = class_name
      @gem = gem

      @file = file
      @location = location

      @memsize = memsize
      @string_value = string_value
    end

  end
end

module MemoryProfiler
  class StatHash < Hash
    include TopN
  end
end

module MemoryProfiler
  class Results
    UNIT_PREFIXES = {
      0 => 'B',
      3 => 'kB',
      6 => 'MB',
      9 => 'GB',
      12 => 'TB',
      15 => 'PB',
      18 => 'EB',
      21 => 'ZB',
      24 => 'YB'
    }.freeze

    TYPES   = ["allocated", "retained"].freeze
    METRICS = ["memory", "objects"].freeze
    NAMES   = ["gem", "file", "location", "class"].freeze

    def self.register_type(name, stat_attribute)
      @@lookups ||= []
      @@lookups << [name, stat_attribute]

      TYPES.each do |type|
        METRICS.each do |metric|
          attr_accessor "#{type}_#{metric}_by_#{name}"
        end
      end
    end

    register_type 'gem', :gem
    register_type 'file', :file
    register_type 'location', :location
    register_type 'class', :class_name

    attr_accessor :strings_retained, :strings_allocated
    attr_accessor :total_retained, :total_allocated
    attr_accessor :total_retained_memsize, :total_allocated_memsize

    def register_results(allocated, retained, top)

      @@lookups.each do |name, stat_attribute|

        memsize_results, count_results = allocated.top_n(top, stat_attribute)

        self.send("allocated_memory_by_#{name}=", memsize_results)
        self.send("allocated_objects_by_#{name}=", count_results)

        memsize_results, count_results = retained.top_n(top, stat_attribute)

        self.send("retained_memory_by_#{name}=", memsize_results)
        self.send("retained_objects_by_#{name}=", count_results)
      end

      self.strings_allocated = string_report(allocated, top)
      self.strings_retained = string_report(retained, top)

      self.total_allocated = allocated.size
      self.total_allocated_memsize = allocated.values.map!(&:memsize).inject(0, :+)
      self.total_retained = retained.size
      self.total_retained_memsize = retained.values.map!(&:memsize).inject(0, :+)

      self
    end

    def scale_bytes(bytes)
      return "0 B" if bytes.zero?

      scale = Math.log10(bytes).div(3) * 3
      scale = 24 if scale > 24
      "%.2f #{UNIT_PREFIXES[scale]}" % (bytes / 10.0**scale)
    end

    def string_report(data, top)
      grouped_strings = data.values
        .keep_if  { |stat| stat.string_value }
        .group_by { |stat| stat.string_value.object_id }
        .values

      if grouped_strings.size > top
        cutoff = grouped_strings.sort_by!(&:size)[-top].size
        grouped_strings.keep_if { |list| list.size >= cutoff }
      end

      grouped_strings
        .sort! { |a, b| a.size == b.size ? a[0].string_value <=> b[0].string_value : b.size <=> a.size }
        .first(top)
        .map! do |list|
          # Return array of [string, [[location, count], [location, count], ...]
          [
            list[0].string_value,
            list.group_by { |stat| stat.location }
              .map { |location, stat_list| [location, stat_list.size] }
              .sort_by!(&:last)
              .reverse!
          ]
        end
    end

    # Output the results of the report
    # @param [Hash] options the options for output
    # @option opts [String] :to_file a path to your log file
    # @option opts [Boolean] :color_output a flag for whether to colorize output
    # @option opts [Integer] :retained_strings how many retained strings to print
    # @option opts [Integer] :allocated_strings how many allocated strings to print
    # @option opts [Boolean] :detailed_report should report include detailed information
    # @option opts [Boolean] :scale_bytes calculates unit prefixes for the numbers of bytes
    # @option opts [Boolean] :normalize_paths print location paths relative to gem's source directory.
    def pretty_print(io = $stdout, **options)
      # Handle the special case that Ruby PrettyPrint expects `pretty_print`
      # to be a customized pretty printing function for a class
      return io.pp_object(self) if defined?(PP) && io.is_a?(PP)

      io = File.open(options[:to_file], "w") if options[:to_file]

      color_output = options.fetch(:color_output) { io.respond_to?(:isatty) && io.isatty }
      @colorize = color_output ? Polychrome.new : Monochrome.new

      if options[:scale_bytes]
        total_allocated_output = scale_bytes(total_allocated_memsize)
        total_retained_output  = scale_bytes(total_retained_memsize)
      else
        total_allocated_output = "#{total_allocated_memsize} bytes"
        total_retained_output  = "#{total_retained_memsize} bytes"
      end

      io.puts "Total allocated: #{total_allocated_output} (#{total_allocated} objects)"
      io.puts "Total retained:  #{total_retained_output} (#{total_retained} objects)"

      unless options[:detailed_report] == false
        TYPES.each do |type|
          METRICS.each do |metric|
            NAMES.each do |name|
              dump_data(io, type, metric, name, options)
            end
          end
        end
      end

      io.puts
      print_string_reports(io, options)

      io.close if io.is_a? File
    end

    def print_string_reports(io, options)
      TYPES.each do |type|
        dump_opts = {
          normalize_paths: options[:normalize_paths],
          limit: options["#{type}_strings".to_sym]
        }
        dump_strings(io, type, dump_opts)
      end
    end

    def normalize_path(path)
      @normalize_path ||= {}
      @normalize_path[path] ||= begin
        if %r!(/gems/.*)*/gems/(?<gemname>[^/]+)(?<rest>.*)! =~ path
          "#{gemname}#{rest}"
        elsif %r!ruby/2\.[^/]+/(?<stdlib>[^/.]+)(?<rest>.*)! =~ path
          "ruby/lib/#{stdlib}#{rest}"
        elsif %r!(?<app>[^/]+/(bin|app|lib))(?<rest>.*)! =~ path
          "#{app}#{rest}"
        else
          path
        end
      end
    end

    private

    def print_title(io, title)
      io.puts
      io.puts title
      io.puts @colorize.line("-----------------------------------")
    end

    def print_output(io, topic, detail)
      io.puts "#{@colorize.path(topic.to_s.rjust(10))}  #{detail}"
    end

    def dump_data(io, type, metric, name, options)
      print_title  io, "#{type} #{metric} by #{name}"
      data = self.send "#{type}_#{metric}_by_#{name}"

      scale_data = metric == "memory" && options[:scale_bytes]
      normalize_paths = options[:normalize_paths]

      if data && !data.empty?
        data.each do |item|
          count = scale_data ? scale_bytes(item[:count]) : item[:count]
          value = normalize_paths ? normalize_path(item[:data]) : item[:data]
          print_output io, count, value
        end
      else
        io.puts "NO DATA"
      end

      nil
    end

    def dump_strings(io, type, options)
      strings = self.send("strings_#{type}") || []
      return if strings.empty?

      options = {} unless options.is_a?(Hash)

      if (limit = options[:limit])
        return if limit == 0
        strings = strings[0...limit]
      end

      normalize_paths = options[:normalize_paths]

      print_title(io, "#{type.capitalize} String Report")
      strings.each do |string, stats|
        print_output io, (stats.reduce(0) { |a, b| a + b[1] }), @colorize.string(string.inspect)
        stats.sort_by { |x, y| [-y, x] }.each do |location, count|
          location = normalize_path(location) if normalize_paths
          print_output io, count, location
        end
        io.puts
      end

      nil
    end

  end

end

require 'objspace'
module MemoryProfiler
  # Reporter is the top level API used for generating memory reports.
  #
  # @example Measure object allocation in a block
  #   report = Reporter.report(top: 50) do
  #     5.times { "foo" }
  #   end
  class Reporter
    class << self
      attr_accessor :current_reporter
    end

    attr_reader :top, :trace, :generation, :report_results

    def initialize(opts = {})
      @top          = opts[:top] || 50
      @trace        = opts[:trace] && Array(opts[:trace])
      @ignore_files = opts[:ignore_files] && Regexp.new(opts[:ignore_files])
      @allow_files  = opts[:allow_files] && /#{Array(opts[:allow_files]).join('|')}/
    end

    # Helper for generating new reporter and running against block.
    # @param [Hash] opts the options to create a report with
    # @option opts :top max number of entries to output
    # @option opts :trace a class or an array of classes you explicitly want to trace
    # @option opts :ignore_files a regular expression used to exclude certain files from tracing
    # @option opts :allow_files a string or array of strings to selectively include in tracing
    # @return [MemoryProfiler::Results]
    def self.report(opts = {}, &block)
      self.new(opts).run(&block)
    end

    def start
      GC.start
      GC.start
      GC.start
      GC.disable

      @generation = GC.count
      ObjectSpace.trace_object_allocations_start
    end

    def stop
      ObjectSpace.trace_object_allocations_stop
      allocated = object_list(generation)
      retained = StatHash.new.compare_by_identity

      GC.enable
      GC.start
      GC.start
      GC.start

      # Caution: Do not allocate any new Objects between the call to GC.start and the completion of the retained
      #          lookups. It is likely that a new Object would reuse an object_id from a GC'd object.

      ObjectSpace.each_object do |obj|
        next unless ObjectSpace.allocation_generation(obj) == generation
        found = allocated[obj.__id__]
        retained[obj.__id__] = found if found
      end
      ObjectSpace.trace_object_allocations_clear

      @report_results = Results.new
      @report_results.register_results(allocated, retained, top)
    end

    # Collects object allocation and memory of ruby code inside of passed block.
    def run(&block)
      start
      begin
        yield
      rescue Exception
        ObjectSpace.trace_object_allocations_stop
        GC.enable
        raise
      else
        stop
      end
    end

    private

    # Iterates through objects in memory of a given generation.
    # Stores results along with meta data of objects collected.
    def object_list(generation)

      rvalue_size = GC::INTERNAL_CONSTANTS[:RVALUE_SIZE]
      helper = Helpers.new

      result = StatHash.new.compare_by_identity

      ObjectSpace.each_object do |obj|
        next unless ObjectSpace.allocation_generation(obj) == generation

        file = ObjectSpace.allocation_sourcefile(obj) || "(no name)"
        next if @ignore_files && @ignore_files =~ file
        next if @allow_files && !(@allow_files =~ file)

        klass = obj.class rescue nil
        unless Class === klass
          # attempt to determine the true Class when .class returns something other than a Class
          klass = Kernel.instance_method(:class).bind(obj).call
        end
        next if @trace && !trace.include?(klass)

        begin
          line       = ObjectSpace.allocation_sourceline(obj)
          location   = helper.lookup_location(file, line)
          class_name = helper.lookup_class_name(klass)
          gem        = helper.guess_gem(file)

          # we do memsize first to avoid freezing as a side effect and shifting
          # storage to the new frozen string, this happens on @hash[s] in lookup_string
          memsize = ObjectSpace.memsize_of(obj)
          string = klass == String ? helper.lookup_string(obj) : nil

          # compensate for API bug
          memsize = rvalue_size if memsize > 100_000_000_000
          result[obj.__id__] = MemoryProfiler::Stat.new(class_name, gem, file, location, memsize, string)
        rescue
          # give up if any any error occurs inspecting the object
        end
      end

      result
    end
  end
end

module MemoryProfiler
  def self.report(opts = {}, &block)
    Reporter.report(opts, &block)
  end

  def self.start(opts = {})
    unless Reporter.current_reporter
      Reporter.current_reporter = Reporter.new(opts)
      Reporter.current_reporter.start
    end
  end

  def self.stop
    Reporter.current_reporter.stop if Reporter.current_reporter
  ensure
    Reporter.current_reporter = nil
  end
end