fix: adjust incremental loader to find the largest chain (#411)

* fix: adjust incremental loader to find the largest chain

* Update oci_python_image/hello_world/integration_test.py

Co-authored-by: aspect-workflows[bot] <143031405+aspect-workflows[bot]@users.noreply.github.com>

---------

Co-authored-by: Alex Eagle <[email protected]>
Co-authored-by: aspect-workflows[bot] <143031405+aspect-workflows[bot]@users.noreply.github.com>

Author: 3 people

oci_python_image/hello_world/integration_test.py

@@ -3,37 +3,52 @@
 # TODO: upstream some of this to the testcontainers library to simplify similar code for users.
 from testcontainers.core.container import DockerContainer
 from testcontainers.core.waiting_utils import wait_for_logs
-import docker
-import tarfile
-import json
-import os
-import tempfile
-import io
-
-
-def add_json_file(tar, name, contents):
-    content = json.dumps(contents).encode("utf-8")
-    info = tarfile.TarInfo(name=name)
-    info.size = len(content)
-    tar.addfile(info, fileobj=io.BytesIO(content))
-
-
-def add_file(tar, name, fileobj):
-    info = tarfile.TarInfo(name=name)
-    info.size = os.fstat(fileobj.fileno()).st_size
-    tar.addfile(info, fileobj=fileobj)
-    fileobj.close()
-
-
-def get_blob_path(image, digest):
-    return "%s/blobs/%s" % (image, digest.replace(":", "/"))
+import logging
 
+# This function is self container for easy-to-use by just copy-pasting. 
+def OCIImageContainer(image):
+    import docker
+    import tarfile
+    import json
+    import io
+    import logging
+
+    logger = logging.getLogger("incremental_loader")
+    def tar(*args):
+        buffer = io.BytesIO()
+        with tarfile.open(fileobj=buffer, mode="w:") as t:
+            for name, size, contents in args:
+                info = tarfile.TarInfo(name=name)
+                info.size = size
+                if isinstance(contents, list) or isinstance(contents, dict):
+                    content = json.dumps(contents).encode("utf-8")
+                    info.size = len(content)
+                    contents = io.BytesIO(content)
+    
+                t.addfile(info, fileobj=contents)
+                contents.close()
+        return buffer
+
+    def config_json(diff_ids):
+        return ("config.json", None, {
+            "rootfs": {
+                "type": "layers",
+                "diff_ids": diff_ids
+            }
+        })
+    
+    def manifest_json(layers):
+        return ("manifest.json", None, [{
+            "Config": "config.json",
+            "RepoTags": [],
+            "Layers": layers
+        }])
 
-def open_blob(image, digest):
-    return open(get_blob_path(image, digest), "rb")
 
+    def open_blob(image, digest):
+        blob_path = "%s/blobs/%s" % (image, digest.replace(":", "/"))
+        return open(blob_path, "rb")
 
-def OCIImageContainer(image):
     with open("%s/index.json" % image) as indexp:
         indexjson = json.load(indexp)
 
@@ -45,72 +60,64 @@ def OCIImageContainer(image):
 
     client = docker.from_env()
 
-    # Probe and layer loading phase
-    layers = manifest["layers"]
-    needed = []
-
-    # Probing phase
-    for i, layer in enumerate(layers):
-        tmp = tempfile.NamedTemporaryFile(suffix=".tar")
-        tar = tarfile.open(fileobj=tmp, mode="w")
-        add_json_file(
-            tar,
-            name="manifest.json",
-            contents=[
-                {
-                    "Config": "config.json",
-                    "RepoTags": [],
-                    "Layers": [layer["digest"]],
-                }
-            ],
-        )
-        add_json_file(
-            tar,
-            name="config.json",
-            contents={
-                "rootfs": {
-                    "type": "layers",
-                    "diff_ids": [config["rootfs"]["diff_ids"][i]],
-                }
-            },
+    # Happy path: check if the image exists in the cache
+    try:
+        probe = tar(
+            manifest_json(layers=[layer["digest"] + ".tar" for layer in manifest["layers"]]),
+            ("config.json", manifest["config"]["size"], open_blob(image, manifest["config"]["digest"]))
         )
+        client.images.load(probe.getvalue())
+        logger.debug("Image was a cache hit")
+        return DockerContainer(manifest["config"]["digest"])
+    except docker.errors.ImageLoadError as e:
+        logger.debug("Image was a cache miss because %s" % e)
+
+    
+    # Unhappy path: image is not in the cache, so we need to load it incrementally
+    layers = manifest["layers"]
+    start_from = 0
 
-        tar.close()
+    # Probe the cache for the first layer that is not in the cache
+    for i in range(1, len(layers)):
+        previous_layer = layers[i-1]
+        layer = layers[i]
 
+        probe = tar(
+            manifest_json(layers=[layer["digest"] + ".tar" for layer in manifest["layers"][:i]]),
+            config_json(diff_ids=config["rootfs"]["diff_ids"][:i])
+        )
+    
         try:
-            # os.system("tar -tvf %s" % tmp.name)
-            client.images.load(
-                open(tmp.name, "rb"),
-            )
-        except docker.errors.ImageLoadError:
-            needed.append(layer["digest"])
-
-    # Loading phase
-    tmp = tempfile.NamedTemporaryFile(suffix=".tar")
-    tar = tarfile.open(fileobj=tmp, mode="w")
-    add_json_file(
-        tar,
-        name="manifest.json",
-        contents=[
-            {
-                "Config": "config.json",
-                "RepoTags": [],
-                "Layers": list(map(lambda x: x["digest"], manifest["layers"])),
-            }
-        ],
-    )
-    add_file(
-        tar, name="config.json", fileobj=open_blob(image, manifest["config"]["digest"])
+            probe_res = client.images.load(probe.getvalue())
+            logger.debug("Layer %s was a cache hit" % layer["digest"], probe_res[0].id)
+        except docker.errors.ImageLoadError as e:
+            logger.debug("Layer %s was a cache miss because %s" % (layer["digest"], e))
+            # From this point on we need to load everything.
+            start_from = i-1
+            break
+
+    logger.debug("Need to load all the layers after: ", start_from)
+
+    # Load all layers at once
+    load = tar(
+        # Send only needed layers
+        *[(layer["digest"] + ".tar", layer["size"], open_blob(image, layer["digest"])) for layer in layers[start_from:]],
+        manifest_json(layers=[layer["digest"] + ".tar" for layer in manifest["layers"]]),
+        ("config.json", manifest["config"]["size"], open_blob(image, manifest["config"]["digest"]))
     )
-    for layer in needed:
-        add_file(tar, name=layer, fileobj=open_blob(image, layer))
 
-    tar.close()
-    r = client.images.load(open(tmp.name, "rb"))
-    return DockerContainer(r[0].id)
+    load_res = client.images.load(load.getvalue())
+
+    # Extra safety check: assert that the final image id is equal to the one in the manifest
+    assert load_res[0].id == manifest["config"]["digest"], "final image id should match manifest %s == %s" % (load_res[0].id, manifest["config"]["digest"])
+
+    return DockerContainer(manifest["config"]["digest"])
+
+
 
 
 def test_wait_for_hello():
+    logging.basicConfig(level=logging.DEBUG)
     print("Starting container")
     with OCIImageContainer("oci_python_image/hello_world/image") as container:
         wait_for_logs(container, "hello py_image_layer!")