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