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