• 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
5"""This throttler tries to reduce result size by compress files to tgz file.
6"""
7
8import re
9import os
10import tarfile
11
12import throttler_lib
13import utils_lib
14
15
16# File with extensions that can not be zipped or compression won't reduce file
17# size further.
18UNZIPPABLE_EXTENSIONS = set([
19        '.gz',
20        '.jpg',
21        '.png',
22        '.tgz',
23        '.xz',
24        '.zip',
25        ])
26
27# Regex for files that should not be compressed.
28UNZIPPABLE_FILE_PATTERNS = [
29        'BUILD_INFO-.*' # ACTS test result files.
30        ]
31
32# Default threshold of file size in byte for it to be qualified for compression.
33# Files smaller than the threshold will not be compressed.
34DEFAULT_FILE_SIZE_THRESHOLD_BYTE = 100 * 1024
35
36def _zip_file(file_info):
37    """Zip the file to reduce the file size.
38
39    @param file_info: A ResultInfo object containing summary for the file to be
40            shrunk.
41    """
42    utils_lib.LOG('Compressing file %s' % file_info.path)
43    parent_result_info = file_info.parent_result_info
44    new_name = file_info.name + '.tgz'
45    new_path = os.path.join(os.path.dirname(file_info.path), new_name)
46    if os.path.exists(new_path):
47        utils_lib.LOG('File %s already exists, removing...' % new_path)
48        if not throttler_lib.try_delete_file_on_disk(new_path):
49            return
50        parent_result_info.remove_file(new_name)
51    with tarfile.open(new_path, 'w:gz') as tar:
52        tar.add(file_info.path, arcname=os.path.basename(file_info.path))
53    stat = os.stat(file_info.path)
54    if not throttler_lib.try_delete_file_on_disk(file_info.path):
55        # Clean up the intermediate file.
56        throttler_lib.try_delete_file_on_disk(new_path)
57        utils_lib.LOG('Failed to compress %s' % file_info.path)
58        return
59
60    # Modify the new file's timestamp to the old one.
61    os.utime(new_path, (stat.st_atime, stat.st_mtime))
62    # Get the original file size before compression.
63    original_size = file_info.original_size
64    parent_result_info.remove_file(file_info.name)
65    parent_result_info.add_file(new_name)
66    new_file_info = parent_result_info.get_file(new_name)
67    # Set the original size to be the size before compression.
68    new_file_info.original_size = original_size
69    # Set the trimmed size to be the physical file size of the compressed file.
70    new_file_info.trimmed_size = new_file_info.size
71
72
73def _get_zippable_files(file_infos, file_size_threshold_byte):
74    """Filter the files that can be throttled.
75
76    @param file_infos: A list of ResultInfo objects.
77    @param file_size_threshold_byte: Threshold of file size in byte for it to be
78            qualified for compression.
79    @yield: ResultInfo objects that can be shrunk.
80    """
81    for info in file_infos:
82        ext = os.path.splitext(info.name)[1].lower()
83        if ext in UNZIPPABLE_EXTENSIONS:
84            continue
85
86        match_found = False
87        for pattern in UNZIPPABLE_FILE_PATTERNS:
88            if re.match(pattern, info.name):
89                match_found = True
90                break
91        if match_found:
92            continue
93
94        if info.trimmed_size <= file_size_threshold_byte:
95            continue
96
97        yield info
98
99
100def throttle(summary, max_result_size_KB,
101             file_size_threshold_byte=DEFAULT_FILE_SIZE_THRESHOLD_BYTE,
102             skip_autotest_log=False):
103    """Throttle the files in summary by compressing file.
104
105    Stop throttling until all files are processed or the result file size is
106    already reduced to be under the given max_result_size_KB.
107
108    @param summary: A ResultInfo object containing result summary.
109    @param max_result_size_KB: Maximum test result size in KB.
110    @param file_size_threshold_byte: Threshold of file size in byte for it to be
111            qualified for compression.
112    @param skip_autotest_log: True to skip shrink Autotest logs, default is
113            False.
114    """
115    file_infos, _ = throttler_lib.sort_result_files(summary)
116    extra_patterns = ([throttler_lib.AUTOTEST_LOG_PATTERN] if skip_autotest_log
117                      else [])
118    file_infos = throttler_lib.get_throttleable_files(
119            file_infos, extra_patterns)
120    file_infos = _get_zippable_files(file_infos, file_size_threshold_byte)
121    for info in file_infos:
122        _zip_file(info)
123
124        if throttler_lib.check_throttle_limit(summary, max_result_size_KB):
125            return
126