<reponame>wacuuu/workload-collocation-agent
# Copyright (c) 2020 Intel Corporation
#
# 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 base64
import copy
from enum import Enum
import logging
from typing import Dict, List, Optional

import jsonpatch
from flask import Flask, jsonify, request

from wca.admission_controller.app_data_provider import AppDataProvider

log = logging.getLogger(__name__)


class AnnotatingService:
    """
    """
    def __init__(self, configuration: Dict[str, str]):
        """
        self.dram_only_threshold - set to 100.0, to not use this functionality, otherwise it can
            be decided that despite having wss/rss < 100% all memory will be allocated on DRAM.
        self.cold_start_duration - whether to start the workload on PMEM memory, if memory type
            for workload is DRAM,PMEM, duration of that period if set
        self.if_toptier_limit - whether to add toptier_limit annotation
        """
        self.data_provider: AppDataProvider = configuration['data_provider']
        self.app = Flask(__name__)

        self.dram_only_threshold: float = configuration.get('dram_threshold', 100.0)
        self.cold_start_duration: Optional[int] = None
        self.if_toptier_limit: bool = True

        self.monitored_namespaces: List[str] = \
            configuration.get('monitored_namespaces', ['default'])

        @self.app.route("/mutate", methods=["POST"])
        def mutate():
            return self._mutate_pod()

    def _create_patch(self, spec, modified_spec, message='Patching pod'):
        patch = jsonpatch.JsonPatch.from_diff(spec, modified_spec)
        return jsonify(
            {
                "response": {
                    "allowed": True,
                    "status": {"message": message},
                    "patch": base64.b64encode(str(patch).encode()).decode(),
                    "patchtype": "JSONPatch",
                }
            }
        )

    def _get_wss(self, app_name) -> Optional[str]:
        wss_per_app = self.data_provider.get_wss()
        log.debug('wss={}, app={}'.format(wss_per_app, app_name))
        if app_name in wss_per_app:
            return str(wss_per_app[app_name])
        return None

    def _get_rss(self, app_name) -> Optional[str]:
        rss_per_app = self.data_provider.get_rss()
        log.debug('rss={}, app={}'.format(rss_per_app, app_name))
        if app_name in rss_per_app:
            return str(rss_per_app[app_name])
        return None

    def _get_memory_type(self, wss: str, rss: str):
        ratio = float(wss) / float(rss) * 100
        if ratio > self.dram_only_threshold:
            memory_type = MemoryType.DRAM
        else:
            memory_type = MemoryType.HMEM
        return memory_type

    def _mutate_pod(self):
        annotations_key = "annotations"
        spec = request.json["request"]["object"]
        modified_spec = copy.deepcopy(spec)
        annotations = {}

        log.debug("[_mutate_pod] modified_spec={}".format(modified_spec))
        log.debug("[_mutate_pod] request={}".format(request.json["request"]))

        if request.json["request"]["namespace"] not in self.monitored_namespaces:
            return self._create_patch(spec, modified_spec)

        app_name: str = modified_spec["metadata"]["labels"]["app"]
        wss: Optional[str] = self._get_wss(app_name)
        rss: Optional[str] = self._get_rss(app_name)

        if wss is not None:
            if self.if_toptier_limit:
                annotations.update({'toptierlimit.cri-resource-manager.intel.com/pod':
                                    '{}G'.format(wss)})

        if wss is not None and rss is not None:
            memory_type = self._get_memory_type(wss, rss)
            annotations.update({'cri-resource-manager.intel.com/memory-type': memory_type})

            if memory_type == MemoryType.HMEM and self.cold_start_duration is not None:
                annotations.update({'cri-resource-manager.intel.com/cold-start':
                                   {'duration': self.cold_start_duration}})

        if not modified_spec["metadata"].get(annotations_key) and annotations:
            modified_spec["metadata"].update({annotations_key: {}})
        if annotations:
            modified_spec["metadata"][annotations_key].update(annotations)

        return self._create_patch(spec, modified_spec)


class MemoryType(str, Enum):
    DRAM = 'dram'
    HMEM = 'dram,pmem'
