[PATCH 2/2] Add "btrsquash" style ramdisk

[Will Woods <wwoods redhat com>]

This adds a "btrsquash" style for ramdisks - a squashfs-compressed btrfs
image which gets loaded by the same dracut 'dmsquash-live' module used by
our LiveOS images.

This saves a good ~256MB RAM, which allows us to install on systems with
512MB RAM again.
---
 src/pylorax/__init__.py    |    5 ++
 src/pylorax/constants.py   |    2 +
 src/pylorax/installtree.py |  102 ++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 109 insertions(+), 0 deletions(-)

diff --git a/src/pylorax/__init__.py b/src/pylorax/__init__.py
index de74b28..a575031 100644
--- a/src/pylorax/__init__.py
+++ b/src/pylorax/__init__.py
@@ -283,6 +283,11 @@ class Lorax(BaseLoraxClass):
         logger.info("moving stubs")
         self.installtree.move_stubs()
 
+        # if needed: create initramfs (before modules get shuffled around)
+        if self.conf.get("ramdisk", "style") == "btrsquash":
+            logger.info("creating dracut initramfs")
+            self.installtree.make_squash_initramfs()
+
         # get the list of required modules
         logger.info("getting list of required modules")
         modules = [f[1:] for f in template if f[0] == "module"]
diff --git a/src/pylorax/constants.py b/src/pylorax/constants.py
index 547a938..c46699f 100644
--- a/src/pylorax/constants.py
+++ b/src/pylorax/constants.py
@@ -47,7 +47,9 @@ class LoraxRequiredCommands(dict):
         self["LOCALEDEF"] = "localedef"
         self["LOSETUP"] = "losetup"
         self["MKDOSFS"] = "mkdosfs"
+        self["MKFS_BTRFS"] = "mkfs.btrfs"
         self["MKISOFS"] = "mkisofs"
+        self["MKSQUASHFS"] = "mksquashfs"
         self["MODINFO"] = "modinfo"
         self["MOUNT"] = "mount"
         self["PARTED"] = "parted"
diff --git a/src/pylorax/installtree.py b/src/pylorax/installtree.py
index 9704769..60d6470 100644
--- a/src/pylorax/installtree.py
+++ b/src/pylorax/installtree.py
@@ -26,6 +26,7 @@ import sys
 import os
 import shutil
 import gzip
+import lzma
 import re
 import glob
 import time
@@ -52,6 +53,9 @@ class LoraxInstallTree(BaseLoraxClass):
 
         if self.style == 'initramfs':
             self.make_initrd = self.make_initramfs
+        if self.style == 'btrsquash':
+            self._mkfs = [self.lcmds.MKFS_BTRFS, "-L", "Anaconda"]
+            self.make_initrd = self.make_live_squashfs
 
     def compress(self, initrd, kernel, compression="xz"):
         start = time.time()
@@ -539,6 +543,104 @@ class LoraxInstallTree(BaseLoraxClass):
 
         return True
 
+    def make_squash_initramfs(self):
+        outfile = "/tmp/initramfs.img"
+        logger.debug("chrooting into installtree to create initramfs.img")
+        # NOTE: ext[34] gets included by default, but we have to ask for btrfs
+        subprocess.check_call(["chroot", self.root,
+                               "/sbin/dracut", "--nomdadmconf", "--nolvmconf",
+                               "--modules", "base btrfs dmsquash-live",
+                               outfile, self.kernels[0].version])
+        # move output file into installtree workdir, and repack with xz.
+        # NOTE: for some reason the concatenated image will fail to boot if we
+        # leave it gzipped? Weird.
+        self.initramfs = joinpaths(self.workdir, "initramfs.img")
+        gzip_in = gzip.open(joinpaths(self.root, outfile))
+        xz_out = lzma.LZMAFile(self.initramfs, "w",
+                               options={'format':'xz', 'level':9})
+        xz_out.write(gzip_in.read())
+
+
+    def make_live_squashfs(self, initrd, kernel, compression):
+        '''This is a little complicated, but dracut wants to find a squashfs
+        image named "squashfs.img" which contains a filesystem image named
+        "LiveOS/rootfs.img".
+        Placing squashfs.img inside a cpio image and concatenating that
+        with the existing initramfs.img will make squashfs.img appear inside
+        initramfs at boot time.'''
+        # These exact names are required by dracut
+        squashname = "squashfs.img"
+        imgname = "LiveOS/rootfs.img"
+
+        # Create fs image of installtree
+        fsimage = joinpaths(self.workdir, "installtree.img")
+        open(fsimage, "wb").truncate(2*1024**3)
+        mountpoint = joinpaths(self.workdir, "rootfs")
+        os.mkdir(mountpoint, 0755)
+        logger.debug("formatting rootfs image: %s %s",
+                     " ".join(self._mkfs), fsimage)
+        subprocess.check_call(self._mkfs + [fsimage], stdout=subprocess.PIPE)
+        logger.debug("mounting rootfs image at %s", mountpoint)
+        subprocess.check_call([self.lcmds.MOUNT, "-o", "loop",
+                               fsimage, mountpoint])
+        try:
+            logger.info("copying installtree into rootfs image")
+            srcfiles = [joinpaths(self.root, f) for f in os.listdir(self.root)]
+            subprocess.check_call(["cp", "-a"] + srcfiles + [mountpoint])
+        finally:
+            logger.debug("unmounting rootfs image")
+            rc = subprocess.call([self.lcmds.UMOUNT, mountpoint])
+        if rc != 0:
+            logger.critical("umount %s failed (returncode %i)", mountpoint, rc)
+            sys.exit(rc)
+        os.rmdir(mountpoint)
+
+        # Make squashfs with rootfs image inside
+        logger.info("creating %s containing %s", squashname, imgname)
+        squashtree = joinpaths(self.workdir, "squashfs")
+        os.makedirs(joinpaths(squashtree, os.path.dirname(imgname)))
+        shutil.move(fsimage, joinpaths(squashtree, imgname))
+        squashimage = joinpaths(self.workdir, squashname)
+        subprocess.check_call([self.lcmds.MKSQUASHFS, squashtree, squashimage,
+                               "-comp", compression])
+        shutil.rmtree(squashtree)
+
+        # Put squashimage in a new cpio image with dracut config
+        logger.debug("creating cpio image containing %s", squashname)
+        initramfsdir = joinpaths(self.workdir, "initramfs")
+        squash_cpio = joinpaths(self.workdir, "squashfs.cpio")
+        cmdline = joinpaths(initramfsdir, "etc/cmdline")
+        os.makedirs(os.path.dirname(cmdline))
+        # write boot cmdline for dracut
+        with open(cmdline, "wb") as fobj:
+            fobj.write("root=live:/{0}\n".format(squashname))
+            if self.style == "btrsquash":
+                fobj.write("rootflags=compress\n")
+        # add squashimage
+        shutil.move(squashimage, initramfsdir)
+        # create cpio container
+        chdir = lambda: os.chdir(initramfsdir)
+        find = subprocess.Popen([self.lcmds.FIND, "."], stdout=subprocess.PIPE,
+                                preexec_fn=chdir)
+        cpio = subprocess.Popen([self.lcmds.CPIO, "--quiet", "-c", "-o"],
+                                stdin=find.stdout,
+                                stdout=open(squash_cpio, "wb"),
+                                preexec_fn=chdir)
+        cpio.communicate()
+        shutil.rmtree(initramfsdir)
+
+        # create final image
+        logger.debug("concatenating initramfs.img and squashfs cpio")
+        logger.debug("initramfs.img size = %i", os.stat(self.initramfs).st_size)
+        with open(initrd.fpath, "wb") as output:
+            with open(self.initramfs, "rb") as fobj:
+                output.write(fobj.read())
+            with open(squash_cpio, "rb") as fobj:
+                output.write(fobj.read())
+        os.remove(self.initramfs)
+        os.remove(squash_cpio)
+
+        return True
 
     @property
     def kernels(self):
-- 
1.7.4