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