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