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"""Help functions used by different throttlers."""
6
7import os
8import re
9
10try:
11    from autotest_lib.client.bin.result_tools import utils_lib
12except ImportError:
13    import utils_lib
14
15
16# A list of file names that should not be throttled, that is, not modified by
17# deletion, trimming or compression.
18NON_THROTTLEABLE_FILE_NAMES = set([
19        '.autoserv_execute',
20        '.parse.lock',
21        '.parse.log',
22        '.parser_execute',
23        'control',
24        'control.srv',
25        'host_keyvals',
26        'job_report.html',
27        'keyval',
28        'profiling',
29        'result_summary.html',
30        'sponge_invocation.xml',
31        'status',
32        'status.log',
33
34        # ACTS related files:
35        'test_run_details.txt',
36        'test_run_error.txt',
37        'test_run_info.txt',
38        'test_run_summary.json',
39        ])
40
41# A list of file name patterns that should not be throttled, that is, not
42# modified by deletion, deduping, trimming or compression.
43NON_THROTTLEABLE_FILE_PATTERNS = [
44        '.*/BUILD_INFO-.*',   # ACTS test result files.
45        '.*/AndroidDevice.*', # ACTS test result files.
46        ]
47
48# Regex of result files sorted based on priority. Files can be throttled first
49# have higher priority.
50RESULT_THROTTLE_PRIORITY = [
51        '(.*/)?sysinfo/var/log/.*',
52        '(.*/)?sysinfo/var/log_diff/.*',
53        '(.*/)?sysinfo/.*',
54        # The last one matches any file.
55        '.*',
56        ]
57
58# Regex of file names for Autotest debug logs. These files should be preserved
59# without throttling if possible.
60AUTOTEST_LOG_PATTERN ='.*\.(DEBUG|ERROR|INFO|WARNING)$'
61
62def _list_files(files, all_files=None):
63    """Get all files in the given directories.
64
65    @param files: A list of ResultInfo objects for files in a directory.
66    @param all_files: A list of ResultInfo objects collected so far.
67    @return: A list of all collected ResultInfo objects.
68    """
69    if all_files is None:
70        all_files = []
71    for info in files:
72        if info.is_dir:
73            _list_files(info.files, all_files)
74        else:
75            all_files.append(info)
76    return all_files
77
78
79def sort_result_files(summary):
80    """Sort result infos based on priority.
81
82    @param summary: A ResultInfo object containing result summary.
83    @return: A tuple of (sorted_files, grouped_files)
84            sorted_files: A list of ResultInfo, sorted based on file size and
85                priority based on RESULT_THROTTLE_PRIORITY.
86            grouped_files: A dictionary of ResultInfo grouped by each item in
87                RESULT_THROTTLE_PRIORITY.
88    """
89    all_files = _list_files(summary.files)
90
91    # Scan all file paths and group them based on the throttle priority.
92    grouped_files = {pattern: [] for pattern in RESULT_THROTTLE_PRIORITY}
93    for info in all_files:
94        for pattern in RESULT_THROTTLE_PRIORITY:
95            if re.match(pattern, info.path):
96                grouped_files[pattern].append(info)
97                break
98
99    sorted_files = []
100    for pattern in RESULT_THROTTLE_PRIORITY:
101        # Sort the files in each group by file size, largest first.
102        infos = grouped_files[pattern]
103        infos.sort(key=lambda info: -info.trimmed_size)
104        sorted_files.extend(infos)
105
106    return sorted_files, grouped_files
107
108
109def get_throttleable_files(file_infos, extra_patterns=[]):
110    """Filter the files can be throttled.
111
112    @param file_infos: A list of ResultInfo objects.
113    @param extra_patterns: Extra patterns of file path that should not be
114            throttled.
115    @yield: ResultInfo objects that can be throttled.
116    """
117    for info in file_infos:
118        # Skip the files being deleted in earlier throttling.
119        if info.trimmed_size == 0:
120            continue
121        if info.name in NON_THROTTLEABLE_FILE_NAMES:
122            continue
123        pattern_matched = False
124        for pattern in extra_patterns + NON_THROTTLEABLE_FILE_PATTERNS:
125            if re.match(pattern, info.path):
126                pattern_matched = True
127                break
128
129        if not pattern_matched:
130            yield info
131
132
133def check_throttle_limit(summary, max_result_size_KB):
134    """Check if throttling is enough already.
135
136    @param summary: A ResultInfo object containing result summary.
137    @param max_result_size_KB: Maximum test result size in KB.
138    @return: True if the result directory has been trimmed to be smaller than
139            max_result_size_KB.
140    """
141    if (summary.trimmed_size <= max_result_size_KB * 1024):
142        utils_lib.LOG('Maximum result size is reached (current result'
143                      'size is %s (limit is %s).' %
144                      (utils_lib.get_size_string(summary.trimmed_size),
145                       utils_lib.get_size_string(max_result_size_KB * 1024)))
146        return True
147    else:
148        return False
149
150
151def try_delete_file_on_disk(path):
152    """Try to delete the give file on disk.
153
154    @param path: Path to the file.
155    @returns: True if the file is deleted, False otherwise.
156    """
157    try:
158        utils_lib.LOG('Deleting file %s.' % path)
159        os.remove(path)
160        return True
161    except OSError as e:
162        utils_lib.LOG('Failed to delete file %s, Error: %s' % (path, e))