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