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