1#!/usr/bin/python2
2# Copyright 2018 The Chromium OS Authors. All rights reserved.
3# Use of this source code is governed by a BSD-style license that can be
4# found in the LICENSE file.
5import collections
6import logging
7import re
8import threading
9
10from autotest_lib.client.bin import utils
11
12class SystemSampler():
13    """A sampler class used to probe various system stat along with FPS.
14
15    Sample usage:
16        def get_consumer_pid():
17            # Returns all memory consumers whose memory usage we care about.
18
19        sampler = SystemSampler(get_consumer_pid)
20        fps = ...
21        sampler.sample(fps)
22    """
23    Sample = collections.namedtuple(
24        'Sample', ['pswpin', 'pswpout', 'free_mem', 'buff_mem', 'cached_mem',
25                   'anon_mem', 'file_mem', 'swap_free', 'swap_used',
26                   'consumer_num', 'consumer_rss', 'consumer_swap',
27                   'cpuload', 'fps'])
28
29    VMStat = collections.namedtuple('VMStat', ['pswpin', 'pswpout'])
30
31    def __init__(self, get_consumer_pid_func):
32        self._get_consumer_pid_func = get_consumer_pid_func
33        self._samples_lock = threading.Lock()
34        self._samples = []
35        self.reset()
36
37    def reset(self):
38        """Resets its internal stats."""
39        self._prev_vmstat = self.get_vmstat()
40        self._prev_cpustat = utils.get_cpu_usage()
41
42    def get_samples(self):
43        with self._samples_lock:
44            return self._samples
45
46    def get_last_avg_fps(self, num):
47        """Gets average fps rate of last |num| samples.
48
49        Returns None if there's not enough samples in hand.
50        """
51        if num <= 0:
52            logging.warning('Num of samples must be > 0')
53            return
54
55        with self._samples_lock:
56            if len(self._samples) >= num:
57                return sum([s.fps for s in self._samples[-num:]])/float(num)
58
59        logging.warning('Not enough samples')
60        return None
61
62    def sample(self, fps_info):
63        """Gets a fps sample with system state."""
64        vmstat = self.get_vmstat()
65        vmstat_diff = self.VMStat(*[(end - start)
66            for start, end in zip(self._prev_vmstat, vmstat)])
67        self._prev_vmstat = vmstat
68
69        cpustat = utils.get_cpu_usage()
70        cpuload = utils.compute_active_cpu_time(
71            self._prev_cpustat, cpustat)
72        self._prev_cpustat = cpustat
73
74        mem_info_in_kb = utils.get_meminfo()
75        # Converts mem_info from KB to MB.
76        mem_info = collections.namedtuple('MemInfo', mem_info_in_kb._fields)(
77            *[v/float(1024) for v in mem_info_in_kb])
78
79        consumer_pids = self._get_consumer_pid_func()
80        logging.info('Consumers %s', consumer_pids)
81        consumer_rss, consumer_swap = self.get_consumer_meminfo(consumer_pids)
82
83        # fps_info = (frame_info, frame_times)
84        fps_count = len([f for f in fps_info[0] if f != ' '])
85
86        sample = self.Sample(
87            pswpin=vmstat_diff.pswpin,
88            pswpout=vmstat_diff.pswpout,
89            free_mem=mem_info.MemFree,
90            buff_mem=mem_info.Buffers,
91            cached_mem=mem_info.Cached,
92            anon_mem=mem_info.Active_anon + mem_info.Inactive_anon,
93            file_mem=mem_info.Active_file + mem_info.Inactive_file,
94            swap_free=mem_info.SwapFree,
95            swap_used=mem_info.SwapTotal - mem_info.SwapFree,
96            consumer_num=len(consumer_pids),
97            consumer_rss=consumer_rss,
98            consumer_swap=consumer_swap,
99            cpuload=cpuload,
100            fps=fps_count)
101
102        logging.info(sample)
103
104        with self._samples_lock:
105            self._samples.append(sample)
106
107    @staticmethod
108    def parse_meminfo_from_proc_entry(pid):
109        """Parses memory related info in /proc/<pid>/totmaps like:
110
111        Rss:              144956 kB
112        Pss:               74923 kB
113        Shared_Clean:      50596 kB
114        Shared_Dirty:      41660 kB
115        Private_Clean:      1032 kB
116        Private_Dirty:     51668 kB
117        Referenced:       137424 kB
118        Anonymous:         91772 kB
119        AnonHugePages:     30720 kB
120        Swap:                  0 kB
121        """
122        mem_info = {}
123        line_pattern = re.compile(r'^(\w+):\s+(\d+)\s+kB')
124        proc_entry = '/proc/%s/totmaps' % pid
125        try:
126            with open(proc_entry) as f:
127                for line in f:
128                    m = line_pattern.match(line)
129                    if m:
130                        key, value = m.groups()
131                        mem_info[key] = float(value)/1024
132        except IOError as e:
133            logging.warning('Failed to open %s: %s', proc_entry, e)
134        return mem_info
135
136    @classmethod
137    def get_consumer_meminfo(cls, pids):
138        rss = 0.0
139        swap = 0.0
140        for pid in pids:
141            mem_info = cls.parse_meminfo_from_proc_entry(pid)
142            rss += mem_info.get('Rss', 0)
143            swap += mem_info.get('Swap', 0)
144        return rss, swap
145
146    @classmethod
147    def get_vmstat(cls):
148        with open('/proc/vmstat') as f:
149            lines = f.readlines()
150        all_fields = dict([l.strip().split(' ') for l in lines])
151        return cls.VMStat(
152            *[int(all_fields.get(f, 0)) for f in cls.VMStat._fields])
153