path: root/tools/run_tests/xds_k8s_test_driver/framework/infrastructure/gcp/compute.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 dataclasses
import enum
import logging
from typing import Any, Dict, List, Optional

from googleapiclient import discovery
import googleapiclient.errors
# TODO(sergiitk): replace with tenacity
import retrying

from framework.infrastructure import gcp

logger = logging.getLogger(__name__)


class ComputeV1(gcp.api.GcpProjectApiResource):
    # TODO(sergiitk): move someplace better
    _WAIT_FOR_BACKEND_SEC = 60 * 10
    _WAIT_FOR_OPERATION_SEC = 60 * 5

    @dataclasses.dataclass(frozen=True)
    class GcpResource:
        name: str
        url: str

    @dataclasses.dataclass(frozen=True)
    class ZonalGcpResource(GcpResource):
        zone: str

    def __init__(self, api_manager: gcp.api.GcpApiManager, project: str):
        super().__init__(api_manager.compute('v1'), project)

    class HealthCheckProtocol(enum.Enum):
        TCP = enum.auto()
        GRPC = enum.auto()

    class BackendServiceProtocol(enum.Enum):
        HTTP2 = enum.auto()
        GRPC = enum.auto()

    def create_health_check(self,
                            name: str,
                            protocol: HealthCheckProtocol,
                            *,
                            port: Optional[int] = None) -> GcpResource:
        if protocol is self.HealthCheckProtocol.TCP:
            health_check_field = 'tcpHealthCheck'
        elif protocol is self.HealthCheckProtocol.GRPC:
            health_check_field = 'grpcHealthCheck'
        else:
            raise TypeError(f'Unexpected Health Check protocol: {protocol}')

        health_check_settings = {}
        if port is None:
            health_check_settings['portSpecification'] = 'USE_SERVING_PORT'
        else:
            health_check_settings['portSpecification'] = 'USE_FIXED_PORT'
            health_check_settings['port'] = port

        return self._insert_resource(
            self.api.healthChecks(), {
                'name': name,
                'type': protocol.name,
                health_check_field: health_check_settings,
            })

    def delete_health_check(self, name):
        self._delete_resource(self.api.healthChecks(), 'healthCheck', name)

    def create_firewall_rule(self, name: str, network_url: str,
                             source_ranges: List[str],
                             ports: List[str]) -> Optional[GcpResource]:
        try:
            return self._insert_resource(
                self.api.firewalls(), {
                    "allowed": [{
                        "IPProtocol": "tcp",
                        "ports": ports
                    }],
                    "direction": "INGRESS",
                    "name": name,
                    "network": network_url,
                    "priority": 1000,
                    "sourceRanges": source_ranges,
                    "targetTags": ["allow-health-checks"]
                })
        except googleapiclient.errors.HttpError as http_error:
            # TODO(lidiz) use status_code() when we upgrade googleapiclient
            if http_error.resp.status == 409:
                logger.debug('Firewall rule %s already existed', name)
                return
            else:
                raise

    def delete_firewall_rule(self, name):
        self._delete_resource(self.api.firewalls(), 'firewall', name)

    def create_backend_service_traffic_director(
            self,
            name: str,
            health_check: GcpResource,
            affinity_header: str = None,
            protocol: Optional[BackendServiceProtocol] = None) -> GcpResource:
        if not isinstance(protocol, self.BackendServiceProtocol):
            raise TypeError(f'Unexpected Backend Service protocol: {protocol}')
        body = {
            'name': name,
            'loadBalancingScheme': 'INTERNAL_SELF_MANAGED',  # Traffic Director
            'healthChecks': [health_check.url],
            'protocol': protocol.name,
        }
        # If affinity header is specified, config the backend service to support
        # affinity, and set affinity header to the one given.
        if affinity_header:
            body['sessionAffinity'] = 'HEADER_FIELD'
            body['localityLbPolicy'] = 'RING_HASH'
            body['consistentHash'] = {
                'httpHeaderName': affinity_header,
            }
        return self._insert_resource(self.api.backendServices(), body)

    def get_backend_service_traffic_director(self, name: str) -> GcpResource:
        return self._get_resource(self.api.backendServices(),
                                  backendService=name)

    def patch_backend_service(self, backend_service, body, **kwargs):
        self._patch_resource(collection=self.api.backendServices(),
                             backendService=backend_service.name,
                             body=body,
                             **kwargs)

    def backend_service_patch_backends(
            self,
            backend_service,
            backends,
            max_rate_per_endpoint: Optional[int] = None):
        if max_rate_per_endpoint is None:
            max_rate_per_endpoint = 5
        backend_list = [{
            'group': backend.url,
            'balancingMode': 'RATE',
            'maxRatePerEndpoint': max_rate_per_endpoint
        } for backend in backends]

        self._patch_resource(collection=self.api.backendServices(),
                             body={'backends': backend_list},
                             backendService=backend_service.name)

    def backend_service_remove_all_backends(self, backend_service):
        self._patch_resource(collection=self.api.backendServices(),
                             body={'backends': []},
                             backendService=backend_service.name)

    def delete_backend_service(self, name):
        self._delete_resource(self.api.backendServices(), 'backendService',
                              name)

    def create_url_map(
        self,
        name: str,
        matcher_name: str,
        src_hosts,
        dst_default_backend_service: GcpResource,
        dst_host_rule_match_backend_service: Optional[GcpResource] = None,
    ) -> GcpResource:
        if dst_host_rule_match_backend_service is None:
            dst_host_rule_match_backend_service = dst_default_backend_service
        return self._insert_resource(
            self.api.urlMaps(), {
                'name':
                    name,
                'defaultService':
                    dst_default_backend_service.url,
                'hostRules': [{
                    'hosts': src_hosts,
                    'pathMatcher': matcher_name,
                }],
                'pathMatchers': [{
                    'name': matcher_name,
                    'defaultService': dst_host_rule_match_backend_service.url,
                }],
            })

    def create_url_map_with_content(self, url_map_body: Any) -> GcpResource:
        return self._insert_resource(self.api.urlMaps(), url_map_body)

    def patch_url_map(self, url_map: GcpResource, body, **kwargs):
        self._patch_resource(collection=self.api.urlMaps(),
                             urlMap=url_map.name,
                             body=body,
                             **kwargs)

    def delete_url_map(self, name):
        self._delete_resource(self.api.urlMaps(), 'urlMap', name)

    def create_target_grpc_proxy(
        self,
        name: str,
        url_map: GcpResource,
    ) -> GcpResource:
        return self._insert_resource(self.api.targetGrpcProxies(), {
            'name': name,
            'url_map': url_map.url,
            'validate_for_proxyless': True,
        })

    def delete_target_grpc_proxy(self, name):
        self._delete_resource(self.api.targetGrpcProxies(), 'targetGrpcProxy',
                              name)

    def create_target_http_proxy(
        self,
        name: str,
        url_map: GcpResource,
    ) -> GcpResource:
        return self._insert_resource(self.api.targetHttpProxies(), {
            'name': name,
            'url_map': url_map.url,
        })

    def delete_target_http_proxy(self, name):
        self._delete_resource(self.api.targetHttpProxies(), 'targetHttpProxy',
                              name)

    def create_forwarding_rule(
        self,
        name: str,
        src_port: int,
        target_proxy: GcpResource,
        network_url: str,
    ) -> GcpResource:
        return self._insert_resource(
            self.api.globalForwardingRules(),
            {
                'name': name,
                'loadBalancingScheme':
                    'INTERNAL_SELF_MANAGED',  # Traffic Director
                'portRange': src_port,
                'IPAddress': '0.0.0.0',
                'network': network_url,
                'target': target_proxy.url,
            })

    def exists_forwarding_rule(self, src_port) -> bool:
        # TODO(sergiitk): Better approach for confirming the port is available.
        #   It's possible a rule allocates actual port range, e.g 8000-9000,
        #   and this wouldn't catch it. For now, we assume there's no
        #   port ranges used in the project.
        filter_str = (f'(portRange eq "{src_port}-{src_port}") '
                      f'(IPAddress eq "0.0.0.0")'
                      f'(loadBalancingScheme eq "INTERNAL_SELF_MANAGED")')
        return self._exists_resource(self.api.globalForwardingRules(),
                                     filter=filter_str)

    def delete_forwarding_rule(self, name):
        self._delete_resource(self.api.globalForwardingRules(),
                              'forwardingRule', name)

    @staticmethod
    def _network_endpoint_group_not_ready(neg):
        return not neg or neg.get('size', 0) == 0

    def wait_for_network_endpoint_group(self, name, zone):

        @retrying.retry(retry_on_result=self._network_endpoint_group_not_ready,
                        stop_max_delay=60 * 1000,
                        wait_fixed=2 * 1000)
        def _wait_for_network_endpoint_group_ready():
            try:
                neg = self.get_network_endpoint_group(name, zone)
                logger.debug(
                    'Waiting for endpoints: NEG %s in zone %s, '
                    'current count %s', neg['name'], zone, neg.get('size'))
            except googleapiclient.errors.HttpError as error:
                # noinspection PyProtectedMember
                reason = error._get_reason()
                logger.debug('Retrying NEG load, got %s, details %s',
                             error.resp.status, reason)
                raise
            return neg

        network_endpoint_group = _wait_for_network_endpoint_group_ready()
        # TODO(sergiitk): dataclass
        return self.ZonalGcpResource(network_endpoint_group['name'],
                                     network_endpoint_group['selfLink'], zone)

    def get_network_endpoint_group(self, name, zone):
        neg = self.api.networkEndpointGroups().get(project=self.project,
                                                   networkEndpointGroup=name,
                                                   zone=zone).execute()
        # TODO(sergiitk): dataclass
        return neg

    def wait_for_backends_healthy_status(
        self,
        backend_service,
        backends,
        timeout_sec=_WAIT_FOR_BACKEND_SEC,
        wait_sec=4,
    ):
        pending = set(backends)

        @retrying.retry(retry_on_result=lambda result: not result,
                        stop_max_delay=timeout_sec * 1000,
                        wait_fixed=wait_sec * 1000)
        def _retry_backends_health():
            for backend in pending:
                result = self.get_backend_service_backend_health(
                    backend_service, backend)

                if 'healthStatus' not in result:
                    logger.debug('Waiting for instances: backend %s, zone %s',
                                 backend.name, backend.zone)
                    continue

                backend_healthy = True
                for instance in result['healthStatus']:
                    logger.debug(
                        'Backend %s in zone %s: instance %s:%s health: %s',
                        backend.name, backend.zone, instance['ipAddress'],
                        instance['port'], instance['healthState'])
                    if instance['healthState'] != 'HEALTHY':
                        backend_healthy = False

                if backend_healthy:
                    logger.info('Backend %s in zone %s reported healthy',
                                backend.name, backend.zone)
                    pending.remove(backend)

            return not pending

        _retry_backends_health()

    def get_backend_service_backend_health(self, backend_service, backend):
        return self.api.backendServices().getHealth(
            project=self.project,
            backendService=backend_service.name,
            body={
                "group": backend.url
            }).execute()

    def _get_resource(self, collection: discovery.Resource,
                      **kwargs) -> GcpResource:
        resp = collection.get(project=self.project, **kwargs).execute()
        logger.info('Loaded compute resource:\n%s',
                    self.resource_pretty_format(resp))
        return self.GcpResource(resp['name'], resp['selfLink'])

    def _exists_resource(self, collection: discovery.Resource,
                         filter: str) -> bool:
        resp = collection.list(
            project=self.project, filter=filter,
            maxResults=1).execute(num_retries=self._GCP_API_RETRIES)
        if 'kind' not in resp:
            # TODO(sergiitk): better error
            raise ValueError('List response "kind" is missing')
        return 'items' in resp and resp['items']

    def _insert_resource(self, collection: discovery.Resource,
                         body: Dict[str, Any]) -> GcpResource:
        logger.info('Creating compute resource:\n%s',
                    self.resource_pretty_format(body))
        resp = self._execute(collection.insert(project=self.project, body=body))
        return self.GcpResource(body['name'], resp['targetLink'])

    def _patch_resource(self, collection, body, **kwargs):
        logger.info('Patching compute resource:\n%s',
                    self.resource_pretty_format(body))
        self._execute(
            collection.patch(project=self.project, body=body, **kwargs))

    def _delete_resource(self, collection: discovery.Resource,
                         resource_type: str, resource_name: str) -> bool:
        try:
            params = {"project": self.project, resource_type: resource_name}
            self._execute(collection.delete(**params))
            return True
        except googleapiclient.errors.HttpError as error:
            if error.resp and error.resp.status == 404:
                logger.info(
                    'Resource %s "%s" not deleted since it does not exist',
                    resource_type, resource_name)
            else:
                logger.warning('Failed to delete %s "%s", %r', resource_type,
                               resource_name, error)
        return False

    @staticmethod
    def _operation_status_done(operation):
        return 'status' in operation and operation['status'] == 'DONE'

    def _execute(self,
                 request,
                 *,
                 test_success_fn=None,
                 timeout_sec=_WAIT_FOR_OPERATION_SEC):
        operation = request.execute(num_retries=self._GCP_API_RETRIES)
        logger.debug('Response %s', operation)

        # TODO(sergiitk) try using wait() here
        # https://googleapis.github.io/google-api-python-client/docs/dyn/compute_v1.globalOperations.html#wait
        operation_request = self.api.globalOperations().get(
            project=self.project, operation=operation['name'])

        if test_success_fn is None:
            test_success_fn = self._operation_status_done

        logger.debug('Waiting for global operation %s, timeout %s sec',
                     operation['name'], timeout_sec)
        response = self.wait_for_operation(operation_request=operation_request,
                                           test_success_fn=test_success_fn,
                                           timeout_sec=timeout_sec)

        if 'error' in response:
            logger.debug('Waiting for global operation failed, response: %r',
                         response)
            raise Exception(f'Operation {operation["name"]} did not complete '
                            f'within {timeout_sec}s, error={response["error"]}')
        return response