paasta-tools 1.21.3__py3-none-any.whl

paasta_tools/docker_wrapper.py

@@ -0,0 +1,234 @@
+#!/usr/bin/env python
+# flake8: noqa: E402
+""" Meant to be used by mesos-slave instead of the /usr/bin/docker executable
+directly This will parse the CLI arguments intended for docker, extract
+environment variable settings related to the actual node hostname and mesos
+task ID, and use those as an additional --hostname argument when calling the
+underlying docker command.
+
+If the environment variables are unspecified, or if --hostname is already
+specified, this does not change any arguments and just directly calls docker
+as-is.
+"""
+import logging
+import os
+import re
+import socket
+import sys
+
+
+if "PATH" not in os.environ:
+    # This command is sometimes executed in a sanitized environment
+    # which does not have the path, which causes the following imports
+    # to fail.
+    # To compensate, we set a minimal path to get off the ground.
+    os.environ["PATH"] = "/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin"
+
+
+LOCK_DIRECTORY = "/var/lib/paasta/mac-address"
+ENV_MATCH_RE = re.compile(r"^(-\w*e\w*|--env(?P<file>-file)?)(=(?P<arg>\S.*))?$")
+MAX_HOSTNAME_LENGTH = 60
+
+
+def parse_env_args(args):
+    result = dict(os.environ.items())
+    in_env = False
+    in_file = False
+    for arg in args:
+        if not in_env:
+            match = ENV_MATCH_RE.match(arg)
+            if not match:
+                continue
+            arg = match.group("arg") or ""
+            in_file = bool(match.group("file"))
+            if not arg:
+                in_env = True
+                continue
+
+        in_env = False
+
+        if in_file:
+            result.update(read_env_file(arg))
+            in_file = False
+            continue
+
+        try:
+            k, v = arg.split("=", 1)
+        except ValueError:
+            continue
+
+        result[k] = v
+
+    return result
+
+
+def read_env_file(filename):
+    # Parse a file where each line is KEY=VALUE
+    # return contents in dictionary form
+    result = {}
+    with open(filename) as f:
+        for line in f:
+            try:
+                k, v = line.split("=", 1)
+            except ValueError:
+                continue
+            result[k] = v.strip()
+    return result
+
+
+def can_add_hostname(args):
+    # return False if --hostname is already specified or if --network=host
+    if is_network_host(args):
+        return False
+
+    for index, arg in enumerate(args):
+
+        # Check for --hostname and variants
+        if arg == "-h":
+            return False
+        if arg.startswith("--hostname"):
+            return False
+        if len(arg) > 1 and arg[0] == "-" and arg[1] != "-":
+            # several short args
+            arg = arg.partition("=")[0]
+            if "h" in arg:
+                return False
+
+    return True
+
+
+def is_network_host(args):
+    for index, arg in enumerate(args):
+        # Check for --network=host and variants
+        if arg in ("--net=host", "--network=host"):
+            return True
+        try:
+            if arg in ("--net", "--network") and args[index + 1] == "host":
+                return True
+        except IndexError:
+            pass
+
+    return False
+
+
+def is_run(args):
+    try:
+        list(args).index("run")
+        return True
+    except ValueError:
+        return False
+
+
+def can_add_mac_address(args):
+    # return False if --mac-address is already specified or if --network=host
+    if is_network_host(args) or not is_run(args):
+        return False
+
+    for index, arg in enumerate(args):
+        # Check for --mac-address
+        if arg.startswith("--mac-address"):
+            return False
+
+    return True
+
+
+def generate_hostname_task_id(hostname, mesos_task_id):
+    task_id = mesos_task_id.rpartition(".")[2]
+
+    hostname_task_id = hostname + "-" + task_id
+
+    # hostnames can only contain alphanumerics and dashes and must be no more
+    # than 60 characters
+    hostname_task_id = re.sub("[^a-zA-Z0-9-]+", "-", hostname_task_id)[
+        :MAX_HOSTNAME_LENGTH
+    ]
+
+    # hostnames can also not end with dashes as per RFC952
+    hostname_task_id = hostname_task_id.rstrip("-")
+
+    return hostname_task_id
+
+
+def add_argument(args, argument):
+    # Add an argument immediately after 'run' command if it exists
+    args = list(args)
+    try:
+        run_index = args.index("run")
+    except ValueError:
+        pass
+    else:
+        args.insert(run_index + 1, argument)
+    return args
+
+
+def arg_collision(new_args, current_args):
+    # Returns True if one of the new arguments is already in the
+    # current argument list.
+    cur_arg_keys = []
+    for c in current_args:
+        cur_arg_keys.append(c.split("=")[0])
+    return bool(set(new_args).intersection(set(cur_arg_keys)))
+
+
+def add_firewall(argv, service, instance):
+    # Delayed import to improve performance when add_firewall is not used
+    from paasta_tools.docker_wrapper_imports import DEFAULT_SYNAPSE_SERVICE_DIR
+    from paasta_tools.docker_wrapper_imports import firewall_flock
+    from paasta_tools.docker_wrapper_imports import prepare_new_container
+    from paasta_tools.docker_wrapper_imports import reserve_unique_mac_address
+    from paasta_tools.docker_wrapper_imports import DEFAULT_SOA_DIR
+
+    output = ""
+    try:
+        mac_address, lockfile = reserve_unique_mac_address(LOCK_DIRECTORY)
+    except Exception as e:
+        output = f"Unable to add mac address: {e}"
+    else:
+        argv = add_argument(argv, f"--mac-address={mac_address}")
+        try:
+
+            with firewall_flock():
+                prepare_new_container(
+                    DEFAULT_SOA_DIR,
+                    DEFAULT_SYNAPSE_SERVICE_DIR,
+                    service,
+                    instance,
+                    mac_address,
+                )
+        except Exception as e:
+            output = f"Unable to add firewall rules: {e}"
+
+    if output:
+        print(output, file=sys.stderr)
+
+    return argv
+
+
+def main(argv=None):
+    argv = argv if argv is not None else sys.argv
+
+    env_args = parse_env_args(argv)
+
+    # Marathon sets MESOS_TASK_ID
+    mesos_task_id = env_args.get("MESOS_TASK_ID")
+
+    fqdn = socket.getfqdn()
+    hostname = fqdn.partition(".")[0]
+    if mesos_task_id and can_add_hostname(argv):
+        argv = add_argument(argv, f"-e=PAASTA_HOST={fqdn}")
+        hostname_task_id = generate_hostname_task_id(hostname, mesos_task_id)
+        argv = add_argument(argv, f"--hostname={hostname_task_id }")
+    elif can_add_hostname(argv):
+        argv = add_argument(argv, f"-e=PAASTA_HOST={fqdn}")
+        argv = add_argument(argv, f"--hostname={hostname}")
+
+    paasta_firewall = env_args.get("PAASTA_FIREWALL")
+    service = env_args.get("PAASTA_SERVICE")
+    instance = env_args.get("PAASTA_INSTANCE")
+    if paasta_firewall and service and instance and can_add_mac_address(argv):
+        try:
+            argv = add_firewall(argv, service, instance)
+        except Exception as e:
+            print(f"Unhandled exception in add_firewall: {e}", file=sys.stderr)
+
+    os.execlp("docker", "docker", *argv[1:])

paasta_tools/docker_wrapper_imports.py

@@ -0,0 +1,13 @@
+#!/usr/bin/env python
+# flake8: noqa: E402
+""" Delayed imports for paasta_tools.docker_wrapper, meant to improve the performance
+of docker_wrapper for execution paths when we do not use add_firewall. It turns out
+that the imports needed add a fair amount of overhead to "docker inspect" command
+which does not use add_firewall, and Mesos executes this command quite a lot. We can
+make it faster by moving all the paasta_tools imports to separate file, i.e. here
+"""
+from paasta_tools.firewall import DEFAULT_SYNAPSE_SERVICE_DIR
+from paasta_tools.firewall import firewall_flock
+from paasta_tools.firewall import prepare_new_container
+from paasta_tools.mac_address import reserve_unique_mac_address
+from paasta_tools.utils import DEFAULT_SOA_DIR

paasta_tools/drain_lib.py

@@ -0,0 +1,351 @@
+# Copyright 2015-2016 Yelp Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import asyncio
+import time
+from typing import Any
+from typing import Awaitable
+from typing import Callable
+from typing import Dict
+from typing import List
+from typing import Optional
+from typing import Sequence
+from typing import Set
+from typing import Type
+from typing import TypeVar
+
+import aiohttp
+from mypy_extensions import TypedDict
+
+from paasta_tools.hacheck import get_spool
+from paasta_tools.hacheck import post_spool
+from paasta_tools.utils import get_user_agent
+
+_drain_methods: Dict[str, Type["DrainMethod"]] = {}
+
+
+_RegisterDrainMethod_T = TypeVar("_RegisterDrainMethod_T", bound=Type["DrainMethod"])
+T = TypeVar("T")
+
+
+def register_drain_method(
+    name: str,
+) -> Callable[[_RegisterDrainMethod_T], _RegisterDrainMethod_T]:
+    """Returns a decorator that registers a DrainMethod subclass at a given name
+    so get_drain_method/list_drain_methods can find it."""
+
+    def outer(drain_method: _RegisterDrainMethod_T) -> _RegisterDrainMethod_T:
+        _drain_methods[name] = drain_method
+        return drain_method
+
+    return outer
+
+
+def get_drain_method(
+    name: str,
+    service: str,
+    instance: str,
+    registrations: List[str],
+    drain_method_params: Optional[Dict] = None,
+) -> "DrainMethod":
+    return _drain_methods[name](
+        service, instance, registrations, **(drain_method_params or {})
+    )
+
+
+def list_drain_methods() -> List[str]:
+    return sorted(_drain_methods.keys())
+
+
+DrainTask = TypeVar("DrainTask", bound=Any)
+
+
+class DrainMethod:
+    """A drain method is a way of stopping new traffic to tasks without killing them. For example, you might take a task
+    out of a load balancer by causing its healthchecks to fail.
+
+    A drain method must have the following methods:
+     - drain(task): Begin draining traffic from a task. This should be idempotent.
+     - stop_draining(task): Stop draining traffic from a task. This should be idempotent.
+     - is_draining(task): Whether a task has already been marked as downed. Note that this state should be stored out of
+                          process, because a bounce may take multiple runs of setup_marathon_job to complete.
+     - is_safe_to_kill(task): Return True if this task is safe to kill, False otherwise.
+
+    When implementing a drain method, be sure to decorate with @register_drain_method(name).
+    """
+
+    def __init__(
+        self, service: str, instance: str, registrations: List[str], **kwargs: Dict
+    ) -> None:
+        self.service = service
+        self.instance = instance
+        self.registrations = registrations
+
+    async def drain(self, task: DrainTask) -> None:
+        """Make a task stop receiving new traffic."""
+        raise NotImplementedError()
+
+    async def stop_draining(self, task: DrainTask) -> None:
+        """Make a task that has previously been downed start receiving traffic again."""
+        raise NotImplementedError()
+
+    async def is_draining(self, task: DrainTask) -> bool:
+        """Return whether a task is being drained."""
+        raise NotImplementedError()
+
+    async def is_safe_to_kill(self, task: DrainTask) -> bool:
+        """Return True if a task is drained and ready to be killed, or False if we should wait."""
+        raise NotImplementedError()
+
+
+@register_drain_method("noop")
+class NoopDrainMethod(DrainMethod):
+    """This drain policy does nothing and assumes every task is safe to kill."""
+
+    async def drain(self, task: DrainTask) -> None:
+        pass
+
+    async def stop_draining(self, task: DrainTask) -> None:
+        pass
+
+    async def is_draining(self, task: DrainTask) -> bool:
+        return False
+
+    async def is_safe_to_kill(self, task: DrainTask) -> bool:
+        return True
+
+
+@register_drain_method("test")
+class TestDrainMethod(DrainMethod):
+    """This drain policy is meant for integration testing. Do not use."""
+
+    # These are variables on the class for ease of use in testing.
+    downed_task_ids: Set[str] = set()
+    safe_to_kill_task_ids: Set[str] = set()
+
+    async def drain(self, task: DrainTask) -> None:
+        if task.id not in self.safe_to_kill_task_ids:
+            self.downed_task_ids.add(task.id)
+
+    async def stop_draining(self, task: DrainTask) -> None:
+        self.downed_task_ids -= {task.id}
+        self.safe_to_kill_task_ids -= {task.id}
+
+    async def is_draining(self, task: DrainTask) -> bool:
+        return task.id in (self.downed_task_ids | self.safe_to_kill_task_ids)
+
+    async def is_safe_to_kill(self, task: DrainTask) -> bool:
+        return task.id in self.safe_to_kill_task_ids
+
+    @classmethod
+    def mark_arbitrary_task_as_safe_to_kill(cls) -> None:
+        cls.safe_to_kill_task_ids.add(cls.downed_task_ids.pop())
+
+
+@register_drain_method("crashy_drain")
+class CrashyDrainDrainMethod(NoopDrainMethod):
+    async def drain(self, task: DrainTask) -> None:
+        raise Exception("Intentionally crashing for testing purposes")
+
+    async def is_safe_to_kill(self, task: DrainTask) -> bool:
+        raise Exception("Intentionally crashing for testing purposes")
+
+    async def stop_draining(self, task: DrainTask) -> None:
+        raise Exception("Intentionally crashing for testing purposes")
+
+    async def is_draining(self, task: DrainTask) -> bool:
+        raise Exception("Intentionally crashing for testing purposes")
+
+
+@register_drain_method("hacheck")
+class HacheckDrainMethod(DrainMethod):
+    """This drain policy issues a POST to hacheck's /spool/{service}/{port}/status endpoint to cause healthchecks to
+    fail. It considers tasks safe to kill if they've been down in hacheck for more than a specified delay."""
+
+    def __init__(
+        self,
+        service: str,
+        instance: str,
+        registrations: List[str],
+        delay: float = 240,
+        hacheck_port: int = 6666,
+        expiration: float = 0,
+        **kwargs: Dict,
+    ) -> None:
+        super().__init__(service, instance, registrations)
+        self.delay = float(delay)
+        self.hacheck_port = hacheck_port
+        self.expiration = float(expiration) or float(delay) * 10
+
+    def spool_urls(self, task: DrainTask) -> List[str]:
+        return [
+            "http://%(task_host)s:%(hacheck_port)d/spool/%(registration)s/%(task_port)d/status"
+            % {
+                "task_host": task.host,
+                "task_port": task.ports[0],
+                "hacheck_port": self.hacheck_port,
+                "registration": registration,
+            }
+            for registration in self.registrations
+        ]
+
+    async def for_each_registration(
+        self, task: DrainTask, func: Callable[..., Awaitable[T]]
+    ) -> Sequence[T]:
+        if task.ports == []:
+            return None
+        futures = [func(url) for url in self.spool_urls(task)]
+        return await asyncio.gather(*futures)
+
+    async def drain(self, task: DrainTask) -> None:
+        await self.for_each_registration(task, self.down)
+
+    async def up(self, url: str) -> None:
+        await post_spool(url=url, status="up", data={"status": "up"})
+
+    async def down(self, url: str) -> None:
+        await post_spool(
+            url=url,
+            status="down",
+            data={
+                "status": "down",
+                "expiration": str(time.time() + self.expiration),
+                "reason": "Drained by Paasta",
+            },
+        )
+
+    async def stop_draining(self, task: DrainTask) -> None:
+        await self.for_each_registration(task, self.up)
+
+    async def is_draining(self, task: DrainTask) -> bool:
+        results = await self.for_each_registration(task, get_spool)
+        return not all([res is None or res["state"] == "up" for res in results])
+
+    async def is_safe_to_kill(self, task: DrainTask) -> bool:
+        results = await self.for_each_registration(task, lambda url: get_spool(url))
+        if all([res is None or res["state"] == "up" for res in results]):
+            return False
+        else:
+            return all(
+                [res.get("since", 0) < (time.time() - self.delay) for res in results]
+            )
+
+
+class StatusCodeNotAcceptableError(Exception):
+    pass
+
+
+UrlSpec = TypedDict(
+    "UrlSpec", {"url_format": str, "method": str, "success_codes": str}, total=False
+)
+
+
+@register_drain_method("http")
+class HTTPDrainMethod(DrainMethod):
+    """This drain policy issues arbitrary HTTP calls to arbitrary URLs specified by the parameters. The URLs are
+    specified as format strings, and will have variables such as {host}, {port}, etc. filled in."""
+
+    def __init__(
+        self,
+        service: str,
+        instance: str,
+        registrations: List[str],
+        drain: UrlSpec,
+        stop_draining: UrlSpec,
+        is_draining: UrlSpec,
+        is_safe_to_kill: UrlSpec,
+    ) -> None:
+        super().__init__(service, instance, registrations)
+        self.drain_url_spec = drain
+        self.stop_draining_url_spec = stop_draining
+        self.is_draining_url_spec = is_draining
+        self.is_safe_to_kill_url_spec = is_safe_to_kill
+
+    def get_format_params(self, task: DrainTask) -> List[Dict[str, Any]]:
+        return [
+            {
+                "host": task.host,
+                "port": task.ports[0],
+                "service": self.service,
+                "instance": self.instance,
+                "nerve_ns": nerve_ns,
+            }
+            for nerve_ns in self.registrations
+        ]
+
+    def format_url(self, url_format: str, format_params: Dict[str, Any]) -> str:
+        return url_format.format(**format_params)
+
+    def parse_success_codes(self, success_codes_str: str) -> Set[int]:
+        """Expand a string like 200-399,407-409,500 to a set containing all the integers in between."""
+        acceptable_response_codes: Set[int] = set()
+        for series_str in str(success_codes_str).split(","):
+            if "-" in series_str:
+                start, end = series_str.split("-")
+                acceptable_response_codes.update(range(int(start), int(end) + 1))
+            else:
+                acceptable_response_codes.add(int(series_str))
+        return acceptable_response_codes
+
+    def check_response_code(self, status_code: int, success_codes_str: str) -> bool:
+        acceptable_response_codes = self.parse_success_codes(success_codes_str)
+        return status_code in acceptable_response_codes
+
+    async def issue_request(self, url_spec: UrlSpec, task: DrainTask) -> None:
+        """Issue a request to the URL specified by url_spec regarding the task given."""
+        format_params = self.get_format_params(task)
+        urls = [
+            self.format_url(url_spec["url_format"], param) for param in format_params
+        ]
+        method = url_spec.get("method", "GET").upper()
+
+        async with aiohttp.ClientSession() as session:
+            reqs = [
+                session.request(
+                    method=method,
+                    url=url,
+                    headers={"User-Agent": get_user_agent()},
+                    timeout=15,
+                )
+                for url in urls
+            ]
+            res = await asyncio.gather(*reqs)
+            for response in res:
+                if not self.check_response_code(
+                    response.status, url_spec["success_codes"]
+                ):
+                    raise StatusCodeNotAcceptableError(
+                        f"Unacceptable status code {response.status} not in {url_spec['success_codes']} when hitting {response.url}"
+                    )
+
+    async def drain(self, task: DrainTask) -> None:
+        return await self.issue_request(self.drain_url_spec, task)
+
+    async def stop_draining(self, task: DrainTask) -> None:
+        return await self.issue_request(self.stop_draining_url_spec, task)
+
+    async def is_draining(self, task: DrainTask) -> bool:
+        try:
+            await self.issue_request(self.is_draining_url_spec, task)
+        except StatusCodeNotAcceptableError:
+            return False
+        else:
+            return True
+
+    async def is_safe_to_kill(self, task: DrainTask) -> bool:
+        try:
+            await self.issue_request(self.is_safe_to_kill_url_spec, task)
+        except StatusCodeNotAcceptableError:
+            return False
+        else:
+            return True

paasta_tools/dump_locally_running_services.py

@@ -0,0 +1,71 @@
+#!/usr/bin/env python
+# Copyright 2015-2019 Yelp Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+Usage: ./paasta_dump_locally_running_services.py [options]
+
+Outputs a JSON-encoded list of services that are running on this host along
+with the host port that each service is listening on.
+
+Command line options:
+
+- -d <SOA_DIR>, --soa-dir <SOA_DIR>: Specify a SOA config dir to read from
+"""
+import argparse
+import json
+import sys
+from typing import List
+from typing import Optional
+from typing import Sequence
+from typing import Tuple
+
+from paasta_tools.kubernetes_tools import get_kubernetes_services_running_here_for_nerve
+from paasta_tools.long_running_service_tools import ServiceNamespaceConfig
+from paasta_tools.puppet_service_tools import get_puppet_services_running_here_for_nerve
+from paasta_tools.utils import DEFAULT_SOA_DIR
+
+
+def parse_args(argv: Optional[Sequence[str]]) -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description="Dumps information about locally running services."
+    )
+    parser.add_argument(
+        "-d",
+        "--soa-dir",
+        dest="soa_dir",
+        metavar="SOA_DIR",
+        default=DEFAULT_SOA_DIR,
+        help="define a different soa config directory",
+    )
+    return parser.parse_args(argv)
+
+
+def main(argv: Optional[Sequence[str]] = None) -> None:
+    args = parse_args(argv)
+    soa_dir = args.soa_dir
+
+    service_dump: List[
+        Tuple[str, ServiceNamespaceConfig]
+    ] = get_puppet_services_running_here_for_nerve(
+        soa_dir=soa_dir
+    ) + get_kubernetes_services_running_here_for_nerve(
+        cluster=None, soa_dir=soa_dir
+    )
+
+    print(json.dumps(service_dump))
+    sys.exit(0)
+
+
+if __name__ == "__main__":
+    main()