#!/usr/bin/env ruby
# Analyze heap dumps for memory leaks
#
# Gather the memory dumps with
# bundle exec rbtrace — timeout 30 -e ‘load “#{Rails.root}/scripts/heap_dump.rb”’ -p $RAILS_PID

require 'set'
require 'json'

if ARGV.length != 3
  puts "Usage: detect_leaks [FIRST.json] [SECOND.json] [THIRD.json]"
  exit 1
end

first_addrs = Set.new
third_addrs = Set.new

# Get a list of memory addresses from the first dump
ignored_lines = 0
File.open(ARGV[0], "r").each_line do |line|
  begin
    parsed = JSON.parse(line)
    first_addrs << parsed["address"] if parsed && parsed["address"]
  rescue
    ignored_lines += 1
  end
end
puts "Ignored #{ignored_lines} lines in first file"

# Get a list of memory addresses from the last dump
ignored_lines = 0
File.open(ARGV[2], "r").each_line do |line|
  begin
    parsed = JSON.parse(line)
    third_addrs << parsed["address"] if parsed && parsed["address"]
  rescue
    ignored_lines += 1
  end
end
puts "Ignored #{ignored_lines} lines in last file"

diff = []

# Get a list of all items present in both the second and
# third dumps but not in the first.
ignored_lines = 0
File.open(ARGV[1], "r").each_line do |line|
  begin
    parsed = JSON.parse(line)
    if parsed && parsed["address"]
      if !first_addrs.include?(parsed["address"]) && third_addrs.include?(parsed["address"])
        diff << parsed
      end
    end
  rescue
    ignored_lines += 1
  end
end
puts "Ignored #{ignored_lines} lines in second file"

# Group items
diff.group_by do |x|
  [x["type"], x["file"], x["line"]]
end.map do |x,y|
  # Collect memory size
  [x, y.count, y.inject(0){|sum,i| sum + (i['bytesize'] || 0) }, y.inject(0){|sum,i| sum + (i['memsize'] || 0) }]
end.sort do |a,b|
  b[1] <=> a[1]
end.each do |x,y,bytesize,memsize|
  # Output information about each potential leak
  puts "Leaked #{y} #{x[0]} objects of size #{bytesize}/#{memsize} at: #{x[1]}:#{x[2]}"
end

# Also output total memory usage, because why not?
memsize = diff.inject(0){|sum,i| sum + (i['memsize'] || 0) }
bytesize = diff.inject(0){|sum,i| sum + (i['bytesize'] || 0) }
puts "\n\nTotal Size: #{bytesize}/#{memsize}"