test/dtrace/helper.rb

   1 # -*- coding: us-ascii -*-
   2 # frozen_string_literal: false
   3 require 'test/unit'
   4 require 'tempfile'
   5
   6 if Process.euid == 0
   7   ok = true
   8 elsif (sudo = ENV["SUDO"]) and !sudo.empty? and (`#{sudo} echo ok` rescue false)
   9   ok = true
  10 else
  11   ok = false
  12 end
  13
  14 impl = :dtrace
  15
  16 # GNU/Linux distros with Systemtap support allows unprivileged users
  17 # in the stapusr and statdev groups to work.
  18 if RUBY_PLATFORM =~ /linux/
  19   impl = :stap
  20   begin
  21     require 'etc'
  22     ok = (%w[stapusr stapdev].map {|g|(Etc.getgrnam(g) || raise(ArgumentError)).gid} & Process.groups).size == 2
  23   rescue LoadError, ArgumentError
  24   end unless ok
  25 end
  26
  27 if ok
  28   case RUBY_PLATFORM
  29   when /darwin/i
  30     begin
  31       require 'pty'
  32     rescue LoadError
  33     end
  34   end
  35 end
  36
  37 # use miniruby to reduce the amount of trace data we don't care about
  38 rubybin = "miniruby#{RbConfig::CONFIG["EXEEXT"]}"
  39 rubybin = File.join(File.dirname(EnvUtil.rubybin), rubybin)
  40 rubybin = EnvUtil.rubybin unless File.executable?(rubybin)
  41
  42 # make sure ruby was built with --enable-dtrace and we can run
  43 # dtrace(1) or stap(1):
  44 cmd = "#{rubybin} --disable=gems -eexit"
  45 case impl
  46 when :dtrace; cmd = %W(dtrace -l -n ruby$target:::gc-sweep-end -c #{cmd})
  47 when :stap; cmd = %W(stap -l process.mark("gc__sweep__end") -c #{cmd})
  48 else
  49   warn "don't know how to check if built with #{impl} support"
  50   cmd = false
  51 end
  52
  53 NEEDED_ENVS = [RbConfig::CONFIG["LIBPATHENV"], "RUBY", "RUBYOPT"].compact
  54
  55 if cmd and ok
  56   sudocmd = []
  57   if sudo
  58     sudocmd << sudo
  59     NEEDED_ENVS.each {|name| val = ENV[name] and sudocmd << "#{name}=#{val}"}
  60   end
  61   ok = system(*sudocmd, *cmd, err: IO::NULL, out: IO::NULL)
  62 end
  63
  64 module DTrace
  65   class TestCase < Test::Unit::TestCase
  66     INCLUDE = File.expand_path('..', File.dirname(__FILE__))
  67
  68     case RUBY_PLATFORM
  69     when /solaris/i
  70       # increase bufsize to 8m (default 4m on Solaris)
  71       DTRACE_CMD = %w[dtrace -b 8m]
  72     when /darwin/i
  73       READ_PROBES = proc do |cmd|
  74         lines = nil
  75         PTY.spawn(*cmd) do |io, _, pid|
  76           lines = io.readlines.each {|line| line.sub!(/\r$/, "")}
  77           Process.wait(pid)
  78         end
  79         lines
  80       end if defined?(PTY)
  81     end
  82
  83     # only handles simple cases, use a Hash for d_program
  84     # if there are more complex cases
  85     def dtrace2systemtap(d_program)
  86       translate = lambda do |str|
  87         # dtrace starts args with '0', systemtap with '1' and prefixes '$'
  88         str = str.gsub(/\barg(\d+)/) { "$arg#{$1.to_i + 1}" }
  89         # simple function mappings:
  90         str.gsub!(/\bcopyinstr\b/, 'user_string')
  91         str.gsub!(/\bstrstr\b/, 'isinstr')
  92         str
  93       end
  94       out = ''
  95       cond = nil
  96       d_program.split(/^/).each do |l|
  97         case l
  98         when /\bruby\$target:::([a-z-]+)/
  99           name = $1.gsub(/-/, '__')
 100           out << %Q{probe process.mark("#{name}")\n}
 101         when %r{/(.+)/}
 102           cond = translate.call($1)
 103         when "{\n"
 104           out << l
 105           out << "if (#{cond}) {\n" if cond
 106         when "}\n"
 107           out << "}\n" if cond
 108           out << l
 109         else
 110           out << translate.call(l)
 111         end
 112       end
 113       out
 114     end
 115
 116     DTRACE_CMD ||= %w[dtrace]
 117
 118     READ_PROBES ||= proc do |cmd|
 119       IO.popen(cmd, err: [:child, :out], &:readlines)
 120     end
 121
 122     def trap_probe d_program, ruby_program
 123       if Hash === d_program
 124         d_program = d_program[IMPL] or
 125           omit "#{d_program} not implemented for #{IMPL}"
 126       elsif String === d_program && IMPL == :stap
 127         d_program = dtrace2systemtap(d_program)
 128       end
 129       d = Tempfile.new(%w'probe .d')
 130       d.write d_program
 131       d.flush
 132
 133       rb = Tempfile.new(%w'probed .rb')
 134       rb.write ruby_program
 135       rb.flush
 136
 137       d_path  = d.path
 138       rb_path = rb.path
 139       cmd = "#{RUBYBIN} --disable=gems -I#{INCLUDE} #{rb_path}"
 140       if IMPL == :stap
 141         cmd = %W(stap #{d_path} -c #{cmd})
 142       else
 143         cmd = [*DTRACE_CMD, "-q", "-s", d_path, "-c", cmd ]
 144       end
 145       if sudo = @@sudo
 146         NEEDED_ENVS.each do |name|
 147           if val = ENV[name]
 148             cmd.unshift("#{name}=#{val}")
 149           end
 150         end
 151         cmd.unshift(sudo)
 152       end
 153       probes = READ_PROBES.(cmd)
 154       d.close(true)
 155       rb.close(true)
 156       yield(d_path, rb_path, probes)
 157     end
 158   end
 159 end if ok
 160
 161 if ok
 162   DTrace::TestCase.class_variable_set(:@@sudo, sudo)
 163   DTrace::TestCase.const_set(:IMPL, impl)
 164   DTrace::TestCase.const_set(:RUBYBIN, rubybin)
 165 end