From ee0c74b9bde5b9f1db15f7b1f377d353343c2192 Mon Sep 17 00:00:00 2001 From: Daniel Smith Date: Mon, 30 May 2022 11:41:42 +0200 Subject: Add ability to prune branches from state file Without pruning, the state file will grow quite large over time. Add a parameter to allow pruning of all or some branches from the state file to minimize this issue. Change-Id: I61155f4aa11a36bfd6da151173fb0b694fc7103a Reviewed-by: Daniel Smith --- util/dependency_updater/main.py | 10 +++++ util/dependency_updater/tools/state.py | 82 +++++++++++++++++++++++----------- 2 files changed, 66 insertions(+), 26 deletions(-) diff --git a/util/dependency_updater/main.py b/util/dependency_updater/main.py index 5349ad7..42d0752 100644 --- a/util/dependency_updater/main.py +++ b/util/dependency_updater/main.py @@ -17,6 +17,12 @@ def parse_args(print_help: bool = False) -> Namespace: "the final state.") parser.add_argument('--reset', dest='reset', action='store_true', help="Forget current update state of [branch], then exit. Requires branch.") + parser.add_argument('--prune_state', dest="prune_and_keep", + help="Prune the saved state of unused branches, then exit.\n" + "Exclusive with all other options.\n" + "Pass a comma-separated list of branches to keep." + "Everything else gets pruned.\n" + "Pass keyword 'ALL' to clear all branches.") parser.add_argument('--pauseOnFail', dest='pause_on_finish_fail', action='store_true', help="If the round finished with failures in blocking repos, do not reset\n" "the round. Hold the current state until rewound or reset.") @@ -98,6 +104,10 @@ def main(): config = Config._load_config("config.yaml", parse_args()) config.datasources.load_datasources(config) config.state_repo = state.check_create_local_repo(config) + if config.args.prune_and_keep: + print(config.args.prune_and_keep) + state.clear_state(config, config.args.prune_and_keep) + exit() if config.args.reset: state.clear_state(config) exit() diff --git a/util/dependency_updater/tools/state.py b/util/dependency_updater/tools/state.py index 9a15b69..f263f0b 100644 --- a/util/dependency_updater/tools/state.py +++ b/util/dependency_updater/tools/state.py @@ -99,7 +99,8 @@ def update_state_data(old_state: dict[str, Repo], return updated -def save_updates_state(config, _clear_state: bool = False) -> None: +def save_updates_state(config, _clear_state: bool = False, + prune_and_keep: str = None) -> None: """Save updates to the state file""" if not config.args.simulate: if _clear_state: @@ -110,43 +111,72 @@ def save_updates_state(config, _clear_state: bool = False) -> None: data: dict[str, dict[str, Repo]] = {} with open(state_path, 'rb') as state_file: data = pickle.load(state_file) - data[config.args.branch] = config.state_data + if prune_and_keep: + keep_branches = prune_and_keep.split(",") + print(f"Found branches: {data.keys()}") + print(f"Retaining branches: {keep_branches}") + for branch in list(data.keys()): + if branch not in keep_branches: + del data[branch] + else: + data[config.args.branch] = config.state_data with open(state_path, 'wb') as state_file: pickle.dump(data, state_file) config.state_repo.index.add("state.bin") config.state_repo.index.commit("Update state") if config._state_ref: config.state_repo.remotes.origin.push(['-f', f"HEAD:{config._state_ref}"]) + print("Done saving state.") elif config.args.no_state: print("Running in no-state mode. Not saving state!") -def clear_state(config) -> None: - """Clear state data. All branches are wiped if not specified!""" - print("Clearing state and resetting updates...") - if config.args.branch: - config.state_data = {} - save_updates_state(config) - print(f"Clearing branch state for {config.args.branch}") - return +def clear_state(config, prune_and_keep: str = None) -> None: + """Clear state data. Currently operating branch is cleared if not pruning! - if config._state_ref: - try: - config.state_repo.remotes.origin.push(['-f', f":{config._state_ref}"]) - print("Cleared remote state on codereview...") - except git.exc.GitCommandError: - print( - "WARN: Failed to push an empty commit, probably because the state is already clear.") - del config.state_repo # Need to tear down the instance of PyGit to close the file handle. - sleep(5) # workaround for sometimes slow closing of git handles. + Parameters: + prune_and_keep: str: List of branches to keep. + Prune all others. Set to "ALL" clears all branches. + """ + + def clear_one(): + print(f"Clearing state and resetting updates for {config.args.branch}...") + if config.args.branch: + config.state_data = {} + save_updates_state(config) + print(f"Clearing branch state for {config.args.branch}") + return + + def clear_all(): + if config._state_ref: + try: + config.state_repo.remotes.origin.push(['-f', f":{config._state_ref}"]) + print("Cleared remote state on codereview...") + except git.exc.GitCommandError: + print( + "WARN: Failed to push an empty commit, probably because the state is already clear.") + del config.state_repo # Need to tear down the instance of PyGit to close the file handle. + sleep(5) # workaround for sometimes slow closing of git handles. + else: + print("\nWARN: No state remote ref set! Only deleting local state.bin file.\n" + "WARN: Run this script again with --reset after configuring an ssh user\n" + "WARN: in ~/.ssh/config for your gerrit host as set by 'GERRIT_HOST' in config.yaml.\n" + "WARN: If a remote state exists next time this script is run, it will likely\n" + "WARN: cause unexpected behavior!") + shutil.rmtree(Path(config.cwd, "_state"), onerror=_unlink_file) + print("Deleted local state files.") + + def clear_some(): + save_updates_state(config, prune_and_keep=prune_and_keep) + + if prune_and_keep: + print("Pruning state data.") + if prune_and_keep == "ALL": + clear_all() + else: + clear_some() else: - print("\nWARN: No state remote ref set! Only deleting local state.bin file.\n" - "WARN: Run this script again with --reset after configuring an ssh user\n" - "WARN: in ~/.ssh/config for your gerrit host as set by 'GERRIT_HOST' in config.yaml.\n" - "WARN: If a remote state exists next time this script is run, it will likely\n" - "WARN: cause unexpected behavior!") - shutil.rmtree(Path(config.cwd, "_state"), onerror=_unlink_file) - print("Deleted local state files.") + clear_one() def _unlink_file(function, path, excinfo): -- cgit v1.2.1