A very rough draft of a ruby heap analyzer

heap.rb

# This is just a sketch of some useful tools for analyzing ruby
# heap dump files while investigating memory leaks.  Maybe I'll
# clean this up and make a gem out of it some day.

require 'json'
require 'pp'
require 'set'

def for_each_line(method)
  File.open('/tmp/heap.json', 'r') do |f|
    f.each_line do |line|
      method.call(JSON.parse(line))
    end
  end
end

def prepare_hash_for_output(o)
  class_name = $map_address_to_class_name[o['class']]
  o['class'] = "#{o['class']} (#{class_name})" if class_name
  o['file'] = "#{o['file']}:#{o['line']}" if o['file']
  o.delete 'bfs_parent'
  o
end

def find_shortest_path(heap, from_address, to_address)
  visited_addresses = Set.new
  current_level_addresses = [from_address]
  next_level_addresses = []
  found = false

  loop do

    while (node_address = current_level_addresses.pop) do
      # puts "searching at node #{node_address}"

      if node_address == to_address
        found = true
        # puts "found!"
        break
      end

      if visited_addresses.include? node_address
        # puts "skipping visited node"
        next
      end
      visited_addresses.add node_address

      children = heap[node_address]['references'] || []
      children.each do |child_address|
        next if visited_addresses.include? child_address
        child_node = heap[child_address]
        next unless child_node
        # puts "adding child to next level, child=#{child_address}, parent=#{node_address}"
        next_level_addresses << child_address
        child_node['bfs_parent'] = node_address
      end
    end

    # puts "current level exhausted, found=#{found}"
    break if found || next_level_addresses.length == 0

    # puts "moving to next level"
    current_level_addresses = next_level_addresses
    next_level_addresses = []
  end

  if found
    # puts "building path to found object"
    path = []
    node = heap[to_address]

    while node
      # puts "node #{node['address']} going into path"
      path.unshift node
      parent_address = node['bfs_parent']
      break unless parent_address
      node = heap[parent_address]
    end

    return path
  end

  nil
end


#######

first_arg, *args = ARGV

$map_address_to_class_name = {}
build_map = lambda do |o|
  if o['type'] == 'CLASS' || o['type'] == 'MODULE'
    $map_address_to_class_name[o['address']] = o['name']
  end
end
for_each_line(build_map)

after_all = nil

if first_arg == 'class'
  class_name = args[0]
  each_line = lambda do |o|
    if (o['type'] == 'CLASS' && o['name'] == class_name)
      pp prepare_hash_for_output(o)
    end
  end
elsif first_arg == 'at_address'
  obj_address = args[0]
  each_line = lambda do |o|
    if o['address'] == obj_address
      pp prepare_hash_for_output(o)
    end
  end
elsif first_arg == 'instances_of'
  class_address = args[0]
  each_line = lambda do |o|
    if o['class'] == class_address
      pp prepare_hash_for_output(o)
    end
  end
elsif first_arg == 'count_instances_of'
  class_address = args[0]
  count = 0
  each_line = lambda do |o|
    if o['class'] == class_address
      count += 1
    end
  end
  after_all = lambda do
    puts "Found #{count} instances of #{class_address}"
  end
elsif first_arg == 'refers_to'
  instance_address = args[0]
  each_line = lambda do |o|
    if o['references'] && o['references'].include?(instance_address)
      pp prepare_hash_for_output(o)
    end
  end
elsif first_arg == 'shortest_path'
  from_address = args[0]
  to_address = args[1]
  heap = {}
  each_line = lambda do |o|
    heap[o['address']] = o
  end
  after_all = lambda do
    path = find_shortest_path(heap, from_address, to_address)
    if path
      path.each do |o|
        pp prepare_hash_for_output(o)
      end
    else
      puts "No path from #{from_address} to #{to_address} found"
    end
  end
else
  $stderr.puts 'Usage:'
  $stderr.puts '  heap.rb class <class_name>'
  $stderr.puts '  heap.rb at_address <address>'
  $stderr.puts '  heap.rb instances_of <class_address>'
  $stderr.puts '  heap.rb count_instances_of <class_address>'
  $stderr.puts '  heap.rb refers_to <address>'
  $stderr.puts '  heap.rb shortest_path <from_address> <to_address>'
  exit 1
end

for_each_line(each_line)
after_all.call if after_all