blob: fd56e10e328dfedd72e7fae491cea26d0aa2fb35 [file] [log] [blame]
# Copyright 2018 The Chromium OS Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
import logging
import subprocess
import time
import threading
from autotest_lib.client.bin import utils
class MemoryEater(object):
"""A util class which run programs to consume memory in the background.
Sample usage:
with MemoryEator() as memory_eater:
# Allocate mlocked memory.
# Allocate memory and sequentially traverse them over and over.
When it goes out of the "with" context or the object is destructed, all
allocated memory are released.
memory_eater_locked = 'memory-eater-locked'
memory_eater = 'memory-eater'
_all_instances = []
def __init__(self):
self._locked_consumers = []
self._active_consumers_lock = threading.Lock()
self._active_consumers = []
def __enter__(self):
return self
def cleanup_consumers(consumers):
"""Kill all processes in |consumers|
@param consumers: The list of consumers to clean.
while len(consumers):
job = consumers.pop()'Killing %d',
def cleanup(self):
"""Releases all allocated memory."""
# Kill all hanging jobs.'Cleaning hanging memory consuming processes...')
with self._active_consumers_lock:
def __exit__(self, type, value, traceback):
def __del__(self):
if self in self._all_instances:
def consume_locked_memory(self, mb):
"""Consume non-swappable memory."""'Consuming locked memory %d MB', mb)
cmd = [self.memory_eater_locked, str(mb)]
p = subprocess.Popen(cmd, stdout=subprocess.PIPE)
# Wait until memory allocation is done.
while True:
line = p.stdout.readline()
if line.find('Done') != -1:
def consume_active_memory(self, mb):
"""Consume active memory."""'Consuming active memory %d MB', mb)
cmd = [self.memory_eater, '--size', str(mb), '--chunk', '128']
p = subprocess.Popen(cmd, stdout=subprocess.PIPE)
with self._active_consumers_lock:
def get_active_consumer_pids(cls):
"""Gets pid of active consumers by all instances of the class."""
all_pids = []
for instance in cls._all_instances:
with instance._active_consumers_lock:
all_pids.extend([ for p in instance._active_consumers])
return all_pids
def consume_free_memory(memory_to_reserve_mb):
"""Consumes free memory until |memory_to_reserve_mb| is remained.
Non-swappable memory is allocated to consume memory.
memory_to_reserve_mb: Consume memory until this amount of free memory
is remained.
@return The MemoryEater() object on which memory is allocated. One can
catch it in a context manager.
consumer = MemoryEater()
while True:
mem_free_mb = utils.read_from_meminfo('MemFree') / 1024'Current Free Memory %d', mem_free_mb)
if mem_free_mb <= memory_to_reserve_mb:
memory_to_consume = min(
2047, mem_free_mb - memory_to_reserve_mb + 1)'Consuming %d MB locked memory', memory_to_consume)
return consumer
class TimeoutException(Exception):
"""Exception to return if timeout happens."""
def __init__(self, message):
super(TimeoutException, self).__init__(message)
class _Timer(object):
"""A simple timer class to check timeout."""
def __init__(self, timeout, des):
@param timeout: Timeout in seconds.
@param des: A short description for this timer.
self.timeout = timeout
self.des = des
if self.timeout:
self.start_time = time.time()
def check_timeout(self):
"""Raise TimeoutException if timeout happens."""
if not self.timeout:
time_delta = time.time() - self.start_time
if time_delta > self.timeout:
err_message = '%s timeout after %s seconds' % (self.des, time_delta)
raise TimeoutException(err_message)
def run_single_memory_pressure(
starting_mb, step_mb, end_condition, duration, cool_down, timeout=None):
"""Runs a single memory consumer to produce memory pressure.
Keep adding memory pressure. In each round, it runs a memory consumer
and waits for a while before checking whether to end the process. If not,
kill current memory consumer and allocate more memory pressure in the next
@param starting_mb: The amount of memory to start with.
@param step_mb: If |end_condition| is not met, allocate |step_mb| more
memory in the next round.
@param end_condition: A boolean function returns whether to end the process.
@param duration: Time (in seconds) to wait between running a memory
consumer and checking |end_condition|.
@param cool_down: Time (in seconds) to wait between each round.
@param timeout: Seconds to stop the function is |end_condition| is not met.
@return The size of memory allocated in the last round.
@raise TimeoutException if timeout.
current_mb = starting_mb
timer = _Timer(timeout, 'run_single_memory_pressure')
while True:
with MemoryEater() as consumer:
if end_condition():
return current_mb
current_mb += step_mb
def run_multi_memory_pressure(size_mb, end_condition, duration, timeout=None):
"""Runs concurrent memory consumers to produce memory pressure.
In each round, it runs a new memory consumer until a certain condition is
@param size_mb: The amount of memory each memory consumer allocates.
@param end_condition: A boolean function returns whether to end the process.
@param duration: Time (in seconds) to wait between running a memory
consumer and checking |end_condition|.
@param timeout: Seconds to stop the function is |end_condition| is not met.
@return Total allocated memory.
@raise TimeoutException if timeout.
total_mb = 0
timer = _Timer(timeout, 'run_multi_memory_pressure')
with MemoryEater() as consumer:
while True:
if end_condition():
return total_mb
total_mb += size_mb