• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
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
5import os
6import re
7
8try:
9    from autotest_lib.client.bin.result_tools import throttler_lib
10    from autotest_lib.client.bin.result_tools import utils_lib
11except ImportError:
12    import throttler_lib
13    import utils_lib
14
15
16# File extensions that can be safely shrunk.
17# Extension matching is case-insensitive but the items in this set must be
18# lowercase to match.
19# Files without an extension and with no alphabetic characters in the extension
20# (e.g. file.20201110) are always shrinkable.
21SHRINKABLE_EXTENSIONS = frozenset([
22        '.log',
23        '.txt',
24        '.debug',
25        '.error',
26        '.info',
27        '.warning',
28])
29
30# Regex for paths that should not be shrunk.
31UNSHRINKABLE_PATH_PATTERNS = [
32        # Files in a log_diff/ directory should already be relatively small,
33        # and trimming them further would be detrimental to debugging. If
34        # they're too large, let other throttlers (e.g., zip_file_ or
35        # delete_file_) deal with them.
36        # Only blocklist a few known-useful log_diff's.
37        '/log_diff/messages$',
38        '/log_diff/net\.log$',
39        # Ramoops files are small but relatively important.
40        # The name of this file has changed starting with linux-3.19.
41        # Use a glob to match all existing records.
42        '/console-ramoops.*',
43        ]
44
45TRIMMED_FILE_HEADER = '!!! This file is trimmed !!!\n'
46ORIGINAL_SIZE_TEMPLATE = 'Original size: %d bytes\n\n'
47# Regex pattern to retrieve the original size of the file.
48ORIGINAL_SIZE_REGEX = 'Original size: (\d+) bytes'
49TRIMMED_FILE_INJECT_TEMPLATE = """
50
51========================================================================
52  < %d > characters are trimmed here.
53========================================================================
54
55"""
56
57# Percent of file content to keep at the beginning and end of the file, default
58# to 20%.
59HEAD_SIZE_PERCENT = 0.20
60
61# Default size in byte to trim the file down to.
62DEFAULT_FILE_SIZE_LIMIT_BYTE = 100 * 1024
63
64def _trim_file(file_info, file_size_limit_byte):
65    """Remove the file content in the middle to reduce the file size.
66
67    @param file_info: A ResultInfo object containing summary for the file to be
68            shrunk.
69    @param file_size_limit_byte: Maximum file size in bytes after trimming.
70    """
71    utils_lib.LOG('Trimming file %s to reduce size from %d bytes to %d bytes' %
72                  (file_info.path, file_info.original_size,
73                   file_size_limit_byte))
74    new_path = os.path.join(os.path.dirname(file_info.path),
75                            file_info.name + '_trimmed')
76    original_size_bytes = file_info.original_size
77    with open(new_path, 'w') as new_file, open(file_info.path) as old_file:
78        # Read the beginning part of the old file, if it's already started with
79        # TRIMMED_FILE_HEADER, no need to add the header again.
80        header =  old_file.read(len(TRIMMED_FILE_HEADER))
81        if header != TRIMMED_FILE_HEADER:
82            new_file.write(TRIMMED_FILE_HEADER)
83            new_file.write(ORIGINAL_SIZE_TEMPLATE % file_info.original_size)
84        else:
85            line = old_file.readline()
86            match = re.match(ORIGINAL_SIZE_REGEX, line)
87            if match:
88                original_size_bytes = int(match.group(1))
89        header_size_bytes = new_file.tell()
90        # Move old file reader to the beginning of the file.
91        old_file.seek(0, os.SEEK_SET)
92
93        new_file.write(old_file.read(
94                int((file_size_limit_byte - header_size_bytes) *
95                    HEAD_SIZE_PERCENT)))
96        # Position to seek from the end of the file.
97        seek_pos = -(file_size_limit_byte - new_file.tell() -
98                     len(TRIMMED_FILE_INJECT_TEMPLATE))
99        bytes_to_skip = original_size_bytes + seek_pos - old_file.tell()
100        # Adjust seek position based on string TRIMMED_FILE_INJECT_TEMPLATE
101        seek_pos += len(str(bytes_to_skip)) - 2
102        bytes_to_skip = original_size_bytes + seek_pos - old_file.tell()
103        new_file.write(TRIMMED_FILE_INJECT_TEMPLATE % bytes_to_skip)
104        old_file.seek(seek_pos, os.SEEK_END)
105        new_file.write(old_file.read())
106    stat = os.stat(file_info.path)
107    if not throttler_lib.try_delete_file_on_disk(file_info.path):
108        # Clean up the intermediate file.
109        throttler_lib.try_delete_file_on_disk(new_path)
110        utils_lib.LOG('Failed to shrink %s' % file_info.path)
111        return
112
113    os.rename(new_path, file_info.path)
114    # Modify the new file's timestamp to the old one.
115    os.utime(file_info.path, (stat.st_atime, stat.st_mtime))
116    # Update the trimmed_size.
117    file_info.trimmed_size = file_info.size
118
119
120def _get_shrinkable_files(file_infos, file_size_limit_byte):
121    """Filter the files that can be throttled.
122
123    @param file_infos: A list of ResultInfo objects.
124    @param file_size_limit_byte: Minimum file size in bytes to be throttled.
125    @yield: ResultInfo objects that can be shrunk.
126    """
127    for info in file_infos:
128        ext = os.path.splitext(info.name)[1].lower()
129        # if ext contains alphabetic characters and is not in the allowlist,
130        # skip the file.
131        # islower() returns false if the string does not contain any alphabetic
132        # characters, e.g. '.20201110'.islower() is False.
133        if ext.islower() and ext not in SHRINKABLE_EXTENSIONS:
134            continue
135
136        match_found = False
137        for pattern in UNSHRINKABLE_PATH_PATTERNS:
138            if re.search(pattern, info.path):
139                match_found = True
140                break
141        if match_found:
142            continue
143
144        if info.trimmed_size <= file_size_limit_byte:
145            continue
146
147        yield info
148
149
150def throttle(summary, max_result_size_KB,
151             file_size_limit_byte=DEFAULT_FILE_SIZE_LIMIT_BYTE,
152             skip_autotest_log=False):
153    """Throttle the files in summary by trimming file content.
154
155    Stop throttling until all files are processed or the result file size is
156    already reduced to be under the given max_result_size_KB.
157
158    @param summary: A ResultInfo object containing result summary.
159    @param max_result_size_KB: Maximum test result size in KB.
160    @param file_size_limit_byte: Limit each file's size in the summary to be
161            under the given threshold, until all files are processed or the
162            result size is under the given max_result_size_KB.
163    @param skip_autotest_log: True to skip shrink Autotest logs, default is
164            False.
165    """
166    file_infos, _ = throttler_lib.sort_result_files(summary)
167    extra_patterns = ([throttler_lib.AUTOTEST_LOG_PATTERN] if skip_autotest_log
168                      else [])
169    file_infos = throttler_lib.get_throttleable_files(
170            file_infos, extra_patterns)
171    file_infos = _get_shrinkable_files(file_infos, file_size_limit_byte)
172    for info in file_infos:
173        _trim_file(info, file_size_limit_byte)
174
175        if throttler_lib.check_throttle_limit(summary, max_result_size_KB):
176            return
177