| #!/usr/bin/python |
| # Copyright 2017 The Chromium OS Authors. All rights reserved. |
| # Use of this source code is governed by a BSD-style license that can be |
| # found in the LICENSE file. |
| |
| """ |
| This is a utility to build a summary of the given directory. and save to a json |
| file. |
| |
| Example usage: |
| result_utils.py -p path |
| |
| The content of the json file looks like: |
| {'default': {'/D': {'control': {'/S': 734}, |
| 'debug': {'/D': {'client.0.DEBUG': {'/S': 5698}, |
| 'client.0.ERROR': {'/S': 254}, |
| 'client.0.INFO': {'/S': 1020}, |
| 'client.0.WARNING': {'/S': 242}}, |
| '/S': 7214} |
| }, |
| '/S': 7948 |
| } |
| } |
| """ |
| |
| import argparse |
| import copy |
| import glob |
| import json |
| import logging |
| import os |
| import time |
| |
| # Do NOT import autotest_lib modules here. This module can be executed without |
| # dependency on other autotest modules. This is to keep the logic of result |
| # trimming on the server side, instead of depending on the autotest client |
| # module. |
| |
| DEFAULT_SUMMARY_FILENAME_FMT = 'dir_summary_%d.json' |
| # Minimum disk space should be available after saving the summary file. |
| MIN_FREE_DISK_BYTES = 10 * 1024 * 1024 |
| |
| # Following are key names for directory summaries. The keys are started with / |
| # so it can be differentiated with a valid file name. The short keys are |
| # designed for smaller file size of the directory summary. |
| |
| # Original size of the directory or file |
| ORIGINAL_SIZE_BYTES = '/S' |
| # Size of the directory or file after trimming |
| TRIMMED_SIZE_BYTES = '/T' |
| # Size of the directory or file being collected from client side |
| COLLECTED_SIZE_BYTES = '/C' |
| # A dictionary of sub-directories' summary: name: {directory_summary} |
| DIRS = '/D' |
| # Default root directory name. To allow summaries to be merged effectively, all |
| # summaries are collected with root directory of '' |
| ROOT_DIR = '' |
| |
| # Autotest uses some state files to track process running state. The files are |
| # deleted from test results. Therefore, these files can be ignored. |
| FILES_TO_IGNORE = set([ |
| 'control.autoserv.state' |
| ]) |
| |
| def get_unique_dir_summary_file(path): |
| """Get a unique file path to save the directory summary json string. |
| |
| @param path: The directory path to save the summary file to. |
| """ |
| summary_file = DEFAULT_SUMMARY_FILENAME_FMT % time.time() |
| # Make sure the summary file name is unique. |
| file_name = os.path.join(path, summary_file) |
| if os.path.exists(file_name): |
| count = 1 |
| name, ext = os.path.splitext(summary_file) |
| while os.path.exists(file_name): |
| file_name = os.path.join(path, '%s_%s%s' % (name, count, ext)) |
| count += 1 |
| return file_name |
| |
| |
| def get_dir_summary(path, top_dir, all_dirs=set()): |
| """Get the directory summary for the given path. |
| |
| @param path: The directory to collect summary. |
| @param top_dir: The top directory to collect summary. This is to check if a |
| directory is a subdir of the original directory to collect summary. |
| @param all_dirs: A set of paths that have been collected. This is to prevent |
| infinite recursive call caused by symlink. |
| |
| @return: A dictionary of the directory summary. |
| """ |
| dir_info = {} |
| dir_info[ORIGINAL_SIZE_BYTES] = 0 |
| summary = {os.path.basename(path): dir_info} |
| |
| if os.path.isfile(path): |
| dir_info[ORIGINAL_SIZE_BYTES] = os.stat(path).st_size |
| else: |
| dir_info[DIRS] = {} |
| real_path = os.path.realpath(path) |
| # The assumption here is that results are copied back to drone by |
| # copying the symlink, not the content, which is true with currently |
| # used rsync in cros_host.get_file call. |
| # Skip scanning the child folders if any of following condition is true: |
| # 1. The directory is a symlink and link to a folder under `top_dir`. |
| # 2. The directory was scanned already. |
| if ((os.path.islink(path) and real_path.startswith(top_dir)) or |
| real_path in all_dirs): |
| return summary |
| |
| all_dirs.add(real_path) |
| for f in sorted(os.listdir(path)): |
| f_summary = get_dir_summary(os.path.join(path, f), top_dir, |
| all_dirs) |
| dir_info[DIRS][f] = f_summary[f] |
| dir_info[ORIGINAL_SIZE_BYTES] += f_summary[f][ORIGINAL_SIZE_BYTES] |
| |
| return summary |
| |
| |
| def build_summary_json(path): |
| """Build summary of files in the given path and return a json string. |
| |
| @param path: The directory to build summary. |
| @return: A json string of the directory summary. |
| @raise IOError: If the given path doesn't exist. |
| """ |
| if not os.path.exists(path): |
| raise IOError('Path %s does not exist.' % path) |
| |
| return get_dir_summary(path, top_dir=path) |
| |
| |
| def _update_sizes(entry): |
| """Update a directory entry's sizes. |
| |
| Values of ORIGINAL_SIZE_BYTES, TRIMMED_SIZE_BYTES and COLLECTED_SIZE_BYTES |
| are re-calculated based on the files under the directory. If the entry is a |
| file, skip the updating. |
| |
| @param entry: A dict of directory entry in a summary. |
| """ |
| if DIRS not in entry: |
| return |
| |
| entry[ORIGINAL_SIZE_BYTES] = sum([entry[DIRS][s][ORIGINAL_SIZE_BYTES] |
| for s in entry[DIRS]]) |
| # Before trimming is implemented, COLLECTED_SIZE_BYTES and |
| # TRIMMED_SIZE_BYTES have the same value of ORIGINAL_SIZE_BYTES. |
| entry[COLLECTED_SIZE_BYTES] = sum([ |
| entry[DIRS][s].get( |
| COLLECTED_SIZE_BYTES, |
| entry[DIRS][s].get(TRIMMED_SIZE_BYTES, |
| entry[DIRS][s][ORIGINAL_SIZE_BYTES])) |
| for s in entry[DIRS]]) |
| entry[TRIMMED_SIZE_BYTES] = sum([ |
| entry[DIRS][s].get(TRIMMED_SIZE_BYTES, |
| entry[DIRS][s][ORIGINAL_SIZE_BYTES]) |
| for s in entry[DIRS]]) |
| |
| |
| def _delete_missing_entries(summary_old, summary_new): |
| """Delete files/directories only exists in old summary. |
| |
| When the new summary is final, i.e., it's built from the final result |
| directory, files or directories missing are considered to be deleted and |
| trimmed to size 0. |
| |
| @param summary_old: Old directory summary. |
| @param summary_new: New directory summary. |
| """ |
| for name in summary_old.keys(): |
| if name not in summary_new: |
| if DIRS in summary_old[name]: |
| # Trim sub-directories. |
| _delete_missing_entries(summary_old[name][DIRS], {}) |
| _update_sizes(summary_old[name]) |
| elif name in FILES_TO_IGNORE: |
| # Remove the file from the summary as it can be ignored. |
| del summary_old[name] |
| else: |
| # Before setting the trimmed size to 0, update the collected |
| # size if it's not set yet. |
| if COLLECTED_SIZE_BYTES not in summary_old[name]: |
| trimmed_size = summary_old[name].get( |
| TRIMMED_SIZE_BYTES, |
| summary_old[name][ORIGINAL_SIZE_BYTES]) |
| summary_old[name][COLLECTED_SIZE_BYTES] = trimmed_size |
| summary_old[name][TRIMMED_SIZE_BYTES] = 0 |
| elif DIRS in summary_old[name]: |
| _delete_missing_entries(summary_old[name][DIRS], |
| summary_new[name][DIRS]) |
| _update_sizes(summary_old[name]) |
| _update_sizes(summary_old) |
| |
| |
| def _merge(summary_old, summary_new, is_final=False): |
| """Merge a new directory summary to an old one. |
| |
| Update the old directory summary with the new summary. Also calculate the |
| total size of results collected from the client side. |
| |
| When merging with previously collected results, any results not existing in |
| the new summary or files with size different from the new files collected |
| are considered as extra results collected or overwritten by the new results. |
| Therefore, the size of the collected result should include such files, and |
| the COLLECTED_SIZE_BYTES can be larger than TRIMMED_SIZE_BYTES. |
| As an example: |
| summary_old: {'file1': {TRIMMED_SIZE_BYTES: 1000, |
| ORIGINAL_SIZE_BYTES: 1000, |
| COLLECTED_SIZE_BYTES: 1000}} |
| This means a result `file1` of original size 1KB was collected with size of |
| 1KB byte. |
| summary_new: {'file1': {TRIMMED_SIZE_BYTES: 1000, |
| ORIGINAL_SIZE_BYTES: 2000, |
| COLLECTED_SIZE_BYTES: 1000}} |
| This means a result `file1` of 2KB was trimmed down to 1KB and was collected |
| with size of 1KB byte. |
| Note that the second result collection has an updated result `file1` |
| (because of the different ORIGINAL_SIZE_BYTES), and it needs to be rsync-ed |
| to the drone. Therefore, the merged summary will be: |
| {'file1': {TRIMMED_SIZE_BYTES: 1000, |
| ORIGINAL_SIZE_BYTES: 2000, |
| COLLECTED_SIZE_BYTES: 2000}} |
| Note that: |
| * TRIMMED_SIZE_BYTES is still at 1KB, which reflects the actual size of the |
| file be collected. |
| * ORIGINAL_SIZE_BYTES is updated to 2KB, which is the size of the file in |
| the new result `file1`. |
| * COLLECTED_SIZE_BYTES is 2KB because rsync will copy `file1` twice as it's |
| changed. |
| |
| @param summary_old: Old directory summary. |
| @param summary_new: New directory summary. |
| @param is_final: True if summary_new is built from the final result folder. |
| Default is set to False. |
| @return: A tuple of (bytes_diff, merged_summary): |
| bytes_diff: The size of results collected based on the diff of the |
| old summary and the new summary. |
| merged_summary: Merged directory summary. |
| """ |
| for name in summary_new: |
| if not name in summary_old: |
| # A file/dir exists in new client dir, but not in the old one, which |
| # means that the file or a directory is newly collected. |
| summary_old[name] = copy.deepcopy(summary_new[name]) |
| elif DIRS in summary_new[name]: |
| # `name` is a directory in new summary, merge the directories of the |
| # old and new summaries under `name`. |
| |
| if DIRS not in summary_old[name]: |
| # If `name` is a file in old summary but a directory in new |
| # summary, the file in the old summary will be overwritten by |
| # the new directory by rsync. Therefore, force it to be an empty |
| # directory in old summary, so that the new directory can be |
| # merged. |
| summary_old[name][ORIGINAL_SIZE_BYTES] = 0 |
| summary_old[name][TRIMMED_SIZE_BYTES] = 0 |
| summary_old[name][COLLECTED_SIZE_BYTES] = 0 |
| summary_old[name][DIRS] = {} |
| |
| _merge(summary_old[name][DIRS], summary_new[name][DIRS], is_final) |
| else: |
| # `name` is a file. Compare the original size, if they are |
| # different, the file was overwritten, so increment the |
| # COLLECTED_SIZE_BYTES. |
| |
| if DIRS in summary_old[name]: |
| # If `name` is a directory in old summary, but a file in the new |
| # summary, rsync will fail to copy the file as it can't |
| # overwrite an directory. Therefore, skip the merge. |
| continue |
| |
| new_size = summary_new[name][ORIGINAL_SIZE_BYTES] |
| old_size = summary_old[name][ORIGINAL_SIZE_BYTES] |
| new_trimmed_size = summary_new[name].get( |
| TRIMMED_SIZE_BYTES, summary_new[name][ORIGINAL_SIZE_BYTES]) |
| old_trimmed_size = summary_old[name].get( |
| TRIMMED_SIZE_BYTES, summary_old[name][ORIGINAL_SIZE_BYTES]) |
| if new_size != old_size: |
| if is_final and new_trimmed_size == old_trimmed_size: |
| # If the file is merged from the final result folder to an |
| # older summary, it's not considered to be trimmed if the |
| # size is not changed. The reason is that the file on the |
| # server side does not have the info of its original size. |
| continue |
| |
| # Before trimming is implemented, COLLECTED_SIZE_BYTES is the |
| # value of ORIGINAL_SIZE_BYTES. |
| new_collected_size = summary_new[name].get( |
| COLLECTED_SIZE_BYTES, |
| summary_new[name].get( |
| TRIMMED_SIZE_BYTES, |
| summary_new[name][ORIGINAL_SIZE_BYTES])) |
| old_collected_size = summary_old[name].get( |
| COLLECTED_SIZE_BYTES, |
| summary_old[name].get( |
| TRIMMED_SIZE_BYTES, |
| summary_old[name][ORIGINAL_SIZE_BYTES])) |
| |
| summary_old[name][COLLECTED_SIZE_BYTES] = ( |
| new_collected_size + old_collected_size) |
| summary_old[name][TRIMMED_SIZE_BYTES] = summary_new[name].get( |
| TRIMMED_SIZE_BYTES, |
| summary_new[name][ORIGINAL_SIZE_BYTES]) |
| summary_old[name][ORIGINAL_SIZE_BYTES] = new_size |
| |
| # Update COLLECTED_SIZE_BYTES and ORIGINAL_SIZE_BYTES based on the |
| # merged directory summary. |
| _update_sizes(summary_old[name]) |
| |
| |
| def merge_summaries(path): |
| """Merge all directory summaries in the given path. |
| |
| This function calculates the total size of result files being collected for |
| the test device and the files generated on the drone. It also returns merged |
| directory summary. |
| |
| @param path: A path to search for directory summaries. |
| @return a tuple of (client_collected_bytes, merged_summary): |
| client_collected_bytes: The total size of results collected from |
| the DUT. The number can be larger than the total file size of |
| the given path, as files can be overwritten or removed. |
| merged_summary: The merged directory summary of the given path. |
| """ |
| # Find all directory summary files and sort them by the time stamp in file |
| # name. |
| summary_files = glob.glob(os.path.join(path, 'dir_summary_*.json')) |
| summary_files = sorted(summary_files, key=os.path.getmtime) |
| |
| all_summaries = [] |
| for summary_file in summary_files: |
| with open(summary_file) as f: |
| all_summaries.append(json.load(f)) |
| |
| # Merge all summaries. |
| merged_summary = (copy.deepcopy(all_summaries[0]) if len(all_summaries) > 0 |
| else {}) |
| for summary in all_summaries[1:]: |
| _merge(merged_summary, summary) |
| # After all summaries from the test device (client side) are merged, we can |
| # get the total size of result files being transfered from the test device. |
| # If there is no directory summary collected, default client_collected_bytes |
| # to 0. |
| client_collected_bytes = 0 |
| if merged_summary: |
| client_collected_bytes = merged_summary[ROOT_DIR][COLLECTED_SIZE_BYTES] |
| |
| # Get the summary of current directory |
| |
| # Make sure the path ends with /, so the top directory in the summary will |
| # be '', which is consistent with other summaries. |
| if not path.endswith(os.sep): |
| path += os.sep |
| |
| last_summary = get_dir_summary(path, top_dir=path) |
| _merge(merged_summary, last_summary, is_final=True) |
| _delete_missing_entries(merged_summary, last_summary) |
| |
| return client_collected_bytes, merged_summary |
| |
| |
| def main(): |
| """main script. """ |
| parser = argparse.ArgumentParser() |
| parser.add_argument('-p', type=str, dest='path', |
| help='Path to build directory summary.') |
| options = parser.parse_args() |
| |
| summary = build_summary_json(options.path) |
| summary_json = json.dumps(summary) |
| summary_file = get_unique_dir_summary_file(options.path) |
| |
| # Make sure there is enough free disk to write the file |
| stat = os.statvfs(options.path) |
| free_space = stat.f_frsize * stat.f_bavail |
| if free_space - len(summary_json) < MIN_FREE_DISK_BYTES: |
| raise IOError('Not enough disk space after saving the summary file. ' |
| 'Available free disk: %s bytes. Summary file size: %s ' |
| 'bytes.' % (free_space, len(summary_json))) |
| |
| with open(summary_file, 'w') as f: |
| f.write(summary_json) |
| logging.info('Directory summary of %s is saved to file %s.', options.path, |
| summary_file) |
| |
| |
| if __name__ == '__main__': |
| main() |
