1 #!/usr/bin/python 2 # Copyright 2017 The Chromium 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 """This tool manages the lxc container pool service.""" 7 8 import argparse 9 import logging 10 import os 11 import signal 12 import time 13 from contextlib import contextmanager 14 15 import common 16 from autotest_lib.client.bin import utils 17 from autotest_lib.client.common_lib import logging_config 18 from autotest_lib.server import server_logging_config 19 from autotest_lib.site_utils import lxc 20 from autotest_lib.site_utils.lxc import container_pool 21 22 try: 23 from chromite.lib import ts_mon_config 24 except ImportError: 25 ts_mon_config = utils.metrics_mock 26 27 28 # Location and base name of log files. 29 _LOG_LOCATION = '/usr/local/autotest/logs' 30 _LOG_NAME = 'lxc_pool.%d' % time.time() 31 32 33 def _start(args): 34 """Starts up the container pool service. 35 36 This function instantiates and starts up the pool service on the current 37 thread (i.e. the function will block, and not return until the service is 38 shut down). 39 """ 40 # TODO(dshi): crbug.com/459344 Set remove this enforcement when test 41 # container can be unprivileged container. 42 if utils.sudo_require_password(): 43 logging.warning('SSP requires root privilege to run commands, please ' 44 'grant root access to this process.') 45 utils.run('sudo true') 46 47 # Configure logging. 48 config = server_logging_config.ServerLoggingConfig() 49 config.configure_logging(verbose=args.verbose) 50 config.add_debug_file_handlers(log_dir=_LOG_LOCATION, log_name=_LOG_NAME) 51 # Pool code is heavily multi-threaded. This will help debugging. 52 logging_config.add_threadname_in_log() 53 54 host_dir = lxc.SharedHostDir() 55 service = container_pool.Service(host_dir) 56 # Catch signals, and send the appropriate stop request to the service 57 # instead of killing the main thread. 58 # - SIGINT is generated by Ctrl-C 59 # - SIGTERM is generated by an upstart stopping event. 60 for sig in (signal.SIGINT, signal.SIGTERM): 61 signal.signal(sig, lambda s, f: service.stop()) 62 63 with ts_mon_config.SetupTsMonGlobalState(service_name='lxc_pool_service', 64 indirect=True, 65 short_lived=False): 66 # Start the service. This blocks and does not return till the service 67 # shuts down. 68 service.start(pool_size=args.size) 69 70 71 def _status(_args): 72 """Requests status from the running container pool. 73 74 The retrieved status is printed out via logging. 75 """ 76 with _create_client() as client: 77 logging.debug('Requesting status...') 78 logging.info(client.get_status()) 79 80 81 def _stop(_args): 82 """Shuts down the running container pool.""" 83 with _create_client() as client: 84 logging.debug('Requesting stop...') 85 logging.info(client.shutdown()) 86 87 88 @contextmanager 89 # TODO(kenobi): Don't hard-code the timeout. 90 def _create_client(timeout=3): 91 logging.debug('Creating client...') 92 address = os.path.join(lxc.SharedHostDir().path, 93 lxc.DEFAULT_CONTAINER_POOL_SOCKET) 94 with container_pool.Client.connect(address, timeout) as connection: 95 yield connection 96 97 98 def parse_args(): 99 """Parse command line inputs. 100 101 @raise argparse.ArgumentError: If command line arguments are invalid. 102 """ 103 parser = argparse.ArgumentParser() 104 105 parser.add_argument('-v', '--verbose', 106 help='Enable verbose output.', 107 action='store_true') 108 109 subparsers = parser.add_subparsers(title='Commands') 110 111 parser_start = subparsers.add_parser('start', 112 help='Start the LXC container pool.') 113 parser_start.set_defaults(func=_start) 114 parser_start.add_argument('--size', 115 type=int, 116 default=lxc.DEFAULT_CONTAINER_POOL_SIZE, 117 help='Pool size (default=%d)' % 118 lxc.DEFAULT_CONTAINER_POOL_SIZE) 119 120 parser_stop = subparsers.add_parser('stop', 121 help='Stop the container pool.') 122 parser_stop.set_defaults(func=_stop) 123 124 parser_status = subparsers.add_parser('status', 125 help='Query pool status.') 126 parser_status.set_defaults(func=_status) 127 128 options = parser.parse_args() 129 return options 130 131 132 def main(): 133 """Main function.""" 134 # Parse args 135 args = parse_args() 136 137 # Dispatch control to the appropriate helper. 138 args.func(args) 139 140 141 if __name__ == '__main__': 142 main() 143