Ben Kwa | 5e2bb5c | 2017-11-01 16:44:16 -0700 | [diff] [blame] | 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 |
Ben Kwa | fe04bb3 | 2017-11-03 10:50:31 -0700 | [diff] [blame] | 12 | import time |
Ben Kwa | 5e2bb5c | 2017-11-01 16:44:16 -0700 | [diff] [blame] | 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 |
Ben Kwa | fe04bb3 | 2017-11-03 10:50:31 -0700 | [diff] [blame] | 18 | from autotest_lib.server import server_logging_config |
Ben Kwa | 5e2bb5c | 2017-11-01 16:44:16 -0700 | [diff] [blame] | 19 | from autotest_lib.site_utils import lxc |
| 20 | from autotest_lib.site_utils.lxc import container_pool |
| 21 | |
Ben Kwa | 9b18c71 | 2017-12-13 15:38:56 -0800 | [diff] [blame] | 22 | try: |
| 23 | from chromite.lib import ts_mon_config |
| 24 | except ImportError: |
| 25 | ts_mon_config = utils.metrics_mock |
| 26 | |
Ben Kwa | 5e2bb5c | 2017-11-01 16:44:16 -0700 | [diff] [blame] | 27 | |
Ben Kwa | fe04bb3 | 2017-11-03 10:50:31 -0700 | [diff] [blame] | 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 | |
Ben Kwa | 181814d | 2017-11-03 09:27:58 -0700 | [diff] [blame] | 33 | def _start(args): |
Ben Kwa | 5e2bb5c | 2017-11-01 16:44:16 -0700 | [diff] [blame] | 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') |
Ben Kwa | 5ab77c8 | 2017-12-01 13:39:40 -0800 | [diff] [blame] | 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 | |
Ben Kwa | 5e2bb5c | 2017-11-01 16:44:16 -0700 | [diff] [blame] | 54 | host_dir = lxc.SharedHostDir() |
| 55 | service = container_pool.Service(host_dir) |
Ben Kwa | 14e76e5 | 2017-11-03 15:18:25 -0700 | [diff] [blame] | 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 | |
Ben Kwa | 9b18c71 | 2017-12-13 15:38:56 -0800 | [diff] [blame] | 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) |
Ben Kwa | 5e2bb5c | 2017-11-01 16:44:16 -0700 | [diff] [blame] | 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...') |
Ben Kwa | 5e2bb5c | 2017-11-01 16:44:16 -0700 | [diff] [blame] | 92 | address = os.path.join(lxc.SharedHostDir().path, |
Ben Kwa | 181814d | 2017-11-03 09:27:58 -0700 | [diff] [blame] | 93 | lxc.DEFAULT_CONTAINER_POOL_SOCKET) |
Ben Kwa | 5e2bb5c | 2017-11-01 16:44:16 -0700 | [diff] [blame] | 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() |
Ben Kwa | 5e2bb5c | 2017-11-01 16:44:16 -0700 | [diff] [blame] | 104 | |
Ben Kwa | 181814d | 2017-11-03 09:27:58 -0700 | [diff] [blame] | 105 | parser.add_argument('-v', '--verbose', |
| 106 | help='Enable verbose output.', |
| 107 | action='store_true') |
Ben Kwa | 5e2bb5c | 2017-11-01 16:44:16 -0700 | [diff] [blame] | 108 | |
Ben Kwa | 181814d | 2017-11-03 09:27:58 -0700 | [diff] [blame] | 109 | subparsers = parser.add_subparsers(title='Commands') |
Ben Kwa | 5e2bb5c | 2017-11-01 16:44:16 -0700 | [diff] [blame] | 110 | |
Ben Kwa | 181814d | 2017-11-03 09:27:58 -0700 | [diff] [blame] | 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) |
Ben Kwa | 5e2bb5c | 2017-11-01 16:44:16 -0700 | [diff] [blame] | 127 | |
| 128 | options = parser.parse_args() |
| 129 | return options |
| 130 | |
| 131 | |
| 132 | def main(): |
| 133 | """Main function.""" |
Ben Kwa | 181814d | 2017-11-03 09:27:58 -0700 | [diff] [blame] | 134 | # Parse args |
| 135 | args = parse_args() |
| 136 | |
Ben Kwa | 181814d | 2017-11-03 09:27:58 -0700 | [diff] [blame] | 137 | # Dispatch control to the appropriate helper. |
Ben Kwa | 5e2bb5c | 2017-11-01 16:44:16 -0700 | [diff] [blame] | 138 | args.func(args) |
| 139 | |
| 140 | |
| 141 | if __name__ == '__main__': |
| 142 | main() |