grpc/tools/run_tests/xds_test_driver/framework/infrastructure/k8s.py

# Copyright 2020 gRPC authors.
#
# 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 functools
import json
import logging
import subprocess
import time
from typing import Optional, List, Tuple

# TODO(sergiitk): replace with tenacity
import retrying
import kubernetes.config
from kubernetes import client
from kubernetes import utils

logger = logging.getLogger(__name__)
# Type aliases
V1Deployment = client.V1Deployment
V1ServiceAccount = client.V1ServiceAccount
V1Pod = client.V1Pod
V1PodList = client.V1PodList
V1Service = client.V1Service
V1Namespace = client.V1Namespace
ApiException = client.ApiException


def simple_resource_get(func):

    def wrap_not_found_return_none(*args, **kwargs):
        try:
            return func(*args, **kwargs)
        except client.ApiException as e:
            if e.status == 404:
                # Ignore 404
                return None
            raise

    return wrap_not_found_return_none


def label_dict_to_selector(labels: dict) -> str:
    return ','.join(f'{k}=={v}' for k, v in labels.items())


class KubernetesApiManager:

    def __init__(self, context):
        self.context = context
        self.client = self._cached_api_client_for_context(context)
        self.apps = client.AppsV1Api(self.client)
        self.core = client.CoreV1Api(self.client)

    def close(self):
        self.client.close()

    @classmethod
    @functools.lru_cache(None)
    def _cached_api_client_for_context(cls, context: str) -> client.ApiClient:
        return kubernetes.config.new_client_from_config(context=context)


class PortForwardingError(Exception):
    """Error forwarding port"""


class KubernetesNamespace:
    NEG_STATUS_META = 'cloud.google.com/neg-status'
    PORT_FORWARD_LOCAL_ADDRESS: str = '127.0.0.1'
    DELETE_GRACE_PERIOD_SEC: int = 5

    def __init__(self, api: KubernetesApiManager, name: str):
        self.name = name
        self.api = api

    def apply_manifest(self, manifest):
        return utils.create_from_dict(self.api.client,
                                      manifest,
                                      namespace=self.name)

    @simple_resource_get
    def get_service(self, name) -> V1Service:
        return self.api.core.read_namespaced_service(name, self.name)

    @simple_resource_get
    def get_service_account(self, name) -> V1Service:
        return self.api.core.read_namespaced_service_account(name, self.name)

    def delete_service(self, name,
                       grace_period_seconds=DELETE_GRACE_PERIOD_SEC):
        self.api.core.delete_namespaced_service(
            name=name,
            namespace=self.name,
            body=client.V1DeleteOptions(
                propagation_policy='Foreground',
                grace_period_seconds=grace_period_seconds))

    def delete_service_account(self,
                               name,
                               grace_period_seconds=DELETE_GRACE_PERIOD_SEC):
        self.api.core.delete_namespaced_service_account(
            name=name,
            namespace=self.name,
            body=client.V1DeleteOptions(
                propagation_policy='Foreground',
                grace_period_seconds=grace_period_seconds))

    @simple_resource_get
    def get(self) -> V1Namespace:
        return self.api.core.read_namespace(self.name)

    def delete(self, grace_period_seconds=DELETE_GRACE_PERIOD_SEC):
        self.api.core.delete_namespace(
            name=self.name,
            body=client.V1DeleteOptions(
                propagation_policy='Foreground',
                grace_period_seconds=grace_period_seconds))

    def wait_for_service_deleted(self, name: str, timeout_sec=60, wait_sec=1):

        @retrying.retry(retry_on_result=lambda r: r is not None,
                        stop_max_delay=timeout_sec * 1000,
                        wait_fixed=wait_sec * 1000)
        def _wait_for_deleted_service_with_retry():
            service = self.get_service(name)
            if service is not None:
                logger.info('Waiting for service %s to be deleted',
                            service.metadata.name)
            return service

        _wait_for_deleted_service_with_retry()

    def wait_for_service_account_deleted(self,
                                         name: str,
                                         timeout_sec=60,
                                         wait_sec=1):

        @retrying.retry(retry_on_result=lambda r: r is not None,
                        stop_max_delay=timeout_sec * 1000,
                        wait_fixed=wait_sec * 1000)
        def _wait_for_deleted_service_account_with_retry():
            service_account = self.get_service_account(name)
            if service_account is not None:
                logger.info('Waiting for service account %s to be deleted',
                            service_account.metadata.name)
            return service_account

        _wait_for_deleted_service_account_with_retry()

    def wait_for_namespace_deleted(self, timeout_sec=240, wait_sec=2):

        @retrying.retry(retry_on_result=lambda r: r is not None,
                        stop_max_delay=timeout_sec * 1000,
                        wait_fixed=wait_sec * 1000)
        def _wait_for_deleted_namespace_with_retry():
            namespace = self.get()
            if namespace is not None:
                logger.info('Waiting for namespace %s to be deleted',
                            namespace.metadata.name)
            return namespace

        _wait_for_deleted_namespace_with_retry()

    def wait_for_service_neg(self, name: str, timeout_sec=60, wait_sec=1):

        @retrying.retry(retry_on_result=lambda r: not r,
                        stop_max_delay=timeout_sec * 1000,
                        wait_fixed=wait_sec * 1000)
        def _wait_for_service_neg():
            service = self.get_service(name)
            if self.NEG_STATUS_META not in service.metadata.annotations:
                logger.info('Waiting for service %s NEG', service.metadata.name)
                return False
            return True

        _wait_for_service_neg()

    def get_service_neg(self, service_name: str,
                        service_port: int) -> Tuple[str, List[str]]:
        service = self.get_service(service_name)
        neg_info: dict = json.loads(
            service.metadata.annotations[self.NEG_STATUS_META])
        neg_name: str = neg_info['network_endpoint_groups'][str(service_port)]
        neg_zones: List[str] = neg_info['zones']
        return neg_name, neg_zones

    @simple_resource_get
    def get_deployment(self, name) -> V1Deployment:
        return self.api.apps.read_namespaced_deployment(name, self.name)

    def delete_deployment(self,
                          name,
                          grace_period_seconds=DELETE_GRACE_PERIOD_SEC):
        self.api.apps.delete_namespaced_deployment(
            name=name,
            namespace=self.name,
            body=client.V1DeleteOptions(
                propagation_policy='Foreground',
                grace_period_seconds=grace_period_seconds))

    def list_deployment_pods(self, deployment: V1Deployment) -> List[V1Pod]:
        # V1LabelSelector.match_expressions not supported at the moment
        return self.list_pods_with_labels(deployment.spec.selector.match_labels)

    def wait_for_deployment_available_replicas(self,
                                               name,
                                               count=1,
                                               timeout_sec=60,
                                               wait_sec=1):

        @retrying.retry(
            retry_on_result=lambda r: not self._replicas_available(r, count),
            stop_max_delay=timeout_sec * 1000,
            wait_fixed=wait_sec * 1000)
        def _wait_for_deployment_available_replicas():
            deployment = self.get_deployment(name)
            logger.info(
                'Waiting for deployment %s to have %s available '
                'replicas, current count %s', deployment.metadata.name, count,
                deployment.status.available_replicas)
            return deployment

        _wait_for_deployment_available_replicas()

    def wait_for_deployment_deleted(self,
                                    deployment_name: str,
                                    timeout_sec=60,
                                    wait_sec=1):

        @retrying.retry(retry_on_result=lambda r: r is not None,
                        stop_max_delay=timeout_sec * 1000,
                        wait_fixed=wait_sec * 1000)
        def _wait_for_deleted_deployment_with_retry():
            deployment = self.get_deployment(deployment_name)
            if deployment is not None:
                logger.info(
                    'Waiting for deployment %s to be deleted. '
                    'Non-terminated replicas: %s', deployment.metadata.name,
                    deployment.status.replicas)
            return deployment

        _wait_for_deleted_deployment_with_retry()

    def list_pods_with_labels(self, labels: dict) -> List[V1Pod]:
        pod_list: V1PodList = self.api.core.list_namespaced_pod(
            self.name, label_selector=label_dict_to_selector(labels))
        return pod_list.items

    def get_pod(self, name) -> client.V1Pod:
        return self.api.core.read_namespaced_pod(name, self.name)

    def wait_for_pod_started(self, pod_name, timeout_sec=60, wait_sec=1):

        @retrying.retry(retry_on_result=lambda r: not self._pod_started(r),
                        stop_max_delay=timeout_sec * 1000,
                        wait_fixed=wait_sec * 1000)
        def _wait_for_pod_started():
            pod = self.get_pod(pod_name)
            logger.info('Waiting for pod %s to start, current phase: %s',
                        pod.metadata.name, pod.status.phase)
            return pod

        _wait_for_pod_started()

    def port_forward_pod(
            self,
            pod: V1Pod,
            remote_port: int,
            local_port: Optional[int] = None,
            local_address: Optional[str] = None,
    ) -> subprocess.Popen:
        """Experimental"""
        local_address = local_address or self.PORT_FORWARD_LOCAL_ADDRESS
        local_port = local_port or remote_port
        cmd = [
            "kubectl", "--context", self.api.context, "--namespace", self.name,
            "port-forward", "--address", local_address,
            f"pod/{pod.metadata.name}", f"{local_port}:{remote_port}"
        ]
        pf = subprocess.Popen(cmd,
                              stdout=subprocess.PIPE,
                              stderr=subprocess.STDOUT,
                              universal_newlines=True)
        # Wait for stdout line indicating successful start.
        expected = (f"Forwarding from {local_address}:{local_port}"
                    f" -> {remote_port}")
        try:
            while True:
                time.sleep(0.05)
                output = pf.stdout.readline().strip()
                if not output:
                    return_code = pf.poll()
                    if return_code is not None:
                        errors = [error for error in pf.stdout.readlines()]
                        raise PortForwardingError(
                            'Error forwarding port, kubectl return '
                            f'code {return_code}, output {errors}')
                elif output != expected:
                    raise PortForwardingError(
                        f'Error forwarding port, unexpected output {output}')
                else:
                    logger.info(output)
                    break
        except Exception:
            self.port_forward_stop(pf)
            raise

        # TODO(sergiitk): return new PortForwarder object
        return pf

    @staticmethod
    def port_forward_stop(pf):
        logger.info('Shutting down port forwarding, pid %s', pf.pid)
        pf.kill()
        stdout, _stderr = pf.communicate(timeout=5)
        logger.info('Port forwarding stopped')
        # TODO(sergiitk): make debug
        logger.info('Port forwarding remaining stdout: %s', stdout)

    @staticmethod
    def _pod_started(pod: V1Pod):
        return pod.status.phase not in ('Pending', 'Unknown')

    @staticmethod
    def _replicas_available(deployment, count):
        return (deployment is not None and
                deployment.status.available_replicas is not None and
                deployment.status.available_replicas >= count)