Clean up the stale patch packages from TiUP

clean_tiup_patch_caches.py

#!/usr/bin/env python3

from __future__ import print_function

__doc__ = """
Put it to the machine deploy TiUP and add it to the crontab.
e.g.

> crontab -e
# clean up the patch caches for TiUP every day
0 1 * * * /opt/clean_tiup_patch_caches.py > /opt/clean_tiup_patch_caches.log 2>&1 &

And you may need to use `service crond status` to ensure `crond` is running normally.
"""

import os
import shutil
import datetime

TIUP_ROOT = "/root/.tiup"
#TIUP_ROOT = "/DATA/disk1/ra_common/.tiup"
DRY_RUN = True
DRY_RUN = False

def get_size(begin_dir_name):
    # for file
    if os.path.isfile(begin_dir_name):
        return os.path.getsize(begin_dir_name)
    # for directory
    total_size = 0
    for dirpath, dirnames, filenames in os.walk(begin_dir_name):
        for f in filenames:
            fp = os.path.join(dirpath, f)
            # skip if it is symbolic link
            if not os.path.islink(fp):
                total_size += os.path.getsize(fp)
    return total_size

def clean_packages(cluster, full_paths, num_to_keep):
    full_paths_with_ctime = sorted(list(map(lambda p: (p, os.path.getmtime(p)), full_paths)), key=lambda p: p[1], reverse=True)
    # Keep the latest package
    for p in full_paths_with_ctime[:num_to_keep]:
        print("Keep the package: {}, created @ {}".format(p[0], datetime.datetime.fromtimestamp(p[1])))
    packs_to_clean = full_paths_with_ctime[num_to_keep:]
    print("Cleaning {} patch packages for cluster: {}".format(len(packs_to_clean), cluster))
    size_removed = 0
    for full_path in packs_to_clean:
        full_name, ctime = full_path
        size_removed += get_size(full_name)
        #print("sum size_removed: {} full_name {}".format(size_removed, full_name))
        if not DRY_RUN:
            if os.path.isdir(full_name): shutil.rmtree(full_name)
            else: os.remove(full_name)
        print("removing", full_name)
    return size_removed

def clean_cluster_package_subdir(cluster_name, sub_dir):
    sub_path = os.path.join(TIUP_ROOT, "storage/cluster/clusters", cluster_name, sub_dir)
    if not os.path.exists(sub_path): return 0
    cached_packages = os.listdir(sub_path)
    if len(cached_packages) == 0: return 0

    packs_by_comp = {}
    for p in cached_packages:
        # tidb-hash, tikv-hash, tiflash-hash
        component = p[:p.rfind('-')]
        full_path = os.path.join(TIUP_ROOT, "storage/cluster/clusters", cluster_name, sub_dir, p)
        if component in packs_by_comp:
            packs_by_comp[component].append(full_path)
        else:
            packs_by_comp[component] = [full_path]
    print("Cleaning outdated packages for cluster: {}, sub_dir: {}".format(cluster_name, sub_dir))
    cluster_size_removed = 0
    for c in packs_by_comp:
        cluster_size_removed += clean_packages(cluster_name, packs_by_comp[c], 1)
    return cluster_size_removed


def clean_cluster_packages():
    size_removed = 0
    cluster_names = os.listdir(os.path.join(TIUP_ROOT, "storage/cluster/clusters"))
    for n in cluster_names:
        # du -sch .tiup/storage/cluster/clusters/*/cache/*
        cluster_size_removed = clean_cluster_package_subdir(n, "cache")
        # du -sch .tiup/storage/cluster/clusters/*/patch/*
        cluster_size_removed += clean_cluster_package_subdir(n, "patch")
        size_removed += cluster_size_removed
        print("Clean outdated packages for cluster: {} done, free {:.2f} GiB\n".format(n, cluster_size_removed / 1024.0 / 1024 / 1024))
    print("All clean done, free {:.2f} GiB. @ {}".format(size_removed / 1024.0 / 1024 / 1024, datetime.datetime.now()))

def clean_nightly_comps():
    components = os.listdir(os.path.join(TIUP_ROOT, 'components'))
    for c in components:
        nightly_ctls = [
            x for x in os.listdir(os.path.join(TIUP_ROOT, 'components', c))
            if x.find('nightly') != -1
        ]
        size_removed = 0
        for n in nightly_ctls:
            abs_path = os.path.join(TIUP_ROOT, 'components', c, n)
            size_removed += get_size(abs_path)
            print("removing", abs_path)
            if not DRY_RUN: shutil.rmtree(abs_path)
        print("Clean outdated components {} done, free {:.2f} GiB. @ {}".format(c, size_removed / 1024.0 / 1024 / 1024, datetime.datetime.now()))
        print()


def main():
    clean_cluster_packages()
    clean_nightly_comps()

if __name__ == '__main__':
    main()