memtrace

#! /usr/bin/env ruby

# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

$brk_base = nil
$brk_max = nil
$mmaped = 0

require 'optparse'
require 'tmpdir'

$opts = {
  :ignore => false,
  :all => false,
  :verbose => false,
  :debug => false,
  :unbalanced => false,
  :output => nil,
  :raw => nil,
  :file => false,
  :gprefix => nil,
}

parser = OptionParser.new do |o|
  o.banner = "#{File.basename($0)} [options] trace_file|command"

  o.on("--ignore-unknown", "Ignore unknown lines") { 
    $opts[:ignore] = true
  }
  o.on("--verbose", "Print details for all subprocesses") {
    $opts[:verbose] = true
  }
  o.on("--all", "Include read only maps in count") {
    $opts[:all] = true
  }
  o.on("--gprefix PREFIX", "Prefix for graph files. Draw with 'gnuplot -p PREFIX.gp'") { |v|
    $opts[:gprefix] = v
  }
  o.on("--unbalanced", "Display unbalanced mmap/munmap calls") {
    $opts[:unbalanced] = true
  }
  o.on("--debug", "Print debugging information") {
    $opts[:debug] = true
  }
  o.on("--output FILE", "-o", "Output to FILE instead of stderr") { |v|
    $opts[:output] = v
  }
  o.on("--raw FILE", "-r", "Output raw strace to FILE") { |v|
    $opts[:raw] = v
  }
  o.on("--file", "-f", "Arguments are strace output files") {
    $opts[:file] = true
  }
  o.on_tail("--help", "This message") {
    puts(o)
    exit(0)
  }
end
parser.order!(ARGV)
if ARGV.empty?
  STDERR.puts("Missing arguments. See --help for details")
  exit(1)
end

class String
  def parse_int
    case self
    when "NULL"
      0
    when /^0x/
      self.to_i(16)
    when /^0/
      self.to_i(8)
    else
      self.to_i
    end
  end
end

class Numeric
  SI_LARGE_SUFFIX = [' ', 'k', 'M', 'G', 'T', 'P', 'E', 'Z', 'Y']
  SI_SMALL_SUFFIX = ['m', 'u', 'n', 'p', 'f', 'a', 'z', 'y']
  def to_SI(multiple = 1000.0)
    nb = self.to_f
    neg = nb < 0
    nb = nb.abs
    return "0 " if nb == 0.0
    multiple = multiple.to_f
    if nb >= 1.0
      SI_LARGE_SUFFIX.each do |s|
        return("%s%.3g%s" % [neg ? "-" : "", nb, s]) if nb < multiple
        nb = nb / multiple
      end
    else
      SI_SMALL_SUFFIX.each do |s|
        nb *= multiple
        return("%s%.3g%s" % [neg ? "-" : "", nb, s]) if nb >= 1.0
      end
    end
  end
end

def print_report(io, title, heap, mapped, vm)
  io.puts("%10s: heap %10d %8s mmap %10d %8s vm %10d %8s" %
          [title[0, 10], heap, "(#{heap.to_SI(1024)})",
           mapped, "(#{mapped.to_SI(1024)})",
           vm, "(#{vm.to_SI(1024)})"])
end

class Graph
  def initialize(prefix)
    @prefix = prefix
    @data = File.open(@prefix + ".data", "w")
    @data.puts("#time heap mmaped vm")
    @base_time = nil
  end

  def add(time, heap, mmapped, vm)
    @base_time = time unless @base_time
    s = "#{time - @base_time} #{heap / 1024} #{mmapped / 1024} #{vm / 1024}"
    @data.puts(s)
  end

  def close
    @data.close
    open(@prefix + ".gp", "w") { |fd|
      fd.puts("set title \"memtrace\"",
              "set xlabel \"time (s)\"",
              "set ylabel \"Memory (Kb)\"",
              "set key on outside",
              "set style data steps",
              "plot \"#{@prefix + ".data"}\" using 1:4 title \"vm\" ls 1, \"\" using 1:2 title \"heap\" ls 2, \"\" using 1:3 title \"mmap\" ls 3",
              "replot \"#{@prefix + ".data"}\" using 1:4 title \"\" ls 1, \"\" using 1:2 title \"\" ls 2, \"\" using 1:3 title \"\" ls 3")
    }
  end
end

class ProcessStats
  def initialize(known_addr, pid)
    @pid = pid
    @brk_base = @brk_cur = nil
    @heap = @heap_max = 0
    @mmaped = @mmaped_max = 0
    @vm = @vm_max = 0
    @known_addr = known_addr
  end
  
  attr_reader :heap_max, :mmaped_max, :vm_max
  attr_reader :heap, :mmaped, :vm
  attr_reader :brk_base, :brk_cur, :pid

  def update_vm
    @vm = @heap + @mmaped
    @vm_max = @vm if @vm > @vm_max
  end

  def brk(args, res, line)
    addr = args.parse_int
    res = res.parse_int
    return if !@brk_cur.nil? && res <= @brk_cur
    if addr == 0
      @brk_cur = res
      @brk_base ||= res
    else
      @brk_cur = res
      @heap = @brk_cur - @brk_base
      @heap_max = @heap if @heap > @heap_max
    end
    update_vm
  end

  def mmap(args, res, line)
    res = res.split[0].parse_int
    return if res == -1
    addr, length, prot, flags, fd, off = args.split(/\s*,\s*/)
    prot = prot.split(/\s*\|\s*/)
    if $opts[:all] || prot.include?("PROT_WRITE")
      @mmaped += length.parse_int
      @mmaped_max = @mmaped if @mmaped > @mmaped_max
      (@known_addr[res] ||= []) << line
    end
    update_vm
  end

  def mremap(args, res, line)
    res = res.split[0].parse_int
    return if res == -1
    old_address, old_size, new_size = args.split(/\s*,\s*/)
    return unless @known_addr.delete(old_address)
    @mmaped += new_size - old_size
    @mmaped_max = @mmaped if @mmaped > @mmaped_max
    (@known_addr[res] ||= []) << line
  end

  def munmap(args, res, line)
    ires = res.split[0].parse_int
    return if ires == -1
    addr, length = args.split(/\s*,\s*/)
    addr = addr.parse_int
    if @known_addr.delete(addr)
      @mmaped -= length.parse_int
    end
    update_vm
  end

  def method_missing(meth, *args)
    unless $opts[:ignore]
      raise "Unknown system call #{meth} '#{args.inspect}'"
    end
  end
end

class ProcessStatHash
  def initialize
    @known_addr = {}
    @hash = Hash.new { |h, k| h[k] = ProcessStats.new(@known_addr, k) }
    @heap = @mmaped = @vm = 0
    @heap_max = @mmaped_max = @vm_max = 0
    @warned = false
    @graph = nil
    @dead = false
  end
  attr_reader :heap, :mmaped, :vm
  attr_reader :heap_max, :mmaped_max, :vm_max
  attr_reader :known_addr
  attr_accessor :graph

  def update(time, pid, cmd, args, res, line)
    ps = @hash[pid]

    case cmd
    when "clone"
      res = res.parse_int
      return if res == -1
      @hash[res] = ps if args =~ /CLONE_VM/

    when "brk", "mmap", "mremap", "munmap"
      ps_heap, ps_mmaped, ps_vm = ps.heap, ps.mmaped, ps.vm
      ps.__send__(cmd, args, res, line)
      @heap += ps.heap - ps_heap
      @mmaped += ps.mmaped - ps_mmaped
      @vm += ps.vm - ps_vm
      @heap_max = @heap if @heap > @heap_max
      @mmaped_max = @mmaped if @mmaped > @mmaped_max
      @vm_max = @vm if @vm > @vm_max
      @graph.add(time, @heap, @mmaped, @vm) if @graph

    when "die", "exit_group"
      if ps.pid == pid && !@dead
        @heap -= ps.heap
        @mmaped -= ps.mmaped
        @vm -= ps.vm
        @graph.add(time, @heap, @mmaped, @vm) if @graph
        @dead = true
      end

    else
      if !@warned
        STDERR.puts("Warning: unknown system calls in input. The results should still be correct unless some of the 'brk', 'mmap', 'mremap', 'munmap' and 'clone' calls are missing")
        @warned = true
      end
    end
  end

  def each(&block); @hash.each(&block); end
end

LRGEXP = /^(\d+\s*)?(\d+.\d+\s*)?(\w+)\s*\((.*?)\)\s*=\s*(.*)$/
UNFEXP = /^(\d+\s*)?(\d+.\d+\s*)?(.*?) <unfinished ...>$/
RESEXP = /^(\d+\s*)?(\d+.\d+\s*)?<... \w+ resumed> (.*)$/
SIGEXP = /^(\d+\s*)?(\d+.\d+\s*)?---[^-]*---$/
KILEXP = /^(\d+\s*)?(\d+.\d+\s*)?\+\+\+ killed by (\w+) \+\+\+/
EXTEXP = /^(\d+\s*)?(\d+.\d+\s*)?\+\+\+ exited with (\d+) \+\+\+/

def parse_file(io, pss, out_io = stderr, raw_io = nil)
  i = 0
  line_buf = {}
  graph = nil
  pss.graph = $opts[:gprefix] && Graph.new($opts[:gprefix])

  io.each_line do |l|
    out_io.puts("--> #{l}") if $opts[:debug]
    if raw_io
      raw_io.print(l)
    end

    case l
    when SIGEXP
      next
      
    when KILEXP, EXTEXP
      pid = $1 ? $1.chomp.parse_int : 0
      time = $2.to_f
      pss.update(time, pid, "die", nil, nil, l)
      next

    when UNFEXP
      pid = $1 ? $1.chomp.parse_int : 0
      line_buf[pid] = $3
      next

    when RESEXP
      pid = $1 ? $1.chomp.parse_int : 0
      line = "#{$1} #{$2} #{line_buf[pid]} #{$3}"

    else
      line = l
    end

    unless line =~ LRGEXP
      raise "Line did not parse: #{line}" unless $opts[:ignore]
      next
    end

    i += 1
    pid, time, cmd, args, res = $1, $2, $3, $4, $5
    time = time.to_f
    pid = pid ? pid.chomp.parse_int : 0
    pss.update(time, pid, cmd, args, res, l)
    if $opts[:debug]
      out_io.puts([pid, cmd, args, res].inspect)
      print_report(out_io, "cur_#{i}", pss.heap, pss.mmaped, pss.vm)
      print_report(out_io, "max_#{i}", pss.heap_max, pss.mmaped_max, pss.vm_max)
    end
  end

ensure
  pss.graph.close if pss.graph
end

exit_value = nil
begin
  pss = ProcessStatHash.new

  output_io = $opts[:file] ? STDOUT : STDERR
  output_io = open($opts[:output], "w") if $opts[:output]


  if $opts[:file]
    ARGV.each do |f|
      open(f) { |io| parse_file(io, pss, output_io) }
    end
  else
    output_raw = nil
    output_raw = open($opts[:raw], "w")  if $opts[:raw]
    r_io, w_io = IO.pipe
    cmd  = ["strace", "-ttt", "-f", "-o", "/dev/fd/#{w_io.fileno}", "-e", "trace=brk,mmap,mremap,munmap,clone,exit_group"]
    cmd += ARGV
    pid = spawn(*cmd, { :close_others => true, w_io => w_io })
    w_io.close
    
    parse_file(r_io, pss, output_io, output_raw)
    Process.wait(pid)
    exit_value = $?
    r_io.close
    output_raw.close if output_raw
  end

  if $opts[:verbose]
    pss.each do |pid, ps|
      print_report(output_io, "max_#{pid}", ps.heap_max, ps.mmaped_max, ps.vm_max)
    end
  end


  print_report(output_io, "max", pss.heap_max, pss.mmaped_max, pss.vm_max)

  if $opts[:unbalanced]
    output_io.puts("Unbalanced: ")
    pss.known_addr.each { |k, v|
      output_io.puts("0x%16x %s" % [k, v.shift])
      v.each { |l| puts("                   #{l}") }
    }
  end
rescue Errno::EPIPE
# ignore
end

exit(0) if exit_value.nil?
if exit_value.exited?
  exit(exit_value.exitstatus)
elsif exit_value.signaled?
  trap(exit_value.termsig, "SYSTEM_DEFAULT")
  Process.kill(exit_value.termsig, 0)
end