| # Copyright (C) 2019 Apple Inc. All rights reserved. |
| # |
| # Redistribution and use in source and binary forms, with or without |
| # modification, are permitted provided that the following conditions |
| # are met: |
| # 1. Redistributions of source code must retain the above copyright |
| # notice, this list of conditions and the following disclaimer. |
| # 2. Redistributions in binary form must reproduce the above copyright |
| # notice, this list of conditions and the following disclaimer in the |
| # documentation and/or other materials provided with the distribution. |
| # |
| # THIS SOFTWARE IS PROVIDED BY APPLE INC. AND ITS CONTRIBUTORS "AS IS" AND |
| # ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED |
| # WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE |
| # DISCLAIMED. IN NO EVENT SHALL APPLE INC. OR ITS CONTRIBUTORS BE LIABLE FOR |
| # ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL |
| # DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR |
| # SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER |
| # CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, |
| # OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE |
| # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. |
| |
| import atexit |
| import os |
| import re |
| import socket |
| import subprocess |
| import time |
| |
| from cassandra.cluster import Cluster, NoHostAvailable |
| |
| |
| class Docker(object): |
| |
| DOCKER_BIN = '/usr/local/bin/docker' |
| DOCKER_COMPOSE = '/usr/local/bin/docker-compose' |
| DEFAULT_PROJECT = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'docker-compose.yml') |
| PS_PORT = re.compile(r'0.0.0.0:(?P<port>\d+)\-\>\d+/tcp') |
| |
| _has_docker = None |
| _docker_stack = 0 |
| _project_stack = {} |
| _instance_for_project = {} |
| |
| @classmethod |
| def installed(cls): |
| if cls._has_docker is not None: |
| return cls._has_docker |
| |
| try: |
| with open(os.devnull, 'w') as devnull: |
| subprocess.check_call([cls.DOCKER_BIN, '-v'], stdout=devnull, stderr=devnull) |
| cls._has_docker = True |
| except (OSError, subprocess.CalledProcessError): |
| print('Docker not installed, required resources will be mocked. Download it at https://download.docker.com/mac/stable/Docker.dmg') |
| cls._has_docker = False |
| return cls._has_docker |
| |
| @classmethod |
| def is_running(cls): |
| try: |
| with open(os.devnull, 'w') as devnull: |
| subprocess.check_call([cls.DOCKER_BIN, 'container', 'ls', '-a'], stdout=devnull, stderr=devnull) |
| return True |
| except (OSError, subprocess.CalledProcessError): |
| return False |
| |
| @classmethod |
| def start(cls, timeout=60): |
| if cls.is_running(): |
| return |
| |
| if not cls.installed(): |
| raise RuntimeError('Docker not installed') |
| |
| with open(os.devnull, 'w') as devnull: |
| subprocess.check_call(['/usr/bin/open', '--background', '-a', 'Docker'], stdout=devnull, stderr=devnull) |
| |
| deadline = None if timeout is None else (time.time() + timeout) |
| while deadline is None or time.time() < deadline: |
| if cls.is_running(): |
| return |
| time.sleep(3) |
| raise RuntimeError('Docker failed to start') |
| |
| @classmethod |
| def stop(cls): |
| if not cls.is_running(): |
| return |
| with open(os.devnull, 'w') as devnull: |
| try: |
| subprocess.check_call(['/usr/bin/killall', 'Docker'], stdout=devnull, stderr=devnull) |
| except subprocess.CalledProcessError: |
| raise RuntimeError('Failed to shutdown Docker, maybe it was in the process of quiting?') |
| |
| @classmethod |
| def start_project(cls, project): |
| project = os.path.abspath(project) |
| if not os.path.isfile(project): |
| raise RuntimeError(f'Cannot find Docker project: {project}') |
| if not cls.is_running(): |
| raise RuntimeError('Cannot start project if Docker is not running') |
| |
| with open(os.devnull, 'w') as devnull: |
| subprocess.check_call([cls.DOCKER_COMPOSE, '-f', project, 'up', '-d', '--quiet-pull', '--no-color'], stdout=devnull, stderr=devnull) |
| |
| @classmethod |
| def wait_for_project(cls, project): |
| project = os.path.abspath(project) |
| if not os.path.isfile(project): |
| raise RuntimeError(f'Cannot find Docker project: {project}') |
| if not cls.is_running(): |
| raise RuntimeError(f'Docker is not running, cannot wait for {project}') |
| |
| has_cassandra = False |
| ports = [] |
| with open(os.devnull, 'w') as devnull: |
| output = subprocess.check_output([cls.DOCKER_COMPOSE, '-f', project, 'ps'], stderr=devnull).decode('utf-8') |
| if len(output.splitlines()) <= 2: |
| raise RuntimeError(f'{project} has not been started, cannot wait for it') |
| for line in output.splitlines()[2:]: |
| if 'cassandra' in line: |
| has_cassandra = True |
| for match in cls.PS_PORT.findall(line): |
| ports.append(int(match)) |
| |
| while True: |
| all_ports_open = True |
| for port in ports: |
| soc = socket.socket(socket.AF_INET, socket.SOCK_STREAM) |
| try: |
| soc.connect(('localhost', port)) |
| soc.shutdown(2) |
| continue |
| except BaseException: |
| all_ports_open = False |
| break |
| if all_ports_open: |
| try: |
| if has_cassandra: |
| connection = Cluster(['localhost']).connect() |
| connection.cluster.shutdown() |
| return |
| except NoHostAvailable: |
| pass |
| time.sleep(3) |
| |
| @classmethod |
| def is_project_running(cls, project): |
| project = os.path.abspath(project) |
| if not os.path.isfile(project): |
| raise RuntimeError(f'Cannot find Docker project: {project}') |
| |
| if not cls.is_running(): |
| return False |
| with open(os.devnull, 'w') as devnull: |
| output = subprocess.check_output([cls.DOCKER_COMPOSE, '-f', project, 'ps'], stderr=devnull) |
| return len(output.splitlines()) > 2 |
| |
| @classmethod |
| def stop_project(cls, project): |
| project = os.path.abspath(project) |
| if not os.path.isfile(project): |
| raise RuntimeError(f'Cannot find Docker project: {project}') |
| if not cls.is_running(): |
| raise RuntimeError('Cannot stop project if Docker is not running') |
| |
| with open(os.devnull, 'w') as devnull: |
| subprocess.check_call([cls.DOCKER_COMPOSE, '-f', project, 'down'], stdout=devnull, stderr=devnull) |
| |
| @classmethod |
| def instance(cls, project=DEFAULT_PROJECT): |
| project = os.path.abspath(project) |
| if project not in cls._instance_for_project: |
| cls._instance_for_project[project] = cls(project) |
| cls._instance_for_project[project].__enter__() |
| atexit.register(cls._instance_for_project[project].__exit__) |
| return cls._instance_for_project[project] |
| |
| def __init__(self, project=DEFAULT_PROJECT): |
| if not self.installed(): |
| raise RuntimeError('Cannot manage a Docker instance if Docker is not installed') |
| self.project = os.path.abspath(project) |
| if not os.path.isfile(self.project): |
| raise RuntimeError(f'Cannot find Docker project: {self.project}') |
| |
| def __enter__(self): |
| if Docker._docker_stack == 0: |
| if Docker.is_running(): |
| Docker._docker_stack += 1 |
| else: |
| Docker.start() |
| Docker._docker_stack += 1 |
| |
| if Docker._project_stack.get(self.project, 0) == 0: |
| if Docker.is_project_running(self.project): |
| Docker._project_stack[self.project] = Docker._project_stack.get(self.project, 0) + 1 |
| else: |
| Docker.start_project(self.project) |
| Docker.wait_for_project(self.project) |
| Docker._project_stack[self.project] = Docker._project_stack.get(self.project, 0) + 1 |
| |
| def __exit__(self, *args): |
| Docker._project_stack[self.project] = Docker._project_stack.get(self.project, 0) - 1 |
| if Docker._project_stack[self.project] <= 0: |
| Docker.stop_project(self.project) |
| |
| Docker._docker_stack -= 1 |
| if Docker._docker_stack <= 0: |
| Docker.stop() |
