Home | History | Annotate | Download | only in site_utils 
      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