• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
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'''
7A library to prespawn autotest processes to minimize startup overhead.
8'''
9
10import cPickle as pickle, os, sys
11from setproctitle import setproctitle
12
13
14if 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
42import logging, subprocess, threading
43from Queue import Queue
44
45
46NUM_PRESPAWNED_PROCESSES = 1
47
48
49class 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