Home | History | Annotate | Download | only in result_tools 
      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 
      8 import re
      9 import os
     10 import tarfile
     11 
     12 import throttler_lib
     13 import utils_lib
     14 
     15 
     16 # File with extensions that can not be zipped or compression won't reduce file
     17 # size further.
     18 UNZIPPABLE_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.
     28 UNZIPPABLE_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.
     34 DEFAULT_FILE_SIZE_THRESHOLD_BYTE = 100 * 1024
     35 
     36 def _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 
     73 def _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 
    100 def 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