src/com/facebook/buck/util/ZipSplitter.java - buck - Git at Google

 /*
  * Copyright 2012-present Facebook, Inc.
  *
  * Licensed under the Apache License, Version 2.0 (the "License"); you may
  * not use this file except in compliance with the License. You may obtain
  * a copy of the License at
  *
  *     http://www.apache.org/licenses/LICENSE-2.0
  *
  * Unless required by applicable law or agreed to in writing, software
  * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
  * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
  * License for the specific language governing permissions and limitations
  * under the License.
  */

 package com.facebook.buck.util;

 import com.google.common.base.Preconditions;
 import com.google.common.base.Predicate;
 import com.google.common.base.Throwables;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Sets;
 import com.google.common.io.ByteStreams;

 import java.io.BufferedOutputStream;
 import java.io.Closeable;
 import java.io.File;
 import java.io.FileNotFoundException;
 import java.io.FileOutputStream;
 import java.io.IOException;
 import java.io.InputStream;
 import java.util.Collection;
 import java.util.Set;
 import java.util.zip.ZipEntry;
 import java.util.zip.ZipOutputStream;

 public class ZipSplitter {

   public static enum DexSplitStrategy {
     MAXIMIZE_PRIMARY_DEX_SIZE,
     MINIMIZE_PRIMARY_DEX_SIZE;
   }

   private final Set<File> inFiles;
   private final File outPrimary;
   private final File outSecondaryDir;
   private final String secondaryPattern;
   private final long zipSizeSoftLimit;
   private final long zipSizeHardLimit;
   private final Predicate<String> requiredInPrimaryZip;
   private final DexSplitStrategy dexSplitStrategy;

   private final ImmutableList.Builder<File> secondaryFiles = ImmutableList.builder();
   private long remainingSize;
   private int currentSecondaryIndex;
   private ZipOutputStreamHelper primaryOut;
   private ZipOutputStreamHelper currentSecondaryOut;
   private boolean newSecondaryOutOnNextEntry;

   private ZipSplitter(
       Set<File> inFiles,
       File outPrimary,
       File outSecondaryDir,
       String secondaryPattern,
       long zipSizeSoftLimit,
       long zipSizeHardLimit,
       Predicate<String> requiredInPrimaryZip,
       DexSplitStrategy dexSplitStrategy) {
     this.inFiles = ImmutableSet.copyOf(inFiles);
     this.outPrimary = Preconditions.checkNotNull(outPrimary);
     this.outSecondaryDir = Preconditions.checkNotNull(outSecondaryDir);
     this.secondaryPattern = Preconditions.checkNotNull(secondaryPattern);
     this.zipSizeSoftLimit = zipSizeSoftLimit;
     this.zipSizeHardLimit = zipSizeHardLimit;
     this.requiredInPrimaryZip = Preconditions.checkNotNull(requiredInPrimaryZip);
     this.dexSplitStrategy = dexSplitStrategy;
   }

   /**
    * Both combines and splits a set of input files into zip files such that no one output zip file
    * has entries that in total exceed {@code zipSizeHardLimit}.  Input files can be themselves zip
    * files, individual class/resource files, or a directory of such of files.  The inputs are
    * "opened", and the files contained within them are individually processed.
    * <p>
    * For example, given a set of inFiles of A.zip and B.zip where A.zip contains { A, B, C }, and
    * B.zip contains { X, Y, Z }, a possible outcome of this method could yield outPrimary.zip
    * containing { A, B }, outSecondary1.zip containing { C, X }, and outSecondary2.zip containing
    * { Y, Z }.
    * <p>
    * This method exists as a critical utility to divide source code so large that dx/dexopt fail
    * due to design constraints.
    *
    * @param inFiles Set of input files (directories or zip files) whose contents should be placed in
    *     the output zip files.
    * @param outPrimary Primary output zip file.
    * @param outSecondaryDir Directory to place secondary output zip files (if any are generated).
    * @param secondaryPattern Pattern containing a single integer (%d) that forms the filename of
    *     output zip files placed in {@code outSecondaryDir}.
    * @param zipSizeSoftLimit Soft limit for the resulting zip file.  Once hit, a new output zip
    *     file will be created after processing the current {@code inFiles} entry.  This soft
    *     limit affects only secondary zip files.  Has no effect if this value is greater than or
    *     equal to {@code zipSizeHardLimit}.
    * @param zipSizeHardLimit Maximum size of the entries in each output zip file.  A new output
    *     zip file will be created immediately before this limit is exceeded regardless of whether
    *     it straddles a single {@code inFiles} entry.
    * @param requiredInPrimaryZip Determine which input <em>entries</em> are necessary in the
    *     primary output zip file.  Note that this is referring to the entries contained within
    *     {@code inFiles}, not the input files themselves.
    * @return Secondary output zip files.
    * @throws IOException
    */
   public static Collection<File> splitZip(
       Set<File> inFiles,
       File outPrimary,
       File outSecondaryDir,
       String secondaryPattern,
       long zipSizeSoftLimit,
       long zipSizeHardLimit,
       Predicate<String> requiredInPrimaryZip,
       DexSplitStrategy dexSplitStrategy) throws IOException {
     ZipSplitter splitter = new ZipSplitter(
         inFiles,
         outPrimary,
         outSecondaryDir,
         secondaryPattern,
         zipSizeSoftLimit,
         zipSizeHardLimit,
         requiredInPrimaryZip,
         dexSplitStrategy);
     return splitter.execute();
   }

   // Not safe to execute multiple times.
   private Collection<File> execute() throws IOException {
     ClasspathTraverser classpathTraverser = new DefaultClasspathTraverser();

     // Compute the total size of the inputs so that we can figure out whether its safe
     // to begin putting non-essential entries into the primary zip.
     // TODO(devjasta): There's a more compact way of doing this by writing the primary zip during
     // this first-pass step then assigning it as the "currentSecondaryOut" to complete the second
     // pass.  We're already tracking unique entries so we would not end up adding those primary
     // entries twice.
     classpathTraverser.traverse(new ClasspathTraversal(inFiles) {
       @Override
       public void visit(FileLike entry) {
         long entrySize = entry.getSize();
         if (entrySize > 0) {
           remainingSize += entrySize;
         }
       }
     });

     currentSecondaryIndex = 0;
     primaryOut = newZipOutput(outPrimary);
     currentSecondaryOut = null;

     try {
       for (File inFile : inFiles) {
         classpathTraverser.traverse(new ClasspathTraversal(ImmutableSet.of(inFile)) {
           @Override
           public void visit(FileLike entry) {
             try {
               processEntry(entry);
             } catch (IOException e) {
               throw new RuntimeException(e);
             }
           }
         });

         // The soft limit was tripped (and not the hard limit).  Flag that the next non-zero length
         // entry should create a new zip.
         if (currentSecondaryOut != null &&
             currentSecondaryOut.getCurrentSize() >= zipSizeSoftLimit) {
           newSecondaryOutOnNextEntry = true;
         }
       }
     } catch (RuntimeException e) {
       Throwables.propagateIfInstanceOf(e.getCause(), IOException.class);
       throw Throwables.propagate(e);
     } finally {
       primaryOut.close();
       if (currentSecondaryOut != null) {
         currentSecondaryOut.close();
       }
     }

     return secondaryFiles.build();
   }

   private void processEntry(FileLike entry) throws IOException {
     long entrySize = entry.getSize();
     if (entrySize <= 0) {
       return;
     }
     if (entrySize > zipSizeHardLimit) {
       throw new IllegalArgumentException(
           "Single entry larger than limit: " + entry);
     }

     ZipOutputStreamHelper targetOut;

     // An entry is placed in the primary zip if either of the following is true:
     //
     // (1) The entry must appear in the first zip according to the EntryProcessor predicate.
     // (2) All of the remaining zip entries fit in the remaining space in the primary zip and
     //     we're trying to maximize the size of the primary zip.
     //
     // Otherwise, the entry will be added to the secondary zip.
     boolean canFitAllRemaining = remainingSize + primaryOut.getCurrentSize() <= zipSizeHardLimit;

     if ((canFitAllRemaining && dexSplitStrategy == DexSplitStrategy.MAXIMIZE_PRIMARY_DEX_SIZE)
         || requiredInPrimaryZip.apply(entry.getRelativePath())        // File must be in primary
         ) {
       // Going to write this entry to the primary zip.
       if (!primaryOut.canPutEntry(entry)) {
         throw new IllegalArgumentException(
             "Unable to fit all required files in primary zip.");
       }
       targetOut = primaryOut;
     } else {
       // Going to write this entry to a secondary zip.
       if (currentSecondaryOut == null ||
           !currentSecondaryOut.canPutEntry(entry) ||
           newSecondaryOutOnNextEntry) {
         if (currentSecondaryOut != null) {
           currentSecondaryOut.close();
         }
         currentSecondaryIndex++;
         File newSecondaryFile = new File(
             outSecondaryDir,
             String.format(secondaryPattern, currentSecondaryIndex));
         secondaryFiles.add(newSecondaryFile);
         currentSecondaryOut = newZipOutput(newSecondaryFile);
         newSecondaryOutOnNextEntry = false;
         // We've already tested for this. It really shouldn't happen.
         Preconditions.checkState(currentSecondaryOut.canPutEntry(entry));
       }
       targetOut = currentSecondaryOut;
     }

     targetOut.putEntry(entry);
     remainingSize -= entrySize;
   }

   private ZipOutputStreamHelper newZipOutput(File file) throws FileNotFoundException {
     return new ZipOutputStreamHelper(
         new ZipOutputStream(new BufferedOutputStream(new FileOutputStream(file))));
   }

   private class ZipOutputStreamHelper implements Closeable {
     private final ZipOutputStream outStream;
     private final Set<String> entryNames = Sets.newHashSet();

     private long currentSize;

     private ZipOutputStreamHelper(ZipOutputStream outStream) {
       this.outStream = outStream;
     }

     public long getCurrentSize() {
       return currentSize;
     }

     private boolean isEntryTooBig(long entrySize) {
       return (currentSize + entrySize > zipSizeHardLimit);
     }

     /**
      * Tests whether the file-like instance can be placed into the zip entry without
      * exceeding the maximum size limit.
      * @param fileLike File-like instance to test.
      * @return True if the file-like instance is small enough to fit; false otherwise.
      * @see #putEntry
      */
     public boolean canPutEntry(FileLike fileLike) {
       return !isEntryTooBig(fileLike.getSize());
     }

     /**
      * Attempt to put the next entry.
      * @param fileLike File-like instance to add as a zip entry.
      * @throws IOException
      * @throws IllegalStateException Thrown if putting this entry would exceed the maximum size
      *     limit.  See {#link #canPutEntry}.
      */
     public void putEntry(FileLike fileLike) throws IOException {
       String name = fileLike.getRelativePath();
       // Tracks unique entry names and avoids duplicates.  This is, believe it or not, how
       // proguard seems to handle merging multiple -injars into a single -outjar.
       if (!entryNames.contains(name)) {
         entryNames.add(name);
         outStream.putNextEntry(new ZipEntry(name));
         InputStream in = fileLike.getInput();
         long entrySize;
         try {
           entrySize = ByteStreams.copy(in, outStream);
         } finally {
           in.close();
         }
         // Make sure FileLike#getSize didn't lie (or we forgot to call canPutEntry).
         Preconditions.checkState(!isEntryTooBig(entrySize),
             "Putting entry %s (%d) exceeded maximum size of %d", name, entrySize, zipSizeHardLimit);
         currentSize += entrySize;
       }
     }

     @Override
     public void close() throws IOException {
       outStream.close();
     }
   }
 }
	/*
	* Copyright 2012-present Facebook, Inc.
	*
	* Licensed under the Apache License, Version 2.0 (the "License"); you may
	* not use this file except in compliance with the License. You may obtain
	* a copy of the License at
	*
	* http://www.apache.org/licenses/LICENSE-2.0
	*
	* Unless required by applicable law or agreed to in writing, software
	* distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
	* WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
	* License for the specific language governing permissions and limitations
	* under the License.
	*/

	package com.facebook.buck.util;

	import com.google.common.base.Preconditions;
	import com.google.common.base.Predicate;
	import com.google.common.base.Throwables;
	import com.google.common.collect.ImmutableList;
	import com.google.common.collect.ImmutableSet;
	import com.google.common.collect.Sets;
	import com.google.common.io.ByteStreams;

	import java.io.BufferedOutputStream;
	import java.io.Closeable;
	import java.io.File;
	import java.io.FileNotFoundException;
	import java.io.FileOutputStream;
	import java.io.IOException;
	import java.io.InputStream;
	import java.util.Collection;
	import java.util.Set;
	import java.util.zip.ZipEntry;
	import java.util.zip.ZipOutputStream;

	public class ZipSplitter {

	public static enum DexSplitStrategy {
	MAXIMIZE_PRIMARY_DEX_SIZE,
	MINIMIZE_PRIMARY_DEX_SIZE;
	}

	private final Set<File> inFiles;
	private final File outPrimary;
	private final File outSecondaryDir;
	private final String secondaryPattern;
	private final long zipSizeSoftLimit;
	private final long zipSizeHardLimit;
	private final Predicate<String> requiredInPrimaryZip;
	private final DexSplitStrategy dexSplitStrategy;

	private final ImmutableList.Builder<File> secondaryFiles = ImmutableList.builder();
	private long remainingSize;
	private int currentSecondaryIndex;
	private ZipOutputStreamHelper primaryOut;
	private ZipOutputStreamHelper currentSecondaryOut;
	private boolean newSecondaryOutOnNextEntry;

	private ZipSplitter(
	Set<File> inFiles,
	File outPrimary,
	File outSecondaryDir,
	String secondaryPattern,
	long zipSizeSoftLimit,
	long zipSizeHardLimit,
	Predicate<String> requiredInPrimaryZip,
	DexSplitStrategy dexSplitStrategy) {
	this.inFiles = ImmutableSet.copyOf(inFiles);
	this.outPrimary = Preconditions.checkNotNull(outPrimary);
	this.outSecondaryDir = Preconditions.checkNotNull(outSecondaryDir);
	this.secondaryPattern = Preconditions.checkNotNull(secondaryPattern);
	this.zipSizeSoftLimit = zipSizeSoftLimit;
	this.zipSizeHardLimit = zipSizeHardLimit;
	this.requiredInPrimaryZip = Preconditions.checkNotNull(requiredInPrimaryZip);
	this.dexSplitStrategy = dexSplitStrategy;
	}

	/**
	* Both combines and splits a set of input files into zip files such that no one output zip file
	* has entries that in total exceed {@code zipSizeHardLimit}. Input files can be themselves zip
	* files, individual class/resource files, or a directory of such of files. The inputs are
	* "opened", and the files contained within them are individually processed.
	* <p>
	* For example, given a set of inFiles of A.zip and B.zip where A.zip contains { A, B, C }, and
	* B.zip contains { X, Y, Z }, a possible outcome of this method could yield outPrimary.zip
	* containing { A, B }, outSecondary1.zip containing { C, X }, and outSecondary2.zip containing
	* { Y, Z }.
	* <p>
	* This method exists as a critical utility to divide source code so large that dx/dexopt fail
	* due to design constraints.
	*
	* @param inFiles Set of input files (directories or zip files) whose contents should be placed in
	* the output zip files.
	* @param outPrimary Primary output zip file.
	* @param outSecondaryDir Directory to place secondary output zip files (if any are generated).
	* @param secondaryPattern Pattern containing a single integer (%d) that forms the filename of
	* output zip files placed in {@code outSecondaryDir}.
	* @param zipSizeSoftLimit Soft limit for the resulting zip file. Once hit, a new output zip
	* file will be created after processing the current {@code inFiles} entry. This soft
	* limit affects only secondary zip files. Has no effect if this value is greater than or
	* equal to {@code zipSizeHardLimit}.
	* @param zipSizeHardLimit Maximum size of the entries in each output zip file. A new output
	* zip file will be created immediately before this limit is exceeded regardless of whether
	* it straddles a single {@code inFiles} entry.
	* @param requiredInPrimaryZip Determine which input <em>entries</em> are necessary in the
	* primary output zip file. Note that this is referring to the entries contained within
	* {@code inFiles}, not the input files themselves.
	* @return Secondary output zip files.
	* @throws IOException
	*/
	public static Collection<File> splitZip(
	Set<File> inFiles,
	File outPrimary,
	File outSecondaryDir,
	String secondaryPattern,
	long zipSizeSoftLimit,
	long zipSizeHardLimit,
	Predicate<String> requiredInPrimaryZip,
	DexSplitStrategy dexSplitStrategy) throws IOException {
	ZipSplitter splitter = new ZipSplitter(
	inFiles,
	outPrimary,
	outSecondaryDir,
	secondaryPattern,
	zipSizeSoftLimit,
	zipSizeHardLimit,
	requiredInPrimaryZip,
	dexSplitStrategy);
	return splitter.execute();
	}

	// Not safe to execute multiple times.
	private Collection<File> execute() throws IOException {
	ClasspathTraverser classpathTraverser = new DefaultClasspathTraverser();

	// Compute the total size of the inputs so that we can figure out whether its safe
	// to begin putting non-essential entries into the primary zip.
	// TODO(devjasta): There's a more compact way of doing this by writing the primary zip during
	// this first-pass step then assigning it as the "currentSecondaryOut" to complete the second
	// pass. We're already tracking unique entries so we would not end up adding those primary
	// entries twice.
	classpathTraverser.traverse(new ClasspathTraversal(inFiles) {
	@Override
	public void visit(FileLike entry) {
	long entrySize = entry.getSize();
	if (entrySize > 0) {
	remainingSize += entrySize;
	}
	}
	});

	currentSecondaryIndex = 0;
	primaryOut = newZipOutput(outPrimary);
	currentSecondaryOut = null;

	try {
	for (File inFile : inFiles) {
	classpathTraverser.traverse(new ClasspathTraversal(ImmutableSet.of(inFile)) {
	@Override
	public void visit(FileLike entry) {
	try {
	processEntry(entry);
	} catch (IOException e) {
	throw new RuntimeException(e);
	}
	}
	});

	// The soft limit was tripped (and not the hard limit). Flag that the next non-zero length
	// entry should create a new zip.
	if (currentSecondaryOut != null &&
	currentSecondaryOut.getCurrentSize() >= zipSizeSoftLimit) {
	newSecondaryOutOnNextEntry = true;
	}
	}
	} catch (RuntimeException e) {
	Throwables.propagateIfInstanceOf(e.getCause(), IOException.class);
	throw Throwables.propagate(e);
	} finally {
	primaryOut.close();
	if (currentSecondaryOut != null) {
	currentSecondaryOut.close();
	}
	}

	return secondaryFiles.build();
	}

	private void processEntry(FileLike entry) throws IOException {
	long entrySize = entry.getSize();
	if (entrySize <= 0) {
	return;
	}
	if (entrySize > zipSizeHardLimit) {
	throw new IllegalArgumentException(
	"Single entry larger than limit: " + entry);
	}

	ZipOutputStreamHelper targetOut;

	// An entry is placed in the primary zip if either of the following is true:
	//
	// (1) The entry must appear in the first zip according to the EntryProcessor predicate.
	// (2) All of the remaining zip entries fit in the remaining space in the primary zip and
	// we're trying to maximize the size of the primary zip.
	//
	// Otherwise, the entry will be added to the secondary zip.
	boolean canFitAllRemaining = remainingSize + primaryOut.getCurrentSize() <= zipSizeHardLimit;

	if ((canFitAllRemaining && dexSplitStrategy == DexSplitStrategy.MAXIMIZE_PRIMARY_DEX_SIZE)
	\|\| requiredInPrimaryZip.apply(entry.getRelativePath()) // File must be in primary
	) {
	// Going to write this entry to the primary zip.
	if (!primaryOut.canPutEntry(entry)) {
	throw new IllegalArgumentException(
	"Unable to fit all required files in primary zip.");
	}
	targetOut = primaryOut;
	} else {
	// Going to write this entry to a secondary zip.
	if (currentSecondaryOut == null \|\|
	!currentSecondaryOut.canPutEntry(entry) \|\|
	newSecondaryOutOnNextEntry) {
	if (currentSecondaryOut != null) {
	currentSecondaryOut.close();
	}
	currentSecondaryIndex++;
	File newSecondaryFile = new File(
	outSecondaryDir,
	String.format(secondaryPattern, currentSecondaryIndex));
	secondaryFiles.add(newSecondaryFile);
	currentSecondaryOut = newZipOutput(newSecondaryFile);
	newSecondaryOutOnNextEntry = false;
	// We've already tested for this. It really shouldn't happen.
	Preconditions.checkState(currentSecondaryOut.canPutEntry(entry));
	}
	targetOut = currentSecondaryOut;
	}

	targetOut.putEntry(entry);
	remainingSize -= entrySize;
	}

	private ZipOutputStreamHelper newZipOutput(File file) throws FileNotFoundException {
	return new ZipOutputStreamHelper(
	new ZipOutputStream(new BufferedOutputStream(new FileOutputStream(file))));
	}

	private class ZipOutputStreamHelper implements Closeable {
	private final ZipOutputStream outStream;
	private final Set<String> entryNames = Sets.newHashSet();

	private long currentSize;

	private ZipOutputStreamHelper(ZipOutputStream outStream) {
	this.outStream = outStream;
	}

	public long getCurrentSize() {
	return currentSize;
	}

	private boolean isEntryTooBig(long entrySize) {
	return (currentSize + entrySize > zipSizeHardLimit);
	}

	/**
	* Tests whether the file-like instance can be placed into the zip entry without
	* exceeding the maximum size limit.
	* @param fileLike File-like instance to test.
	* @return True if the file-like instance is small enough to fit; false otherwise.
	* @see #putEntry
	*/
	public boolean canPutEntry(FileLike fileLike) {
	return !isEntryTooBig(fileLike.getSize());
	}

	/**
	* Attempt to put the next entry.
	* @param fileLike File-like instance to add as a zip entry.
	* @throws IOException
	* @throws IllegalStateException Thrown if putting this entry would exceed the maximum size
	* limit. See {#link #canPutEntry}.
	*/
	public void putEntry(FileLike fileLike) throws IOException {
	String name = fileLike.getRelativePath();
	// Tracks unique entry names and avoids duplicates. This is, believe it or not, how
	// proguard seems to handle merging multiple -injars into a single -outjar.
	if (!entryNames.contains(name)) {
	entryNames.add(name);
	outStream.putNextEntry(new ZipEntry(name));
	InputStream in = fileLike.getInput();
	long entrySize;
	try {
	entrySize = ByteStreams.copy(in, outStream);
	} finally {
	in.close();
	}
	// Make sure FileLike#getSize didn't lie (or we forgot to call canPutEntry).
	Preconditions.checkState(!isEntryTooBig(entrySize),
	"Putting entry %s (%d) exceeded maximum size of %d", name, entrySize, zipSizeHardLimit);
	currentSize += entrySize;
	}
	}

	@Override
	public void close() throws IOException {
	outStream.close();
	}
	}
	}