From 3534e328c536fbe8a0a1da4f29df62d1eae95da9 Mon Sep 17 00:00:00 2001
From: Harshit Bansal <harshitbansal2015@gmail.com>
Date: Fri, 18 Aug 2017 19:42:06 +0000
Subject: [PATCH] clean-venv-cache: Expand `clean-venv-cache` tool.

Expands `clean-venv-cache` tool so that it can be used in prod.
Also moves it from `tools/` to `scripts/`.
---
 .travis.yml              |  2 +-
 scripts/clean-venv-cache | 68 ++++++++++++++++++++++++++++++++++++++++
 tools/clean-venv-cache   | 28 -----------------
 3 files changed, 69 insertions(+), 29 deletions(-)
 create mode 100755 scripts/clean-venv-cache
 delete mode 100755 tools/clean-venv-cache

diff --git a/.travis.yml b/.travis.yml
index 7d9348d74a..a02bfaeff3 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -16,7 +16,7 @@ install:
 
   # Clean any virtualenvs that are not in use to avoid our cache
   # becoming huge.
-  - mispipe "tools/clean-venv-cache --travis" ts
+  - mispipe "scripts/clean-venv-cache --threshold 0" ts
 
   # Clean any npm cache which was generated but not required anymore to avoid
   # our cache becoming huge.
diff --git a/scripts/clean-venv-cache b/scripts/clean-venv-cache
new file mode 100755
index 0000000000..26488cbff6
--- /dev/null
+++ b/scripts/clean-venv-cache
@@ -0,0 +1,68 @@
+#!/usr/bin/env python3
+import argparse
+import datetime
+import os
+import subprocess
+import sys
+import time
+
+if False:
+    from typing import Set, Text
+
+ZULIP_PATH = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+sys.path.append(ZULIP_PATH)
+from scripts.lib.zulip_tools import GENERIC_CACHE_SCRIPT_PARSER, \
+    get_caches_to_be_purged, get_environment, get_recent_deployments
+
+ENV = get_environment()
+VENV_CACHE_DIR = '/srv/zulip-venv-cache'
+if ENV == "travis":
+    VENV_CACHE_DIR = os.path.join(os.environ["HOME"], "zulip-venv-cache")
+
+def parse_args():
+    # type: () -> argparse.Namespace
+    parser = argparse.ArgumentParser(description="This script cleans unused zulip venv caches.",
+                                     parents=[GENERIC_CACHE_SCRIPT_PARSER, ])
+    args = parser.parse_args()
+    return args
+
+def get_caches_in_use(threshold_days):
+    # type: (int) -> Set[Text]
+    setups_to_check = set([ZULIP_PATH, ])
+    caches_in_use = set()
+    if ENV == "prod":
+        setups_to_check |= get_recent_deployments(threshold_days)
+    hash_reqs = os.path.join(ZULIP_PATH, 'scripts', 'lib', 'hash_reqs.py')
+    for path in setups_to_check:
+        for filename in os.listdir(os.path.join(path, "requirements")):
+            requirements_file = os.path.join(path, "requirements", filename)
+            hash_val = subprocess.check_output([hash_reqs, requirements_file]).strip()
+            caches_in_use.add(os.path.join(VENV_CACHE_DIR, hash_val.decode('utf-8')))
+    return caches_in_use
+
+def main():
+    # type: () -> None
+    args = parse_args()
+
+    all_caches = set([os.path.join(VENV_CACHE_DIR, cache) for cache in os.listdir(VENV_CACHE_DIR)])
+    caches_in_use = get_caches_in_use(args.threshold_days)
+    caches_to_purge = get_caches_to_be_purged(VENV_CACHE_DIR, caches_in_use, args.threshold_days)
+    caches_to_keep = all_caches - caches_to_purge
+
+    if args.dry_run:
+        print("Performing a dry run...")
+    else:
+        print("Cleaning unused venv caches...")
+
+    for cache_dir in caches_to_purge:
+        print("Cleaning unused venv %s" % (cache_dir,))
+        if not args.dry_run:
+            subprocess.check_call(["sudo", "rm", "-rf", cache_dir])
+
+    for cache_dir in caches_to_keep:
+        print("Keeping used venv: %s" % (cache_dir,))
+
+    print("Done!\n")
+
+if __name__ == "__main__":
+    main()
diff --git a/tools/clean-venv-cache b/tools/clean-venv-cache
deleted file mode 100755
index 8cecda3257..0000000000
--- a/tools/clean-venv-cache
+++ /dev/null
@@ -1,28 +0,0 @@
-#!/usr/bin/env python3
-from __future__ import absolute_import
-from __future__ import print_function
-import glob
-import os
-import subprocess
-import sys
-
-ZULIP_PATH = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
-VENV_CACHE_DIR = "/srv/zulip-venv-cache"
-if "--travis" in sys.argv:
-    VENV_CACHE_DIR = os.path.join(os.environ["HOME"], "zulip-venv-cache")
-
-used_caches = set()
-hash_reqs = os.path.join(ZULIP_PATH, 'scripts', 'lib', 'hash_reqs.py')
-
-for filename in os.listdir(os.path.join(ZULIP_PATH, "requirements")):
-    requirements_file = os.path.join(ZULIP_PATH, "requirements", filename)
-    hash_val = subprocess.check_output([hash_reqs, requirements_file]).strip()
-    used_caches.add(os.path.join(VENV_CACHE_DIR, hash_val.decode('utf-8')))
-
-for cache_dir_base in os.listdir(VENV_CACHE_DIR):
-    cache_dir = os.path.join(VENV_CACHE_DIR, cache_dir_base)
-    if cache_dir not in used_caches:
-        print("Cleaning unused venv %s" % (cache_dir,))
-        subprocess.check_call(["sudo", "rm", "-rf", cache_dir])
-    else:
-        print("Keeping used venv %s" % (cache_dir,))