2 files changed, 44 insertions(+), 18 deletions(-)

A cephfs-layout-tool/__init__.py
R migrate_pools.py => cephfs-layout-tool/migrate_pools.py

@@ 3,32 3,48 @@ import os
 import shutil
 import sys
 import tempfile
-from collections import namedtuple
-from typing import Optional, Callable
+import functools
+
+# from collections import namedtuple
+from typing import Optional, NamedTuple
 
 import xattr  # type: ignore
 import humanize  # type: ignore
 
-CephLayout = namedtuple("CephLayout", ["stripe_count", "object_size", "pool"])
-
 TMPDIR = "/c/tmp"
 
 OK_POOLS = {"cephfs_crs52data", "cephfs_crs52data2"}
 
 
-def memoize(fn: Callable):
-    """ Memoization decorator for a function taking a single argument """
+class CephLayout(
+    NamedTuple("CephLayout", [("stripe_count", int), ("object_size", int), ("pool", str)])
+):
+    def __eq__(self, other):
+        return (
+            self.stripe_count == other.stripe_count
+            and self.object_size == other.object_size
+            and self.pool == other.pool
+        )
+
+
+def memoize(obj):
+    """Decorator to memoize a function."""
+    cache = obj.cache = {}
 
-    class MemoDict(dict):
-        def __missing__(self, key):
-            ret = self[key] = fn(key)
-            return ret
+    @functools.wraps(obj)
+    def memoizer(*args, **kwargs):
+        key = str(args) + str(kwargs)
+        if key not in cache:
+            cache[key] = obj(*args, **kwargs)
+        return cache[key]
 
-    return MemoDict().__getitem__
+    return memoizer
 
 
 @memoize
 def extract_layout(filename: str) -> Optional[CephLayout]:
+    """Figure out what the file layout for a given directory should be, looking at parent
+       directories if necessary."""
     filetype = "file"
     if os.path.isdir(filename):
         filetype = "dir"


@@ 46,8 62,8 @@ def extract_layout(filename: str) -> Optional[CephLayout]:
             return extract_layout(os.path.dirname(filename))
         return None
     for attr in xattrs:
-        n = attr.split("=")
-        cephlayout[n[0]] = n[1]
+        attr_tuple = attr.split("=")
+        cephlayout[attr_tuple[0]] = attr_tuple[1]
     del cephlayout["stripe_unit"]
     return CephLayout(**cephlayout)
 


@@ 55,6 71,7 @@ def extract_layout(filename: str) -> Optional[CephLayout]:
 # make a temp dir with the same layout as the given dir
 @memoize
 def mkdtemp_layout(layout: CephLayout, prefix: str = TMPDIR) -> str:
+    """Create temporary directory with the given layout"""
     tempdir = tempfile.mkdtemp(dir=prefix)
     xattrs = xattr.xattr(tempdir)
     for attr in layout._fields:


@@ 63,6 80,7 @@ def mkdtemp_layout(layout: CephLayout, prefix: str = TMPDIR) -> str:
 
 
 def main():
+    """entrypoint of script"""
     startdir = sys.argv[1]
 
     total_savings = 0


@@ 86,14 104,15 @@ def main():
             file_layout = extract_layout(filename)
             if not file_layout:
                 continue
+            if dir_layout != file_layout:
+                print("file layout doesn't match dir layout: {}".format(file_layout))
+                tmploc = os.path.join(tmp_layout_dir, name)
+                relayout_file(filename, tmploc)
             if dir_layout.pool != file_layout.pool:
                 print("%s in wrong pool: %s" % (name, file_layout.pool))
                 statinfo = os.stat(filename)
                 tmploc = os.path.join(tmp_layout_dir, name)
-                print("copying {} to temp location {}".format(filename, tmploc))
-                shutil.copy2(filename, tmploc)
-                print("moving back on top of original")
-                shutil.move(tmploc, filename)
+                relayout_file(filename, tmploc)
                 oldusage = (statinfo.st_size / 4) * 6
                 newusage = (statinfo.st_size / 5) * 7
                 savings = oldusage - newusage


@@ 102,7 121,14 @@ def main():
                 print("saved {}".format(humanize.naturalsize(savings)))
 
     print("saved space in total: {}".format(humanize.naturalsize(total_savings)))
-    os.rmdir(session_tmpdir)
+    shutil.rmtree(session_tmpdir)
+
+
+def relayout_file(filename, tmploc):
+    print("copying {} to temp location {}".format(filename, tmploc))
+    shutil.copy2(filename, tmploc)
+    print("moving back on top of original")
+    shutil.move(tmploc, filename)
 
 
 if __name__ == "__main__":