From 700d13876ad2f1542516676875a2286b6d362ff5 Mon Sep 17 00:00:00 2001
From: jpic <jpic@yourlabs.org>
Date: Sat, 30 May 2020 18:34:19 +0200
Subject: [PATCH] Add layer caching

---
 shlax/actions/packages.py |  6 ++++
 shlax/targets/base.py     | 50 +++++++++++++++--------------
 shlax/targets/buildah.py  | 67 ++++++++++++++++++++++++++++++++-------
 3 files changed, 89 insertions(+), 34 deletions(-)

diff --git a/shlax/actions/packages.py b/shlax/actions/packages.py
index cbdb385..523bd1a 100644
--- a/shlax/actions/packages.py
+++ b/shlax/actions/packages.py
@@ -163,3 +163,9 @@ class Packages:
 
     def __repr__(self):
         return f'Packages({self.packages}, upgrade={self.upgrade})'
+
+    def layerhasher(self, parent=None):
+        import hashlib
+        blurb = (parent or '') + repr(self)
+        sha1 = hashlib.sha1(blurb.encode('ascii'))
+        return sha1.hexdigest()
diff --git a/shlax/targets/base.py b/shlax/targets/base.py
index 1ca6402..7c50a59 100644
--- a/shlax/targets/base.py
+++ b/shlax/targets/base.py
@@ -30,30 +30,34 @@ class Target:
             self.parent = target
 
         for action in actions or self.actions:
-            result = Result(self, action)
-            self.output.start(action)
-            try:
-                await action(target=self)
-            except Exception as e:
-                self.output.fail(action, e)
-                result.status = 'failure'
-                result.exception = e
-                if actions:
-                    # nested call, re-raise
-                    raise
-                else:
-                    break
-            else:
-                self.output.success(action)
-                result.status = 'success'
-            finally:
-                self.caller.results.append(result)
+            if await self.action(action, reraise=bool(actions)):
+                break
 
-                clean = getattr(action, 'clean', None)
-                if clean:
-                    action.result = result
-                    self.output.clean(action)
-                    await clean(self)
+    async def action(self, action, reraise=False):
+        result = Result(self, action)
+        self.output.start(action)
+        try:
+            await action(target=self)
+        except Exception as e:
+            self.output.fail(action, e)
+            result.status = 'failure'
+            result.exception = e
+            if reraise:
+                # nested call, re-raise
+                raise
+            else:
+                return True
+        else:
+            self.output.success(action)
+            result.status = 'success'
+        finally:
+            self.caller.results.append(result)
+
+            clean = getattr(action, 'clean', None)
+            if clean:
+                action.result = result
+                self.output.clean(action)
+                await clean(self)
 
     async def rexec(self, *args, **kwargs):
         kwargs['user'] = 'root'
diff --git a/shlax/targets/buildah.py b/shlax/targets/buildah.py
index 179fd69..cfd9a57 100644
--- a/shlax/targets/buildah.py
+++ b/shlax/targets/buildah.py
@@ -1,3 +1,5 @@
+import copy
+import json
 import os
 import sys
 from pathlib import Path
@@ -45,9 +47,50 @@ class Buildah(Target):
             os.execvp('buildah', ['buildah', 'unshare'] + sys.argv)
             # program has been replaced
 
+        actions_done = await self.cache_load(*actions)
+
+        if actions:
+            actions = actions[len(actions_done):]
+        else:
+            self.actions = self.actions[len(actions_done):]
+
         self.ctr = (await self.parent.exec('buildah', 'from', self.base)).out
         self.mnt = Path((await self.parent.exec('buildah', 'mount', self.ctr)).out)
-        await super().__call__()
+        await super().__call__(*actions)
+
+    async def cache_load(self, *actions):
+        actions_done = []
+        result = await self.parent.exec(
+            'podman image list --format json',
+            quiet=True,
+        )
+        result = json.loads(result.out)
+        images = [item for sublist in result for item in sublist['History']]
+        self.hash_previous = None
+        for action in actions or self.actions:
+            hasher = getattr(action, 'layerhasher', None)
+            if not hasher:
+                break
+            layerhash = hasher(self.hash_previous)
+            layerimage = copy.deepcopy(self.image)
+            layerimage.tags = [layerhash]
+            if 'localhost/' + str(layerimage) in images:
+                self.base = str(layerimage)
+                self.hash_previous = layerhash
+                actions_done.append(action)
+                self.output.success(f'Found cached layer for {action}')
+
+        return actions_done
+
+    async def action(self, action, reraise=False):
+        result = await super().action(action, reraise)
+        hasher = getattr(action, 'layerhasher', None)
+        if hasher:
+            layerhash = hasher(self.hash_previous if self.hash_previous else None)
+            layerimage = copy.deepcopy(self.image)
+            layerimage.tags = [layerhash]
+            await self.commit(layerimage)
+        return result
 
     async def clean(self, target):
         for src, dst in self.mounts.items():
@@ -79,25 +122,27 @@ class Buildah(Target):
         _args += [' '.join([str(a) for a in args])]
         return await self.parent.exec(*_args, **kwargs)
 
-    async def commit(self):
-        if not self.image:
+    async def commit(self, image=None):
+        image = image or self.image
+        if not image:
             return
 
-        for key, value in self.config.items():
-            await self.parent.exec(f'buildah config --{key} "{value}" {self.ctr}')
+        if not image:
+            # don't go through that if layer commit
+            for key, value in self.config.items():
+                await self.parent.exec(f'buildah config --{key} "{value}" {self.ctr}')
 
-        self.sha = (await self.exec(
+        self.sha = (await self.parent.exec(
             'buildah',
             'commit',
-            '--format=' + self.image.format,
+            '--format=' + image.format,
             self.ctr,
-            buildah=False,
         )).out
 
-        if self.image.tags:
-            tags = [f'{self.image.repository}:{tag}' for tag in self.image.tags]
+        if image.tags:
+            tags = [f'{image.repository}:{tag}' for tag in image.tags]
         else:
-            tags = [self.image.repository]
+            tags = [image.repository]
 
         for tag in tags:
             await self.parent.exec('buildah', 'tag', self.sha, tag)