class MemoryCgroup(Cgroup):

    CONTROLLER = 'memory'

    MAX_MEM_VAL = INT64_MAX

    def _write_file_mem_val(self, filename, value):
        if isinstance(value, int):
            value = str(min(value, self.MAX_MEM_VAL))
        self._write_file_str(filename, value)

    def read_mem_current(self):
        return self._read_file_int('usage_in_bytes')

    def read_swap_current(self):
        mem = self._read_file_int('usage_in_bytes')
        memsw = self._read_file_int('memsw.usage_in_bytes')
        return max(memsw - mem, 0)

    def read_mem_low(self):
        return self._read_file_int('low')

    def write_mem_low(self, val):
        self._write_file_mem_val('low', val)

    def read_mem_high(self):
        return self._read_file_int('high')

    def write_mem_high(self, val):
        self._write_file_mem_val('high', val)

    def write_mem_config(self, mem, sw):
        mem_old = self._read_file_int('limit_in_bytes')
        memsw = mem + sw
        if mem > mem_old:
            self.write_memsw_limit_in_bytes( memsw)
            self.write_limit_in_bytes(mem)
        else:
            self.write_limit_in_bytes(mem)
            self.write_memsw_limit_in_bytes( memsw)

    def read_mem_max(self):
        return self._read_file_int('limit_in_bytes')

    def write_limit_in_bytes(self, val):
        # Warning: also changes swap size to memsw.limit_in_bytes - val
        self._write_file_mem_val('limit_in_bytes', val)

    def read_swap_max(self):
        mem = self._read_file_int('limit_in_bytes')
        memsw = self._read_file_int('memsw.limit_in_bytes')
        return max(memsw - mem, 0)

    def write_memsw_limit_in_bytes(self, val):
        # Warning: changes swap size to val - limit_in_bytes
        self._write_file_mem_val('memsw.limit_in_bytes', val)

    def write_cache_limit_in_bytes(self, val):
        self._write_file_mem_val('cache.limit_in_bytes', val)

    def read_mem_stat(self):
        return self._read_file_kv('stat')

    def write_oom_guarantee(self, val):
        self._write_file_int('oom_guarantee', val)

    def write_swappiness(self, val):
        self._write_file_int('swappiness', val)

    def read_swappiness(self):
        return self._read_file_int('swappiness')

    def write_oom_control(self, val):
        self._write_file_int('oom_control', val)

    def write_cleancache(self, val):
        self._write_file_int('disable_cleancache', int(not val))

    def read_cache_limit_in_bytes(self):
        return self._read_file_int('cache.limit_in_bytes')

    def get_numa_stats(self):
        stats = self._read_file_str("numa_stat")
        res = {}
        for line in stats.split("\n"):
            if not line:
                continue
            name, data = re.split(" ", line, maxsplit = 1)
            name = "mem" + name.split("=")[0]
            for node, value in re.findall("N(\d+)=(\d+)", data):
                node, value = int(node), int(value)
                if node not in res:
                    res[node] = {}
                res[node][name] = value * PAGE_SIZE
        return res

    def set_node_list(self, nodes):
        self._write_file_str('numa_migrate', ','.join(map(str, nodes))) class BlkIOCgroup(Cgroup):

    CONTROLLER = 'blkio'

    def _get_io_stats(self, filename, keys):
        result = {k: 0 for k in keys}
        kv = self._read_file_kv(filename)
        for k, v in kv.items():
            try:
                result[k.split()[-1]] += v
            except KeyError:
                continue
        return tuple(result[k] for k in keys)

    def get_io_serviced(self):
        '''Return a tuple containing the total number of read and write
        requests issued by this cgroup.
        '''
        return self._get_io_stats('io_serviced', ('Read', 'Write'))

    def get_io_service_bytes(self):
        '''Return a tuple containing the total number of bytes read and written
        by this cgroup.
        '''
        return self._get_io_stats('io_service_bytes', ('Read', 'Write')) class CpuSetCgroup(Cgroup):

    CONTROLLER = 'cpuset'

    def get_cpu_list(self):
        return parse_range_list(self._read_file_str("cpus"))

    def get_node_list(self):
        return parse_range_list(self._read_file_str("mems"))

    def set_memory_migrate(self, val):
        self._write_file_str("memory_migrate", str(int(val)))

    def set_cpu_list(self, cpus):
        self._write_file_str("cpus",",".join(map(str,cpus)))

    def set_node_list(self, nodes):
        self._write_file_str("mems", ",".join(map(str, nodes))) class CpuCgroup(Cgroup):

    CONTROLLER = 'cpu'

    def get_cpu_stats(self):
        names = ["cpuuser", "cpunice", "cpusystem", "cpuidle"]
        stats = self._read_file_str("proc.stat")
        res = {}
        for line in stats.splitlines():
            if not re.search("cpu\d+", line):
                continue
            cpu, data = re.split(" ", line, maxsplit = 1)
            cpu = int(cpu[3:])
            for name, value in zip(names, re.findall("(\d+)", data)):
                if cpu not in res:
                    res[cpu] = {}
                res[cpu][name] = int(value)
        return res

    def get_nr_cpus(self):
        return self._read_file_int("nr_cpus")

    def write_cpu_shares(self, val):
        self._write_file_int('shares', val) class Error(Exception):
    pass class VEImpl:
    '''VE implementation.

    This class defines the interface to an underlying VE implementation
    (such as libvirt or cgroup).

    Any of the functions defined by this interface may raise Error.
    '''

    def __init__(self, name):
        pass

    @staticmethod
    def mem_overhead(config_limit):
        '''Return an estimate of memory overhead.

        This function is supposed to return the amount of memory beyond the
        configured limit which is required to run the VE smoothly. E.g. for
        VMs this should equal expected RSS of the emulator process.
        '''
        return 0

    def get_stats(self):
        '''Return stats dict {name: value}.
        '''
        pass

    def set_mem_protection(self, value):
        '''Set memory best-effort protection.

        If memory usage of a VE is below this value, the VE's memory shouldn't
        be reclaimed on host pressure if memory can be reclaimed from
        unprotected VEs.
        '''
        pass

    def set_mem_target(self, value):
        '''Set memory allocation target.

        This function sets memory consumption target for a VE. Note, it does
        not necessarily mean that the VE memory usage will reach the target
        instantly or even any time soon - in fact, it may not reach it at all
        in case allocation is reduced. However, reducing the value will put the
        VE under heavy local memory pressure forcing it to release its memory
        to the host.
        '''
        pass

    def set_config(self, config):
        '''Set new config.
        '''
        pass

    def pin_node_mem(self, nodes):
        '''
        Should be expanded when memory migration "knob" will be implemented
        in kernel.
        '''
        pass

    def pin_cpu_list(self, cpus):
        '''
        Should be expanded when memory migration "knob" will be implemented
        in kernel.
        '''
        pass def register_ve_impl(ve_impl):
    assert ve_impl.VE_TYPE not in _VE_IMPL_MAP
    _VE_IMPL_MAP[ve_impl.VE_TYPE] = ve_impl VE_TYPE_CT = 0 VE_TYPE_SERVICE = 4 class VCMMDConfig(metaclass=Singleton):
    '''VCMMD config loader.

    This is a singleton class that provides methods for loading VCMMD
    configuration from a file and getting option values by name.
    '''

    def __init__(self, filename):
        self.logger = logging.getLogger('vcmmd.config')
        self._filename = filename
        self._data = None
        self._cache = {}

    def load(self):
        '''Load config from a file.

        The file must be in json format.
        '''
        self._data = None
        self._cache = {}

        self.logger.info("Loading config from file '%s'", self._filename)
        self._data = self.read()

    def _get_legacy_storage_limits(self):
        if os.path.isfile('/etc/vz/vstorage-limits.conf'):
            try:
                with open('/etc/vz/vstorage-limits.conf') as f:
                    return json.load(f)
            except (IOError, ValueError) as err:
                self.logger.error('Error reading vstorage-limits.conf: %s', err)
        return {}

    def read(self):
        '''Read config from a file.

        The file must be in json format.
        '''
        try:
            with open(self._filename, 'r') as f:
                data = json.load(f)
        except IOError as err:
            self.logger.error('Error reading config file: %s', err)
            return None
        except ValueError as err:
            self.logger.error('Error parsing config file: %s', err)
            return None
        limits_config = self._get_default_limits_config()
        if 'Limits' in data:
            limits_config = self._update(limits_config, data['Limits'])
        limits_config = self._update(limits_config, self._get_legacy_storage_limits())
        data['Limits'] = limits_config
        return data

    def dump(self, name, val):
        self._data = self.read()

        try:
            _data = self._data
            key = name.split('.')[-1]
            path = name.split('.')[:-1]
            for k in path:
                _data = _data[k]
            _data[key] = val
        except KeyError as err:
            self.logger.error('Error parsing config file: %s', err)
            return

        try:
            with open(self._filename, 'w') as f:
                f.write(json.dumps(self._data, indent=8))
        except IOError as err:
            self.logger.error('Error writing config file: %s', err)

        self._cache[name] = val

    @classmethod
    def _update(cls, a_dict, b_dict):
        """
        Update a_dict with b_dict values hierarchically.
        """
        for k, v in b_dict.items():
            if isinstance(v, dict):
                a_dict[k] = cls._update(a_dict.get(k, {}), v)
            else:
                a_dict[k] = v
        return a_dict

    @classmethod
    def _exclude(cls, a_dict, b_dict):
        """
        Remove from a_dict items storing under the same hierarchy in b_dict and having the same value.
        """
        for k, v in list(a_dict.items()):
            if k in b_dict:
                if isinstance(b_dict[k], dict) and isinstance(a_dict[k], dict):
                    cls._exclude(a_dict[k], b_dict[k])
                    if not a_dict[k]:
                        del a_dict[k]
                else:
                    if v == b_dict[k]:
                        del a_dict[k]

    @staticmethod
    def _get_default_limits_config():
        def _expand(min, max, share):
            return {"Min": min, "Max": max, "Share": share}

        return {
            "System": {
                "Path": "system.slice",
                "Limit": _expand(0, -1, 1),
                "Guarantee": _expand(320 << 20, 700 << 20, 0.04),
                "Swap": _expand(0, 0, 0),
            },
            "User": {
                "Path": "user.slice",
                "Limit": _expand(0, -1, 1),
                "Guarantee": _expand(32 << 20, 128 << 20, 0.02),
                "Swap": _expand(0, 0, 0),
            },
            "VStorage": {
                "Path": "vstorage.slice/vstorage-services.slice",
                "Limit": _expand(0, -1, 0.7),
                "Guarantee": _expand(0, 0, 0.25),
                "Swap": _expand(0, 0, 0),
            },
        }

    def _get(self, name):
        d = self._data
        for k in name.split('.'):
            if not isinstance(d, dict) or k not in d:
                raise KeyError
            d = d[k]
        return d

    def get(self, name, default=None, checkfn=None):
        '''Get the value of a config option.

        This function lookups a config option by name and returns its value.
        To lookup an option in a sub-section, use dot, e.g. 'section.option'.
        In case the option does not exist, the value of 'default' is returned.

        Unless 'checkfn' is not None, no checks is performed upon the retrieved
        value. To assure the value meets specific requirements, use get_str,
        get_bool, and get_num methods.

        If 'checkfn' argument is not None, it must be a function taking exactly
        one argument. The function will be called to check the retrieved value.
        It may raise ValueError or TypeError, in which case the retrieved value
        will be discarded and 'default' will be returned.

        Note, the value returned by this function is cached, meaning that the
        next call to it with the same 'name' will return the same value
        bypassing any checks.
        '''
        # First, check if we've already fetched the requested value. If this is
        # the case, bypass any checks and return the cached value.
        # TODO: replace caching with @lru_cache after upgrading to Py3
        try:
            return self._cache[name]
        except KeyError:
            pass
        try:
            val = self._get(name)
            if checkfn is not None:
                checkfn(val)
        except (KeyError, TypeError, ValueError) as err:
            # do not complain if the option is absent
            if not isinstance(err, KeyError):
                self.logger.warning("Invalid value for config option '%s': %s",
                                    name, err)
            val = default
        # Save the value to speed up following retrievals and avoid spewing
        # warnings if any over and over again.
        self._cache[name] = val
        self.logger.debug('%s = %r', name, val)
        return val

    def get_str(self, name, default=None):
        def checkfn(val):
            t = type(val)
            if t != str:
                raise TypeError("expected string, got '{}'".format(t.__name__))
        return self.get(name, default, checkfn)

    def get_bool(self, name, default=None):
        def checkfn(val):
            t = type(val)
            if t != bool:
                raise TypeError("expected boolean, got '{}'".format(t.__name__))
        return self.get(name, default, checkfn)

    def get_num(self, name, default=None,
                integer=False, minimum=None, maximum=None):
        def checkfn(val):
            t = type(val)
            if t not in (int, float) or (integer and t == float):
                raise TypeError("expected {}, got "
                                "'{}'".format('integer' if integer else 'number',
                                              t.__name__))
            if minimum is not None and val < minimum:
                raise ValueError("must be >= {}, got {}".format(minimum, val))
            if maximum is not None and val > maximum:
                raise ValueError("must be <= {}, got {}".format(maximum, val))
        return self.get(name, default, checkfn)

    def get_choice(self, name, choices, default=None):
        def checkfn(val):
            t = type(val)
            if t != str:
                raise TypeError("expected string, got '{}'".format(t.__name__))
            if val not in choices:
                raise ValueError("must be one of {}, got "
                                 "{}".format(tuple(choices), str(val)))
        return self.get(name, default, checkfn)

    def _make_config(self):
        cfg_dict = {}
        for name in self._cache:
            x = cfg_dict
            path = name.split('.')
            key = path[-1]
            path = path[:-1]
            for section in path:
                if section not in x:
                    x[section] = {}
                x = x[section]
            x[key] = self._cache[name]
        return cfg_dict

    def report(self, full_config=False):
        config = copy.deepcopy(self._make_config())
        if not full_config and 'Limits' in config:
            default_limits = self._get_default_limits_config()
            self._exclude(config['Limits'], default_limits)
            if not config['Limits']:
                del config['Limits']
        return json.dumps(config) INT64_MAX = int(2 ** 63 - 1) def run_async(func, *args, **kwargs):
    return _thread_pool.apply_async(func, args, kwargs)
import logging
import time
from vcmmd.cgroup import MemoryCgroup, BlkIOCgroup, CpuSetCgroup, CpuCgroup
from vcmmd.ve.base import Error, VEImpl, register_ve_impl
from vcmmd.ve_type import VE_TYPE_CT, VE_TYPE_SERVICE
from vcmmd.config import VCMMDConfig
from vcmmd.util.limits import INT64_MAX
from vcmmd.util.threading import run_async
            return self._memcg.read_mem_stat()['total_rss']
        except (KeyError, IOError) as err:
            raise Error('Cgroup read failed: {}'.format(err))

    def set_mem_protection(self, value):
        """Use memcg/memory.low to protect the CT from host pressure."""
        try:
            self._memcg.write_mem_low(value)
        except IOError as err:
            raise Error('Cgroup write failed: {}'.format(err))

    def set_mem_target(self, value):
        # Decreasing memory.high might take long as it implies memory reclaim,
        # so do it asynchronously.
        #
        # XXX: For the sake of simplicity, we don't care about failures here.
        # It is acceptable, because adjusting memory.high may fail only if
        # the cgroup gets destroyed, which we will see and report anyway from
        # get_stats().
        run_async(self._memcg.write_mem_high, value)
        self.mem_limit = value

    def set_config(self, config):
        try:
            self._memcg.write_oom_guarantee(config.guarantee)
            self._memcg.write_mem_config(config.limit, config.swap)
        except IOError as err:
            raise Error('Cgroup write failed: {}'.format(err))

        self.mem_limit = min(self.mem_limit, config.limit)
        run_async(self._memcg.write_cache_limit_in_bytes, config.cache)


class CTImpl(ABSVEImpl):

    def __init__(self, name):
        super(CTImpl, self).__init__(name)
        self._cpusetcg = lookup_cgroup(CpuSetCgroup, name)
        self._blkcg = lookup_cgroup(BlkIOCgroup, name)
        self._cpucg = lookup_cgroup(CpuCgroup, name)
        self._memcg = lookup_cgroup(MemoryCgroup, name)

    def get_stats(self):
        try:
            current = self._memcg.read_mem_current()
            stat = self._memcg.read_mem_stat()
            io_serviced = self._blkcg.get_io_serviced()
            io_service_bytes = self._blkcg.get_io_service_bytes()
        except IOError as err:
            raise Error('Cgroup read failed: {}'.format(err))

        memtotal = max(self.mem_limit, current)
        memfree = memtotal - current
        memavail = (memfree +
                    stat.get('active_file', 0) +
                    stat.get('inactive_file', 0) +
                    stat.get('slab_reclaimable', 0))

        return {'rss': current,
                'host_mem': current,
                'host_swap': stat.get('swap', -1),
                'actual': memtotal,
                'memfree': memfree,
                'memavail': memavail,
                'swapin': stat.get('pswpin', -1),
                'swapout': stat.get('pswpout', -1),
                'minflt': stat.get('pgfault', -1),
                'majflt': stat.get('pgmajfault', -1),
                'rd_req': io_serviced[0],
                'rd_bytes': io_service_bytes[0],
                'wr_req': io_serviced[1],
                'wr_bytes': io_service_bytes[1],
                'last_update': int(time.time())}

    @property
    def nr_cpus(self):
        try:
            self._nr_cpus = self._cpucg.get_nr_cpus()
            return self._nr_cpus
        except IOError:
            return getattr(self, "_nr_cpus", -1)

    def get_node_list(self):
        """Get list of nodes where CT is running."""
        try:
            node_list = self._cpusetcg.get_node_list()
        except IOError as err:
            raise Error('Cgroup read failed: {}'.format(err))
        return node_list

    def node_mem_migrate(self, nodes):
        try:
            self._memcg.set_node_list(nodes)
        except IOError as err:
            raise Error('Cgroup write failed: {}'.format(err))

    def pin_node_mem(self, nodes):
        """Change list of memory nodes for CT.

        This function changes CT affinity for memory and migrates CT's memory
        accordingly.
        """
        try:
            self._cpusetcg.set_node_list(nodes)
        except IOError as err:
            raise Error('Cgroup write failed: {}'.format(err))

    def pin_cpu_list(self, cpus):
        """Change list of CPUs for CT.

        This function changes CT affinity for CPUs.
        """
        try:
            self._cpusetcg.set_cpu_list(cpus)
        except IOError as err:
            raise Error('Cgroup write failed: {}'.format(err))


class ServiceCTImpl(ABSVEImpl):


    VE_TYPE = VE_TYPE_SERVICE