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.
5
6import logging
7import subprocess
8import time
9import threading
10
11from autotest_lib.client.bin import utils
12
13class MemoryEater(object):
14    """A util class which run programs to consume memory in the background.
15
16    Sample usage:
17    with MemoryEator() as memory_eater:
18      # Allocate mlocked memory.
19      memory_eater.consume_locked_memory(123)
20
21      # Allocate memory and sequentially traverse them over and over.
22      memory_eater.consume_active_memory(500)
23
24    When it goes out of the "with" context or the object is destructed, all
25    allocated memory are released.
26    """
27
28    memory_eater_locked = 'memory-eater-locked'
29    memory_eater = 'memory-eater'
30
31    _all_instances = []
32
33    def __init__(self):
34        self._locked_consumers = []
35        self._active_consumers_lock = threading.Lock()
36        self._active_consumers = []
37        self._all_instances.append(self)
38
39    def __enter__(self):
40        return self
41
42    @staticmethod
43    def cleanup_consumers(consumers):
44        """Kill all processes in |consumers|
45
46        @param consumers: The list of consumers to clean.
47        """
48        while len(consumers):
49            job = consumers.pop()
50            logging.info('Killing %d', job.pid)
51            job.kill()
52
53    def cleanup(self):
54        """Releases all allocated memory."""
55        # Kill all hanging jobs.
56        logging.info('Cleaning hanging memory consuming processes...')
57        self.cleanup_consumers(self._locked_consumers)
58        with self._active_consumers_lock:
59            self.cleanup_consumers(self._active_consumers)
60
61    def __exit__(self, type, value, traceback):
62        self.cleanup()
63
64    def __del__(self):
65        self.cleanup()
66        if self in self._all_instances:
67            self._all_instances.remove(self)
68
69    def consume_locked_memory(self, mb):
70        """Consume non-swappable memory."""
71        logging.info('Consuming locked memory %d MB', mb)
72        cmd = [self.memory_eater_locked, str(mb)]
73        p = subprocess.Popen(cmd, stdout=subprocess.PIPE)
74        self._locked_consumers.append(p)
75        # Wait until memory allocation is done.
76        while True:
77            line = p.stdout.readline()
78            if line.find('Done') != -1:
79                break
80
81    def consume_active_memory(self, mb):
82        """Consume active memory."""
83        logging.info('Consuming active memory %d MB', mb)
84        cmd = [self.memory_eater, '--size', str(mb), '--chunk', '128']
85        p = subprocess.Popen(cmd, stdout=subprocess.PIPE)
86        with self._active_consumers_lock:
87            self._active_consumers.append(p)
88
89    @classmethod
90    def get_active_consumer_pids(cls):
91        """Gets pid of active consumers by all instances of the class."""
92        all_pids = []
93        for instance in cls._all_instances:
94            with instance._active_consumers_lock:
95                all_pids.extend([p.pid for p in instance._active_consumers])
96        return all_pids
97
98
99def consume_free_memory(memory_to_reserve_mb):
100    """Consumes free memory until |memory_to_reserve_mb| is remained.
101
102    Non-swappable memory is allocated to consume memory.
103    memory_to_reserve_mb: Consume memory until this amount of free memory
104        is remained.
105    @return The MemoryEater() object on which memory is allocated. One can
106        catch it in a context manager.
107    """
108    consumer = MemoryEater()
109    while True:
110        mem_free_mb = utils.read_from_meminfo('MemFree') / 1024
111        logging.info('Current Free Memory %d', mem_free_mb)
112        if mem_free_mb <= memory_to_reserve_mb:
113            break
114        memory_to_consume = min(
115            2047, mem_free_mb - memory_to_reserve_mb + 1)
116        logging.info('Consuming %d MB locked memory', memory_to_consume)
117        consumer.consume_locked_memory(memory_to_consume)
118    return consumer
119
120
121class TimeoutException(Exception):
122    """Exception to return if timeout happens."""
123    def __init__(self, message):
124        super(TimeoutException, self).__init__(message)
125
126
127class _Timer(object):
128    """A simple timer class to check timeout."""
129    def __init__(self, timeout, des):
130        """Initializer.
131
132        @param timeout: Timeout in seconds.
133        @param des: A short description for this timer.
134        """
135        self.timeout = timeout
136        self.des = des
137        if self.timeout:
138            self.start_time = time.time()
139
140    def check_timeout(self):
141        """Raise TimeoutException if timeout happens."""
142        if not self.timeout:
143          return
144        time_delta = time.time() - self.start_time
145        if time_delta > self.timeout:
146            err_message = '%s timeout after %s seconds' % (self.des, time_delta)
147            logging.warning(err_message)
148            raise TimeoutException(err_message)
149
150
151def run_single_memory_pressure(
152    starting_mb, step_mb, end_condition, duration, cool_down, timeout=None):
153    """Runs a single memory consumer to produce memory pressure.
154
155    Keep adding memory pressure. In each round, it runs a memory consumer
156    and waits for a while before checking whether to end the process. If not,
157    kill current memory consumer and allocate more memory pressure in the next
158    round.
159    @param starting_mb: The amount of memory to start with.
160    @param step_mb: If |end_condition| is not met, allocate |step_mb| more
161        memory in the next round.
162    @param end_condition: A boolean function returns whether to end the process.
163    @param duration: Time (in seconds) to wait between running a memory
164        consumer and checking |end_condition|.
165    @param cool_down:  Time (in seconds) to wait between each round.
166    @param timeout: Seconds to stop the function is |end_condition| is not met.
167    @return The size of memory allocated in the last round.
168    @raise TimeoutException if timeout.
169    """
170    current_mb = starting_mb
171    timer = _Timer(timeout, 'run_single_memory_pressure')
172    while True:
173        timer.check_timeout()
174        with MemoryEater() as consumer:
175            consumer.consume_active_memory(current_mb)
176            time.sleep(duration)
177            if end_condition():
178                return current_mb
179        current_mb += step_mb
180        time.sleep(cool_down)
181
182
183def run_multi_memory_pressure(size_mb, end_condition, duration, timeout=None):
184    """Runs concurrent memory consumers to produce memory pressure.
185
186    In each round, it runs a new memory consumer until a certain condition is
187    met.
188    @param size_mb: The amount of memory each memory consumer allocates.
189    @param end_condition: A boolean function returns whether to end the process.
190    @param duration: Time (in seconds) to wait between running a memory
191        consumer and checking |end_condition|.
192    @param timeout: Seconds to stop the function is |end_condition| is not met.
193    @return Total allocated memory.
194    @raise TimeoutException if timeout.
195    """
196    total_mb = 0
197    timer = _Timer(timeout, 'run_multi_memory_pressure')
198    with MemoryEater() as consumer:
199        while True:
200            timer.check_timeout()
201            consumer.consume_active_memory(size_mb)
202            time.sleep(duration)
203            if end_condition():
204                return total_mb
205            total_mb += size_mb
206