Home | History | Annotate | Download | only in bin 
      1 # Copyright (c) 2012 The Chromium OS Authors. All rights reserved.
      2 # Use of this source code is governed by a BSD-style license that can be
      3 # found in the LICENSE file.
      4 
      5 
      6 '''
      7 A library to prespawn autotest processes to minimize startup overhead.
      8 '''
      9 
     10 import cPickle as pickle, os, sys
     11 from setproctitle import setproctitle
     12 
     13 
     14 if len(sys.argv) == 2 and sys.argv[1] == '--prespawn_autotest':
     15     # Run an autotest process, and on stdin, wait for a pickled environment +
     16     # argv (as a tuple); see spawn() below.  Once we receive these, start
     17     # autotest.
     18 
     19     # Do common imports (to save startup time).
     20     # pylint: disable=W0611
     21     import common
     22     import autotest_lib.client.bin.job
     23 
     24     if os.environ.get('CROS_DISABLE_SITE_SYSINFO'):
     25         from autotest_lib.client.bin import sysinfo, base_sysinfo
     26         sysinfo.sysinfo = autotest_lib.client.bin.base_sysinfo.base_sysinfo
     27 
     28     # Wait for environment and autotest arguments.
     29     env, sys.argv = pickle.load(sys.stdin)
     30     # Run autotest and exit.
     31     if env:
     32         os.environ.clear()
     33         os.environ.update(env)
     34         proc_title = os.environ.get('CROS_PROC_TITLE')
     35         if proc_title:
     36             setproctitle(proc_title)
     37 
     38         execfile('autotest')
     39     sys.exit(0)
     40 
     41 
     42 import logging, subprocess, threading
     43 from Queue import Queue
     44 
     45 
     46 NUM_PRESPAWNED_PROCESSES = 1
     47 
     48 
     49 class Prespawner():
     50     def __init__(self):
     51         self.prespawned = Queue(NUM_PRESPAWNED_PROCESSES)
     52         self.thread = None
     53         self.terminated = False
     54 
     55     def spawn(self, args, env_additions=None):
     56         '''
     57         Spawns a new autotest (reusing an prespawned process if available).
     58 
     59         @param args: A list of arguments (sys.argv)
     60         @param env_additions: Items to add to the current environment
     61         '''
     62         new_env = dict(os.environ)
     63         if env_additions:
     64             new_env.update(env_additions)
     65 
     66         process = self.prespawned.get()
     67         # Write the environment and argv to the process's stdin; it will launch
     68         # autotest once these are received.
     69         pickle.dump((new_env, args), process.stdin, protocol=2)
     70         process.stdin.close()
     71         return process
     72 
     73     def start(self):
     74         '''
     75         Starts a thread to pre-spawn autotests.
     76         '''
     77         def run():
     78             while not self.terminated:
     79                 process = subprocess.Popen(
     80                     ['python', '-u', os.path.realpath(__file__),
     81                      '--prespawn_autotest'],
     82                     cwd=os.path.dirname(os.path.realpath(__file__)),
     83                     stdin=subprocess.PIPE)
     84                 logging.debug('Pre-spawned an autotest process %d', process.pid)
     85                 self.prespawned.put(process)
     86 
     87             # Let stop() know that we are done
     88             self.prespawned.put(None)
     89 
     90         if not self.thread:
     91             self.thread = threading.Thread(target=run, name='Prespawner')
     92             self.thread.start()
     93 
     94     def stop(self):
     95         '''
     96         Stops the pre-spawn thread gracefully.
     97         '''
     98         if not self.thread:
     99             # Never started
    100             return
    101 
    102         self.terminated = True
    103         # Wait for any existing prespawned processes.
    104         while True:
    105             process = self.prespawned.get()
    106             if not process:
    107                 break
    108             # Send a 'None' environment and arg list to tell the prespawner
    109             # processes to exit.
    110             pickle.dump((None, None), process.stdin, protocol=2)
    111             process.stdin.close()
    112             process.wait()
    113         self.thread = None
    114