1# Copyright 2017 The Chromium OS Authors. All rights reserved. 2# Use of this source code is governed by a BSD-style license that can be 3# found in the LICENSE file. 4 5"""This throttler tries to reduce result size by compress files to tgz file. 6""" 7 8import re 9import os 10import tarfile 11 12try: 13 from autotest_lib.client.bin.result_tools import throttler_lib 14 from autotest_lib.client.bin.result_tools import utils_lib 15except ImportError: 16 import throttler_lib 17 import utils_lib 18 19 20# File with extensions that can not be zipped or compression won't reduce file 21# size further. 22UNZIPPABLE_EXTENSIONS = set([ 23 '.gz', 24 '.jpg', 25 '.png', 26 '.tgz', 27 '.xz', 28 '.zip', 29 ]) 30 31# Regex for files that should not be compressed. 32UNZIPPABLE_FILE_PATTERNS = [ 33 'BUILD_INFO-.*' # ACTS test result files. 34 ] 35 36# Default threshold of file size in byte for it to be qualified for compression. 37# Files smaller than the threshold will not be compressed. 38DEFAULT_FILE_SIZE_THRESHOLD_BYTE = 100 * 1024 39 40def _zip_file(file_info): 41 """Zip the file to reduce the file size. 42 43 @param file_info: A ResultInfo object containing summary for the file to be 44 shrunk. 45 """ 46 utils_lib.LOG('Compressing file %s' % file_info.path) 47 parent_result_info = file_info.parent_result_info 48 new_name = file_info.name + '.tgz' 49 new_path = os.path.join(os.path.dirname(file_info.path), new_name) 50 if os.path.exists(new_path): 51 utils_lib.LOG('File %s already exists, removing...' % new_path) 52 if not throttler_lib.try_delete_file_on_disk(new_path): 53 return 54 parent_result_info.remove_file(new_name) 55 with tarfile.open(new_path, 'w:gz') as tar: 56 tar.add(file_info.path, arcname=os.path.basename(file_info.path)) 57 stat = os.stat(file_info.path) 58 if not throttler_lib.try_delete_file_on_disk(file_info.path): 59 # Clean up the intermediate file. 60 throttler_lib.try_delete_file_on_disk(new_path) 61 utils_lib.LOG('Failed to compress %s' % file_info.path) 62 return 63 64 # Modify the new file's timestamp to the old one. 65 os.utime(new_path, (stat.st_atime, stat.st_mtime)) 66 # Get the original file size before compression. 67 original_size = file_info.original_size 68 parent_result_info.remove_file(file_info.name) 69 parent_result_info.add_file(new_name) 70 new_file_info = parent_result_info.get_file(new_name) 71 # Set the original size to be the size before compression. 72 new_file_info.original_size = original_size 73 # Set the trimmed size to be the physical file size of the compressed file. 74 new_file_info.trimmed_size = new_file_info.size 75 76 77def _get_zippable_files(file_infos, file_size_threshold_byte): 78 """Filter the files that can be throttled. 79 80 @param file_infos: A list of ResultInfo objects. 81 @param file_size_threshold_byte: Threshold of file size in byte for it to be 82 qualified for compression. 83 @yield: ResultInfo objects that can be shrunk. 84 """ 85 for info in file_infos: 86 ext = os.path.splitext(info.name)[1].lower() 87 if ext in UNZIPPABLE_EXTENSIONS: 88 continue 89 90 match_found = False 91 for pattern in UNZIPPABLE_FILE_PATTERNS: 92 if re.match(pattern, info.name): 93 match_found = True 94 break 95 if match_found: 96 continue 97 98 if info.trimmed_size <= file_size_threshold_byte: 99 continue 100 101 yield info 102 103 104def throttle(summary, max_result_size_KB, 105 file_size_threshold_byte=DEFAULT_FILE_SIZE_THRESHOLD_BYTE, 106 skip_autotest_log=False): 107 """Throttle the files in summary by compressing file. 108 109 Stop throttling until all files are processed or the result file size is 110 already reduced to be under the given max_result_size_KB. 111 112 @param summary: A ResultInfo object containing result summary. 113 @param max_result_size_KB: Maximum test result size in KB. 114 @param file_size_threshold_byte: Threshold of file size in byte for it to be 115 qualified for compression. 116 @param skip_autotest_log: True to skip shrink Autotest logs, default is 117 False. 118 """ 119 file_infos, _ = throttler_lib.sort_result_files(summary) 120 extra_patterns = ([throttler_lib.AUTOTEST_LOG_PATTERN] if skip_autotest_log 121 else []) 122 file_infos = throttler_lib.get_throttleable_files( 123 file_infos, extra_patterns) 124 file_infos = _get_zippable_files(file_infos, file_size_threshold_byte) 125 for info in file_infos: 126 _zip_file(info) 127 128 if throttler_lib.check_throttle_limit(summary, max_result_size_KB): 129 return 130