Cache compressed local base image layers

jib-core/src/main/java/com/google/cloud/tools/jib/builder/steps/ExtractTarStep.java

@@ -18,12 +18,15 @@
 
 import com.fasterxml.jackson.databind.MapperFeature;
 import com.fasterxml.jackson.databind.ObjectMapper;
+import com.google.cloud.tools.jib.api.DescriptorDigest;
 import com.google.cloud.tools.jib.blob.Blob;
 import com.google.cloud.tools.jib.blob.BlobDescriptor;
 import com.google.cloud.tools.jib.blob.Blobs;
 import com.google.cloud.tools.jib.builder.ProgressEventDispatcher;
 import com.google.cloud.tools.jib.builder.TimerEventDispatcher;
 import com.google.cloud.tools.jib.builder.steps.ExtractTarStep.LocalImage;
+import com.google.cloud.tools.jib.cache.Cache;
+import com.google.cloud.tools.jib.cache.CacheCorruptedException;
 import com.google.cloud.tools.jib.cache.CachedLayer;
 import com.google.cloud.tools.jib.configuration.BuildConfiguration;
 import com.google.cloud.tools.jib.docker.json.DockerManifestEntryTemplate;
@@ -46,6 +49,7 @@
 import java.nio.file.Path;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Optional;
 import java.util.concurrent.Callable;
 import java.util.zip.GZIPInputStream;
 import java.util.zip.GZIPOutputStream;
@@ -96,7 +100,8 @@ static boolean isGzipped(Path path) throws IOException {
 
   @Override
   public LocalImage call()
-      throws IOException, LayerCountMismatchException, BadContainerConfigurationFormatException {
+      throws IOException, LayerCountMismatchException, BadContainerConfigurationFormatException,
+          CacheCorruptedException {
     Path destination = Files.createTempDirectory("jib-extract-tar");
     try (TimerEventDispatcher ignored =
         new TimerEventDispatcher(
@@ -132,56 +137,23 @@ public LocalImage call()
 
       // Process layer blobs
       // TODO: Optimize; compressing/calculating layer digests is slow
-      //       e.g. parallelize, cache layers, faster compression method
+      //       e.g. parallelize, faster compression method
       try (ProgressEventDispatcher progressEventDispatcher =
           progressEventDispatcherFactory.create(
               "processing base image layers", layerFiles.size())) {
         List<PreparedLayer> layers = new ArrayList<>(layerFiles.size());
         V22ManifestTemplate v22Manifest = new V22ManifestTemplate();
 
-        List<ProgressEventDispatcher.Factory> childProgressFactories = new ArrayList<>();
-        for (String ignored1 : layerFiles) {
-          childProgressFactories.add(progressEventDispatcher.newChildProducer());
-        }
-
         for (int index = 0; index < layerFiles.size(); index++) {
-          Path file = destination.resolve(layerFiles.get(index));
-
-          // Compress layers if necessary and calculate the digest/size
-          Blob blob = Blobs.from(file);
-          try (ProgressEventDispatcher childDispatcher =
-                  childProgressFactories
-                      .get(index)
-                      .create("compressing " + file, Files.size(file));
-              ThrottledAccumulatingConsumer throttledProgressReporter =
-                  new ThrottledAccumulatingConsumer(childDispatcher::dispatchProgress)) {
-            if (!layersAreCompressed) {
-              Path compressedFile = destination.resolve(layerFiles.get(index) + ".compressed");
-              try (GZIPOutputStream compressorStream =
-                      new GZIPOutputStream(Files.newOutputStream(compressedFile));
-                  NotifyingOutputStream notifyingOutputStream =
-                      new NotifyingOutputStream(compressorStream, throttledProgressReporter)) {
-                blob.writeTo(notifyingOutputStream);
-              }
-              blob = Blobs.from(compressedFile);
-            }
-          }
-          BlobDescriptor blobDescriptor = blob.writeTo(ByteStreams.nullOutputStream());
-
-          // 'manifest' contains the layer files in the same order as the diff ids in
-          // 'configuration', so we don't need to recalculate those.
-          // https://containers.gitbook.io/build-containers-the-hard-way/#docker-load-format
+          Path layerFile = destination.resolve(layerFiles.get(index));
           CachedLayer layer =
-              CachedLayer.builder()
-                  .setLayerBlob(blob)
-                  .setLayerDigest(blobDescriptor.getDigest())
-                  .setLayerSize(blobDescriptor.getSize())
-                  .setLayerDiffId(configurationTemplate.getLayerDiffId(index))
-                  .build();
-
+              getCachedTarLayer(
+                  configurationTemplate.getLayerDiffId(index),
+                  layerFile,
+                  layersAreCompressed,
+                  progressEventDispatcher.newChildProducer());
           layers.add(new PreparedLayer.Builder(layer).build());
-          v22Manifest.addLayer(blobDescriptor.getSize(), blobDescriptor.getDigest());
-          progressEventDispatcher.dispatchProgress(1);
+          v22Manifest.addLayer(layer.getSize(), layer.getDigest());
         }
 
         BlobDescriptor configDescriptor =
@@ -193,4 +165,42 @@ public LocalImage call()
       }
     }
   }
+
+  private CachedLayer getCachedTarLayer(
+      DescriptorDigest diffId,
+      Path layerFile,
+      boolean layersAreCompressed,
+      ProgressEventDispatcher.Factory progressEventDispatcherFactory)
+      throws IOException, CacheCorruptedException {
+    try (ProgressEventDispatcher childDispatcher =
+            progressEventDispatcherFactory.create(
+                "compressing layer " + diffId, Files.size(layerFile));
+        ThrottledAccumulatingConsumer throttledProgressReporter =
+            new ThrottledAccumulatingConsumer(childDispatcher::dispatchProgress)) {
+      Cache cache = buildConfiguration.getBaseImageLayersCache();
+
+      // Retrieve pre-compressed layer from cache
+      Optional<CachedLayer> optionalLayer = cache.retrieveTarLayer(diffId);
+      if (optionalLayer.isPresent()) {
+        return optionalLayer.get();
+      }
+
+      // Just write layers that are already compressed
+      if (layersAreCompressed) {
+        return cache.writeTarLayer(diffId, Blobs.from(layerFile));
+      }
+
+      // Compress uncompressed layers while writing
+      Blob compressedBlob =
+          Blobs.from(
+              outputStream -> {
+                try (GZIPOutputStream compressorStream = new GZIPOutputStream(outputStream);
+                    NotifyingOutputStream notifyingOutputStream =
+                        new NotifyingOutputStream(compressorStream, throttledProgressReporter)) {
+                  Blobs.from(layerFile).writeTo(notifyingOutputStream);
+                }
+              });
+      return cache.writeTarLayer(diffId, compressedBlob);
+    }
+  }
 }

jib-core/src/main/java/com/google/cloud/tools/jib/cache/Cache.java

@@ -116,6 +116,20 @@ public CachedLayer writeUncompressedLayer(
         uncompressedLayerBlob, LayerEntriesSelector.generateSelector(layerEntries));
   }
 
+  /**
+   * Caches a layer that was extracted from a local base image, and names the file using the
+   * provided diff id.
+   *
+   * @param diffId the diff id
+   * @param compressedBlob the compressed layer blob
+   * @return the {@link CachedLayer} for the written layer
+   * @throws IOException if an I/O exception occurs
+   */
+  public CachedLayer writeTarLayer(DescriptorDigest diffId, Blob compressedBlob)
+      throws IOException {
+    return cacheStorageWriter.writeTarLayer(diffId, compressedBlob);
+  }
+
   /**
    * Retrieves the cached manifest and container configuration for an image reference.
    *
@@ -160,4 +174,17 @@ public Optional<CachedLayer> retrieve(DescriptorDigest layerDigest)
       throws IOException, CacheCorruptedException {
     return cacheStorageReader.retrieve(layerDigest);
   }
+
+  /**
+   * Retrieves a {@link CachedLayer} for a local base image layer with the given diff id.
+   *
+   * @param diffId the diff id
+   * @return the {@link CachedLayer} with the given diff id
+   * @throws CacheCorruptedException if the cache was found to be corrupted
+   * @throws IOException if an I/O exception occurs
+   */
+  public Optional<CachedLayer> retrieveTarLayer(DescriptorDigest diffId)
+      throws IOException, CacheCorruptedException {
+    return cacheStorageReader.retrieveTarLayer(diffId);
+  }
 }

jib-core/src/main/java/com/google/cloud/tools/jib/cache/CacheStorageFiles.java

@@ -26,6 +26,7 @@
 class CacheStorageFiles {
 
   private static final String LAYERS_DIRECTORY = "layers";
+  private static final String LOCAL_DIRECTORY = "local";
   private static final String IMAGES_DIRECTORY = "images";
   private static final String SELECTORS_DIRECTORY = "selectors";
   private static final String TEMPORARY_DIRECTORY = "tmp";
@@ -54,14 +55,14 @@ static boolean isLayerFile(Path file) {
    * @return the diff ID portion of the layer file filename
    * @throws CacheCorruptedException if no valid diff ID could be parsed
    */
-  DescriptorDigest getDiffId(Path layerFile) throws CacheCorruptedException {
+  DescriptorDigest getDigestFromFilename(Path layerFile) throws CacheCorruptedException {
     try {
-      String diffId = layerFile.getFileName().toString();
-      return DescriptorDigest.fromHash(diffId);
+      String hash = layerFile.getFileName().toString();
+      return DescriptorDigest.fromHash(hash);
 
     } catch (DigestException | IndexOutOfBoundsException ex) {
       throw new CacheCorruptedException(
-          cacheDirectory, "Layer file did not include valid diff ID: " + layerFile, ex);
+          cacheDirectory, "Layer file did not include valid hash: " + layerFile, ex);
     }
   }
 
@@ -125,6 +126,15 @@ Path getLayerDirectory(DescriptorDigest layerDigest) {
     return getLayersDirectory().resolve(layerDigest.getHash());
   }
 
+  /**
+   * Resolves the {@link #LOCAL_DIRECTORY} in the {@link #cacheDirectory}.
+   *
+   * @return the directory containing local base image layers
+   */
+  Path getLocalDirectory() {
+    return cacheDirectory.resolve(LOCAL_DIRECTORY);
+  }
+
   /**
    * Gets the directory to store the image manifest and configuration.
    *

jib-core/src/main/java/com/google/cloud/tools/jib/cache/CacheStorageReader.java

@@ -168,7 +168,7 @@ Optional<CachedLayer> retrieve(DescriptorDigest layerDigest)
       if (layerFiles.size() != 1) {
         throw new CacheCorruptedException(
             cacheStorageFiles.getCacheDirectory(),
-            "No or multiple layer files found for layer with digest "
+            "No or multiple layer files found for layer hash "
                 + layerDigest.getHash()
                 + " in directory: "
                 + layerDirectory);
@@ -180,7 +180,45 @@ Optional<CachedLayer> retrieve(DescriptorDigest layerDigest)
               .setLayerDigest(layerDigest)
               .setLayerSize(Files.size(layerFile))
               .setLayerBlob(Blobs.from(layerFile))
-              .setLayerDiffId(cacheStorageFiles.getDiffId(layerFile))
+              .setLayerDiffId(cacheStorageFiles.getDigestFromFilename(layerFile))
+              .build());
+    }
+  }
+
+  /**
+   * Retrieves the {@link CachedLayer} for the local base image layer with the given diff ID.
+   *
+   * @param diffId the diff ID
+   * @return the {@link CachedLayer} referenced by the diff ID, if found
+   * @throws CacheCorruptedException if the cache was found to be corrupted
+   * @throws IOException if an I/O exception occurs
+   */
+  Optional<CachedLayer> retrieveTarLayer(DescriptorDigest diffId)
+      throws IOException, CacheCorruptedException {
+    Path layerDirectory = cacheStorageFiles.getLocalDirectory().resolve(diffId.getHash());
+    if (!Files.exists(layerDirectory)) {
+      return Optional.empty();
+    }
+
+    try (Stream<Path> files = Files.list(layerDirectory)) {
+      List<Path> layerFiles =
+          files.filter(CacheStorageFiles::isLayerFile).collect(Collectors.toList());
+      if (layerFiles.size() != 1) {
+        throw new CacheCorruptedException(
+            cacheStorageFiles.getCacheDirectory(),
+            "No or multiple layer files found for layer hash "
+                + diffId.getHash()
+                + " in directory: "
+                + layerDirectory);
+      }
+
+      Path layerFile = layerFiles.get(0);
+      return Optional.of(
+          CachedLayer.builder()
+              .setLayerDigest(cacheStorageFiles.getDigestFromFilename(layerFile))
+              .setLayerSize(Files.size(layerFile))
+              .setLayerBlob(Blobs.from(layerFile))
+              .setLayerDiffId(diffId)
               .build());
     }
   }

jib-core/src/main/java/com/google/cloud/tools/jib/cache/CacheStorageWriter.java

@@ -251,6 +251,47 @@ CachedLayer writeUncompressed(Blob uncompressedLayerBlob, @Nullable DescriptorDi
     }
   }
 
+  /**
+   * Saves a local base image layer.
+   *
+   * @param diffId the layer blob's diff ID
+   * @param compressedBlob the blob to save
+   * @throws IOException if an I/O exception occurs
+   */
+  CachedLayer writeTarLayer(DescriptorDigest diffId, Blob compressedBlob) throws IOException {
+    Files.createDirectories(cacheStorageFiles.getLocalDirectory());
+    Files.createDirectories(cacheStorageFiles.getTemporaryDirectory());
+    try (TemporaryDirectory temporaryDirectory =
+        new TemporaryDirectory(cacheStorageFiles.getTemporaryDirectory())) {
+      Path temporaryLayerDirectory = temporaryDirectory.getDirectory();
+      Path temporaryLayerFile = cacheStorageFiles.getTemporaryLayerFile(temporaryLayerDirectory);
+
+      BlobDescriptor layerBlobDescriptor;
+      try (OutputStream fileOutputStream =
+          new BufferedOutputStream(Files.newOutputStream(temporaryLayerFile))) {
+        layerBlobDescriptor = compressedBlob.writeTo(fileOutputStream);
+      }
+
+      // Renames the temporary layer file to its digest
+      // (temp/temp -> temp/<digest>)
+      String fileName = layerBlobDescriptor.getDigest().getHash();
+      Path digestLayerFile = temporaryLayerDirectory.resolve(fileName);
+      moveIfDoesNotExist(temporaryLayerFile, digestLayerFile);
+
+      // Moves the temporary directory to directory named with diff ID
+      // (temp/<digest> -> <diffID>/<digest>)
+      Path destination = cacheStorageFiles.getLocalDirectory().resolve(diffId.getHash());
+      moveIfDoesNotExist(temporaryLayerDirectory, destination);
+
+      return CachedLayer.builder()
+          .setLayerDigest(layerBlobDescriptor.getDigest())
+          .setLayerDiffId(diffId)
+          .setLayerSize(layerBlobDescriptor.getSize())
+          .setLayerBlob(Blobs.from(destination.resolve(fileName)))
+          .build();
+    }
+  }
+
   /**
    * Saves the manifest and container configuration for a V2.2 or OCI image.
    *
@@ -286,7 +327,7 @@ void writeMetadata(ImageReference imageReference, V21ManifestTemplate manifestTe
     Path imageDirectory = cacheStorageFiles.getImageDirectory(imageReference);
     Files.createDirectories(imageDirectory);
 
-    try (LockFile ignored1 = LockFile.lock(imageDirectory.resolve("lock"))) {
+    try (LockFile ignored = LockFile.lock(imageDirectory.resolve("lock"))) {
       writeMetadata(manifestTemplate, imageDirectory.resolve("manifest.json"));
     }
   }

jib-core/src/test/java/com/google/cloud/tools/jib/builder/steps/ExtractTarStepTest.java

@@ -18,6 +18,8 @@
 
 import com.google.cloud.tools.jib.builder.ProgressEventDispatcher;
 import com.google.cloud.tools.jib.builder.steps.ExtractTarStep.LocalImage;
+import com.google.cloud.tools.jib.cache.Cache;
+import com.google.cloud.tools.jib.cache.CacheCorruptedException;
 import com.google.cloud.tools.jib.configuration.BuildConfiguration;
 import com.google.cloud.tools.jib.event.EventHandlers;
 import com.google.cloud.tools.jib.image.LayerCountMismatchException;
@@ -54,7 +56,9 @@ private static Path getResource(String resource) throws URISyntaxException {
   }
 
   @Before
-  public void setup() {
+  public void setup() throws IOException {
+    Mockito.when(buildConfiguration.getBaseImageLayersCache())
+        .thenReturn(Cache.withDirectory(temporaryFolder.newFolder().toPath()));
     Mockito.when(buildConfiguration.getEventHandlers()).thenReturn(eventHandlers);
     Mockito.when(progressEventDispatcherFactory.create(Mockito.anyString(), Mockito.anyLong()))
         .thenReturn(progressEventDispatcher);
@@ -66,7 +70,7 @@ public void setup() {
   @Test
   public void testCall_validDocker()
       throws URISyntaxException, LayerCountMismatchException,
-          BadContainerConfigurationFormatException, IOException {
+          BadContainerConfigurationFormatException, IOException, CacheCorruptedException {
     Path dockerBuild = getResource("core/extraction/docker-save.tar");
     LocalImage result =
         new ExtractTarStep(buildConfiguration, dockerBuild, progressEventDispatcherFactory).call();
@@ -91,7 +95,7 @@ public void testCall_validDocker()
   @Test
   public void testCall_validTar()
       throws URISyntaxException, LayerCountMismatchException,
-          BadContainerConfigurationFormatException, IOException {
+          BadContainerConfigurationFormatException, IOException, CacheCorruptedException {
     Path tarBuild = getResource("core/extraction/jib-image.tar");
     LocalImage result =
         new ExtractTarStep(buildConfiguration, tarBuild, progressEventDispatcherFactory).call();