Enable parallel tracing support on Dtrace

tools/isolate/trace_inputs.py

 #!/usr/bin/env python
 # coding=utf-8
 # Copyright (c) 2012 The Chromium Authors. All rights reserved.
 # Use of this source code is governed by a BSD-style license that can be
 # found in the LICENSE file.
 
 """Traces an executable and its child processes and extract the files accessed
 by them.
 
 The implementation uses OS-specific API. The native Kernel logger and the ETL
@@ skipping 11 matching lines @@
 import csv
 import getpass
 import glob
 import json
 import logging
 import optparse
 import os
 import re
 import subprocess
 import sys
+import tempfile
 import threading
+import time
 import weakref
 
 ## OS-specific imports
 
 if sys.platform == 'win32':
   from ctypes.wintypes import byref, create_unicode_buffer, c_int, c_wchar_p
   from ctypes.wintypes import windll, FormatError  # pylint: disable=E0611
   from ctypes.wintypes import GetLastError  # pylint: disable=E0611
 elif sys.platform == 'darwin':
   import Carbon.File  #  pylint: disable=F0401
@@ skipping 1095 matching lines @@
     RE_EXECVE = re.compile(r'^\"(.+?)\", \[(\d+), (.+)\]$')
     RE_OPEN = re.compile(r'^\"(.+?)\", (0x[0-9a-z]+), (0x[0-9a-z]+)$')
     RE_PROC_START = re.compile(r'^(\d+), \"(.+?)\", (\d+)$')
     RE_RENAME = re.compile(r'^\"(.+?)\", \"(.+?)\"$')
 
     O_DIRECTORY = 0x100000
 
     class Process(ApiBase.Context.Process):
       pass
 
-    def __init__(self, blacklist):
+    def __init__(self, blacklist, tracer_pid, initial_cwd):
+      logging.info(
+          '%s(%d, %s)' % (self.__class__.__name__, tracer_pid, initial_cwd))
       super(Dtrace.Context, self).__init__(blacklist)
       # Process ID of the trace_child_process.py wrapper script instance.
-      self._tracer_pid = None
+      self._tracer_pid = tracer_pid
+      self._initial_cwd = initial_cwd
       # First process to be started by self._tracer_pid.
       self._initial_pid = None
-      self._initial_cwd = None
       self._line_number = 0
 
     def on_line(self, line):
       self._line_number += 1
       match = self.RE_HEADER.match(line)
       if not match:
         raise TracingFailure(
             'Found malformed line: %s' % line,
             None,
             self._line_number,
@@ skipping 30 matching lines @@
       if sorted(self.processes) != sorted(p.pid for p in process.all):
         raise TracingFailure(
             'Found internal inconsitency in process lifetime detection',
             None,
             None,
             None,
             sorted(self.processes),
             sorted(p.pid for p in process.all))
       return Results(process)
 
-    def handle_dtrace_BEGIN(self, pid, args):
-      if self._tracer_pid or self._initial_pid:
+    def handle_dtrace_BEGIN(self, _pid, args):
+      if not self.RE_DTRACE_BEGIN.match(args):
         raise TracingFailure(
             'Found internal inconsitency in dtrace_BEGIN log line',
             None, None, None)
-      self._tracer_pid = pid
-      self._initial_cwd = self.RE_DTRACE_BEGIN.match(args).group(1)
 
     def handle_proc_start(self, pid, args):
       """Transfers cwd.
 
       The dtrace script already takes care of only tracing the processes that
       are child of the traced processes so there is no need to verify the
       process hierarchy.
       """
       if pid in self.processes:
         raise TracingFailure(
             'Found internal inconsitency in proc_start: %d started two times' %
                 pid,
             None, None, None)
       match = self.RE_PROC_START.match(args)
       if not match:
         raise TracingFailure(
             'Failed to parse arguments: %s' % args,
             None, None, None)
       ppid = int(match.group(1))
-      if (ppid == self._tracer_pid) == bool(self._initial_pid):
-        raise TracingFailure(
-            ( 'Parent process is _tracer_pid(%d) but _initial_pid(%d) is '
-              'already set') % (self._tracer_pid, self._initial_pid),
-            None, None, None)
-      if not self._initial_pid:
+      if ppid == self._tracer_pid and not self._initial_pid:
         self._initial_pid = pid
         ppid = None
         cwd = self._initial_cwd
-      else:
+      elif ppid in self.processes:
         parent = self.processes[ppid]
         cwd = parent.cwd
+      else:
+        # Another process tree, ignore.
+        return
       proc = self.processes[pid] = self.Process(self, pid, cwd, ppid)
       proc.cwd = cwd
-      logging.debug('New child: %s -> %d' % (ppid, pid))
+      logging.debug('New child: %s -> %d  cwd:%s' % (ppid, pid, unicode(cwd)))
 
     def handle_proc_exit(self, pid, _args):
       """Removes cwd."""
-      if pid != self._tracer_pid:
-        # self._tracer_pid is not traced itself.
+      if pid in self.processes:
+        # self._tracer_pid is not traced itself and other traces run neither.
         self.processes[pid].cwd = None
 
     def handle_execve(self, pid, args):
       """Sets the process' executable.
 
       TODO(maruel): Read command line arguments.  See
       https://discussions.apple.com/thread/1980539 for an example.
       https://gist.github.com/1242279
 
       Will have to put the answer at http://stackoverflow.com/questions/7556249.
       :)
       """
+      if not pid in self.processes:
+        # Another process tree, ignore.
+        return
       match = self.RE_EXECVE.match(args)
       if not match:
         raise TracingFailure(
             'Failed to parse arguments: %s' % args,
             None, None, None)
       proc = self.processes[pid]
       proc.executable = match.group(1)
       proc.command = process_quoted_arguments(match.group(3))
       if int(match.group(2)) != len(proc.command):
         raise TracingFailure(
             'Failed to parse execve() arguments: %s' % args,
             None, None, None)
 
     def handle_chdir(self, pid, args):
       """Updates cwd."""
-      if not self._tracer_pid or pid not in self.processes:
-        raise TracingFailure(
-            'Found inconsistency in the log, is it an old log?',
-            None, None, None)
+      if pid not in self.processes:
+        # Another process tree, ignore.
+        return
       cwd = self.RE_CHDIR.match(args).group(1)
       if not cwd.startswith('/'):
         cwd2 = os.path.join(self.processes[pid].cwd, cwd)
         logging.debug('handle_chdir(%d, %s) -> %s' % (pid, cwd, cwd2))
       else:
         logging.debug('handle_chdir(%d, %s)' % (pid, cwd))
         cwd2 = cwd
       self.processes[pid].cwd = cwd2
 
     def handle_open_nocancel(self, pid, args):
       """Redirects to handle_open()."""
       return self.handle_open(pid, args)
 
     def handle_open(self, pid, args):
+      if pid not in self.processes:
+        # Another process tree, ignore.
+        return
       match = self.RE_OPEN.match(args)
       if not match:
         raise TracingFailure(
             'Failed to parse arguments: %s' % args,
             None, None, None)
       flag = int(match.group(2), 16)
       if self.O_DIRECTORY & flag == self.O_DIRECTORY:
         # Ignore directories.
         return
       self._handle_file(pid, match.group(1))
 
     def handle_rename(self, pid, args):
+      if pid not in self.processes:
+        # Another process tree, ignore.
+        return
       match = self.RE_RENAME.match(args)
       if not match:
         raise TracingFailure(
             'Failed to parse arguments: %s' % args,
             None, None, None)
       self._handle_file(pid, match.group(1))
       self._handle_file(pid, match.group(2))
 
     def _handle_file(self, pid, filepath):
       if not filepath.startswith('/'):
         filepath = os.path.join(self.processes[pid].cwd, filepath)
       # We can get '..' in the path.
       filepath = os.path.normpath(filepath)
       # Sadly, still need to filter out directories here;
       # saw open_nocancel(".", 0, 0) = 0 lines.
       if os.path.isdir(filepath):
         return
       self.processes[pid].add_file(filepath)
 
+    def handle_ftruncate(self, pid, args):
+      """Just used as a signal to kill dtrace, ignoring."""
+      pass
+
     @staticmethod
     def _handle_ignored(pid, args):
       """Is called for all the event traces that are not handled."""
       raise NotImplementedError('Please implement me')
 
   class Tracer(ApiBase.Tracer):
     # pylint: disable=C0301
     #
     # To understand the following code, you'll want to take a look at:
     # http://developers.sun.com/solaris/articles/dtrace_quickref/dtrace_quickref.html
@@ skipping 28 matching lines @@
     # $ grep syscall probes.txt | less
     #    or use dtrace directly:
     # $ sudo dtrace -l -P syscall | less
     #
     # trackedpid is an associative array where its value can be 0, 1 or 2.
     # 0 is for untracked processes and is the default value for items not
     #   in the associative array.
     # 1 is for tracked processes.
     # 2 is for trace_child_process.py only. It is not tracked itself but
     #   all its decendants are.
+    #
+    # The script will kill itself only once waiting_to_die == 1 and
+    # current_processes == 0, so that both getlogin() was called and that
+    # all traced processes exited.
     D_CODE = """
       dtrace:::BEGIN {
-        /* Since the child process is already started, initialize
-           current_processes to 1. */
-        current_processes = 1;
+        waiting_to_die = 0;
+        current_processes = 0;
         logindex = 0;
-        trackedpid[TRACER_PID] = 2;
         printf("%d %d %s_%s(\\"%s\\")\\n",
-               logindex, TRACER_PID, probeprov, probename, CWD);
+               logindex, PID, probeprov, probename, SCRIPT);
         logindex++;
       }
 
       proc:::start /trackedpid[ppid]/ {
         trackedpid[pid] = 1;
         current_processes += 1;
         printf("%d %d %s_%s(%d, \\"%s\\", %d)\\n",
                logindex, pid, probeprov, probename,
                ppid,
                execname,
                current_processes);
         logindex++;
       }
-      proc:::exit /trackedpid[pid] && current_processes == 1/ {
-        /* Last process is exiting. */
+      /* Should use SCRIPT but there is no access to this variable at that
+       * point. */
+      proc:::start /ppid == PID && execname == "Python"/ {
+        trackedpid[pid] = 2;
+        current_processes += 1;
+        printf("%d %d %s_%s(%d, \\"%s\\", %d)\\n",
+               logindex, pid, probeprov, probename,
+               ppid,
+               execname,
+               current_processes);
+        logindex++;
+      }
+      proc:::exit /trackedpid[pid] &&
+                  current_processes == 1 &&
+                  waiting_to_die == 1/ {
         trackedpid[pid] = 0;
         current_processes -= 1;
         printf("%d %d %s_%s(%d)\\n",
                logindex, pid, probeprov, probename,
                current_processes);
         logindex++;
         exit(0);
       }
       proc:::exit /trackedpid[pid]/ {
         trackedpid[pid] = 0;
         current_processes -= 1;
         printf("%d %d %s_%s(%d)\\n",
                logindex, pid, probeprov, probename,
                current_processes);
         logindex++;
       }
 
+      /* Use an arcane function to detect when we need to die */
+      syscall::ftruncate:entry /pid == PID && arg0 == FILE_ID/ {
+        waiting_to_die = 1;
+        printf("%d %d %s()\\n", logindex, pid, probefunc);
+        logindex++;
+      }
+      syscall::ftruncate:entry /
+          pid == PID && arg0 == FILE_ID && current_processes == 0/ {
+        exit(0);
+      }
+
       syscall::open*:entry /trackedpid[pid] == 1/ {
         self->open_arg0 = arg0;
         self->open_arg1 = arg1;
         self->open_arg2 = arg2;
       }
       syscall::open*:return /trackedpid[pid] == 1 && errno == 0/ {
         this->open_arg0 = copyinstr(self->open_arg0);
         printf("%d %d %s(\\"%s\\", 0x%x, 0x%x)\\n",
                logindex, pid, probefunc,
                this->open_arg0,
@@ skipping 196 matching lines @@
         syscall::removexattr:return,
         syscall::setxattr:return,
         syscall::stat64:return,
         syscall::stat:return,
         syscall::truncate:return,
         syscall::unlink:return,
         syscall::utimes:return,
       */
       """
 
-    @classmethod
-    def code(cls, pid, cwd):
+    def __init__(self, logname):
+      """Starts the log collection with dtrace.
+
+      Requires root access or chmod 4555 on dtrace. dtrace is asynchronous so
+      this needs to wait for dtrace to be "warmed up".
+      """
+      super(Dtrace.Tracer, self).__init__(logname)
+      self._script = os.path.join(BASE_DIR, 'trace_child_process.py')
+      # This unique dummy temp file is used to signal the dtrace script that it
+      # should stop as soon as all the child processes are done. A bit hackish
+      # but works fine enough.
+      self._dummy_file_id, self._dummy_file_name = tempfile.mkstemp(
+          prefix='trace_signal_file')
+
+      # Note: do not use the -p flag. It's useless if the initial process quits
+      # too fast, resulting in missing traces from the grand-children. The D
+      # code manages the dtrace lifetime itself.
+      trace_cmd = [
+        'sudo',
+        'dtrace',
+        '-x', 'dynvarsize=4m',
+        '-x', 'evaltime=exec',
+        '-o', '/dev/stderr',
+        '-q',
+        '-n', self._get_dtrace_code(),
+      ]
+      with open(self._logname + '.log', 'wb') as logfile:
+        self._dtrace = subprocess.Popen(
+            trace_cmd, stdout=logfile, stderr=subprocess.STDOUT)
+      logging.debug('Started dtrace pid: %d' % self._dtrace.pid)
+
+      # Reads until one line is printed, which signifies dtrace is up and ready.
+      with open(self._logname + '.log', 'rb') as logfile:
+        while 'dtrace_BEGIN' not in logfile.readline():
+          if self._dtrace.poll() is not None:
+            # Do a busy wait. :/
+            break
+      logging.debug('dtrace started')
+
+    def _get_dtrace_code(self):
       """Setups the D code to implement child process tracking.
 
-      Injects the pid and the initial cwd into the trace header for context.
-      The reason is that the child process trace_child_process.py is already
-      running at that point so:
-      - no proc_start() is logged for it.
-      - there is no way to figure out the absolute path of cwd in kernel on OSX
+      Injects the cookie in the script so it knows when to stop.
+
+      The script will detect any trace_child_process.py instance and will start
+      tracing it.
       """
-      cwd = os.path.realpath(cwd).replace('\\', '\\\\').replace('%', '%%')
       return (
-          'inline int TRACER_PID = %d;\n'
-          'inline string CWD = "%s";\n'
+          'inline int PID = %d;\n'
+          'inline string SCRIPT = "%s";\n'
+          'inline int FILE_ID = %d;\n'
           '\n'
           '%s\n'
-          '%s') % (pid, cwd, cls.D_CODE, cls.D_CODE_EXECVE)
+          '%s') % (
+              os.getpid(),
+              self._script,
+              self._dummy_file_id,
+              self.D_CODE,
+              self.D_CODE_EXECVE)
 
     def trace(self, cmd, cwd, tracename, output):
       """Runs dtrace on an executable.
 
       This dtruss is broken when it starts the process itself or when tracing
       child processes, this code starts a wrapper process
       trace_child_process.py, which waits for dtrace to start, then
       trace_child_process.py starts the executable to trace.
       """
       logging.info('trace(%s, %s, %s, %s)' % (cmd, cwd, tracename, output))
       logging.info('Running: %s' % cmd)
-      signal = 'Go!'
       logging.debug('Our pid: %d' % os.getpid())
       with self._lock:
         if not self._initialized:
           raise TracingFailure(
               'Called Tracer.trace() on an unitialized object',
               None, None, None, tracename)
         assert tracename not in (i['trace'] for i in self._traces)
 
-      # Part 1: start the child process.
+      # Starts the script wrapper to start the child process. This signals the
+      # dtrace script that this process is to be traced.
       stdout = stderr = None
       if output:
         stdout = subprocess.PIPE
         stderr = subprocess.STDOUT
       child_cmd = [
         sys.executable,
-        os.path.join(BASE_DIR, 'trace_child_process.py'),
-        '--wait',
+        self._script,
+        tracename,
       ]
+      # Call a dummy function so that dtrace knows I'm about to launch a process
+      # that needs to be traced.
+      # Yummy.
       child = subprocess.Popen(
           child_cmd + cmd,
           stdin=subprocess.PIPE,
           stdout=stdout,
           stderr=stderr,
           cwd=cwd)
       logging.debug('Started child pid: %d' % child.pid)
 
-      # Part 2: start dtrace process.
-      # Note: do not use the -p flag. It's useless if the initial process quits
-      # too fast, resulting in missing traces from the grand-children. The D
-      # code manages the dtrace lifetime itself.
-      trace_cmd = [
-        'sudo',
-        'dtrace',
-        '-x', 'dynvarsize=4m',
-        '-x', 'evaltime=exec',
-        '-n', self.code(child.pid, cwd),
-        '-o', '/dev/stderr',
-        '-q',
-      ]
-      logname = self._logname + '.log'
-      with open(logname, 'w') as logfile:
-        dtrace = subprocess.Popen(
-            trace_cmd, stdout=logfile, stderr=subprocess.STDOUT)
-      logging.debug('Started dtrace pid: %d' % dtrace.pid)
-
-      # Part 3: Read until one line is printed, which signifies dtrace is up and
-      # ready.
-      with open(logname, 'r') as logfile:
-        while 'dtrace_BEGIN' not in logfile.readline():
-          if dtrace.poll() is not None:
-            break
-
-      # Part 4: We can now tell our child to go.
-      # TODO(maruel): Another pipe than stdin could be used instead. This
-      # would be more consistent with the other tracing methods.
-      out = child.communicate(signal)[0]
-
-      dtrace.wait()
-      if dtrace.returncode != 0:
-        print 'dtrace failure: %d' % dtrace.returncode
-        with open(logname) as logfile:
-          print ''.join(logfile.readlines()[-100:])
-        # Find a better way.
-        os.remove(logname)
+      out = child.communicate()[0]
+      # This doesn't mean tracing is done, one of the grand-child process may
+      # still be alive. It will be tracked with the dtrace script.
 
       with self._lock:
         assert tracename not in (i['trace'] for i in self._traces)
         self._traces.append(
           {
             'cmd': cmd,
             'cwd': cwd,
             # The pid of strace process, not very useful.
             'pid': child.pid,
             'trace': tracename,
             'output': out,
           })
-      return dtrace.returncode or child.returncode, out
+      return child.returncode, out
+
+    def close(self, timeout=None):
+      """Terminates dtrace."""
+      logging.debug('close(%s)' % timeout)
+      try:
+        try:
+          super(Dtrace.Tracer, self).close(timeout)
+          # Signal dtrace that it should stop now.
+          os.ftruncate(self._dummy_file_id, 0)
+          if timeout:
+            start = time.time()
+            # Use polling. :/
+            while (self._dtrace.poll() is None and
+                   (time.time() - start) < timeout):
+              time.sleep(0.1)
+            self._dtrace.kill()
+          self._dtrace.wait()
+        finally:
+          # Make sure to kill it in any case.
+          if self._dtrace.poll() is None:
+            try:
+              self._dtrace.kill()
+              self._dtrace.wait()
+            except OSError:
+              pass
+
+        if self._dtrace.returncode != 0:
+          # Warn about any dtrace failure but basically ignore it.
+          print 'dtrace failure: %s' % self._dtrace.returncode
+      finally:
+        os.close(self._dummy_file_id)
+        os.remove(self._dummy_file_name)
 
     def post_process_log(self):
       """Sorts the log back in order when each call occured.
 
       dtrace doesn't save the buffer in strict order since it keeps one buffer
       per CPU.
       """
       super(Dtrace.Tracer, self).post_process_log()
       logname = self._logname + '.log'
       with open(logname, 'rb') as logfile:
@@ skipping 18 matching lines @@
     for ext in ('', '.log'):
       if os.path.isfile(logname + ext):
         os.remove(logname + ext)
 
   @classmethod
   def parse_log(cls, logname, blacklist):
     logging.info('parse_log(%s, %s)' % (logname, blacklist))
     data = read_json(logname)
     out = []
     for item in data['traces']:
-      context = cls.Context(blacklist)
+      context = cls.Context(blacklist, item['pid'], item['cwd'])
       for line in open(logname + '.log', 'rb'):
         context.on_line(line)
       out.append(
           {
             'results': context.to_results(),
             'trace': item['trace'],
             'output': item['output'],
           })
     return out
 
@@ skipping 909 matching lines @@
       for fn in dir(sys.modules[__name__])
       if fn.startswith('CMD')))
 
   command = get_command_handler(argv[0] if argv else None)
   parser = gen_parser(command)
   return command(parser, argv[1:])
 
 
 if __name__ == '__main__':
   sys.exit(main(sys.argv[1:]))