From c9ce89d2afcf0b6f4efa6d0cf773ac11041c84ef Mon Sep 17 00:00:00 2001
From: Tristan Van Berkom <tristan.vanberkom@codethink.co.uk>
Date: Fri, 18 Jan 2019 12:17:19 -0500
Subject: _cas/cascache.py: Cleanup directories when removing refs

With out this, empty directories in the refs/heads directory just
grow unconditionally.
---
 buildstream/_cas/cascache.py | 56 ++++++++++++++++++++++++++++++++++++++++----
 1 file changed, 51 insertions(+), 5 deletions(-)

diff --git a/buildstream/_cas/cascache.py b/buildstream/_cas/cascache.py
index adbd34c9e..1092f6f75 100644
--- a/buildstream/_cas/cascache.py
+++ b/buildstream/_cas/cascache.py
@@ -21,6 +21,7 @@ import hashlib
 import itertools
 import os
 import stat
+import errno
 import tempfile
 import uuid
 import contextlib
@@ -545,11 +546,7 @@ class CASCache():
     def remove(self, ref, *, defer_prune=False):
 
         # Remove cache ref
-        refpath = self._refpath(ref)
-        if not os.path.exists(refpath):
-            raise CASCacheError("Could not find ref '{}'".format(ref))
-
-        os.unlink(refpath)
+        self._remove_ref(ref)
 
         if not defer_prune:
             pruned = self.prune()
@@ -626,6 +623,55 @@ class CASCache():
     def _refpath(self, ref):
         return os.path.join(self.casdir, 'refs', 'heads', ref)
 
+    # _remove_ref()
+    #
+    # Removes a ref.
+    #
+    # This also takes care of pruning away directories which can
+    # be removed after having removed the given ref.
+    #
+    # Args:
+    #    ref (str): The ref to remove
+    #
+    # Raises:
+    #    (CASCacheError): If the ref didnt exist, or a system error
+    #                     occurred while removing it
+    #
+    def _remove_ref(self, ref):
+
+        # Remove the ref itself
+        refpath = self._refpath(ref)
+        try:
+            os.unlink(refpath)
+        except FileNotFoundError as e:
+            raise CASCacheError("Could not find ref '{}'".format(ref)) from e
+
+        # Now remove any leading directories
+        basedir = os.path.join(self.casdir, 'refs', 'heads')
+        components = list(os.path.split(ref))
+        while components:
+            components.pop()
+            refdir = os.path.join(basedir, *components)
+
+            # Break out once we reach the base
+            if refdir == basedir:
+                break
+
+            try:
+                os.rmdir(refdir)
+            except FileNotFoundError:
+                # The parent directory did not exist, but it's
+                # parent directory might still be ready to prune
+                pass
+            except OSError as e:
+                if e.errno == errno.ENOTEMPTY:
+                    # The parent directory was not empty, so we
+                    # cannot prune directories beyond this point
+                    break
+
+                # Something went wrong here
+                raise CASCacheError("System error while removing ref '{}': {}".format(ref, e)) from e
+
     # _commit_directory():
     #
     # Adds local directory to content addressable store.
-- 
cgit v1.2.1