Callback in PackWriter & BundleWriter.

Added callback in PackWriter and BundleWriter for the caller to get the
count of objects to write, and a chance to abort the write operation.

Change-Id: I1baeedcc6946b1093652de4a707fe597a577e526
Signed-off-by: Yuxuan 'fishy' Wang <fishywang@google.com>

4 files changed, 192 insertions, 0 deletions

diff --git a/org.eclipse.jgit.test/tst/org/eclipse/jgit/transport/BundleWriterTest.java b/org.eclipse.jgit.test/tst/org/eclipse/jgit/transport/BundleWriterTest.java
index 24cee0a6a1..bd1ec33fa6 100644
--- a/org.eclipse.jgit.test/tst/org/eclipse/jgit/transport/BundleWriterTest.java
+++ b/org.eclipse.jgit.test/tst/org/eclipse/jgit/transport/BundleWriterTest.java
@@ -61,6 +61,8 @@ import java.util.Set;
 import org.eclipse.jgit.errors.MissingBundlePrerequisiteException;
 import org.eclipse.jgit.errors.NotSupportedException;
 import org.eclipse.jgit.errors.TransportException;
+import org.eclipse.jgit.internal.storage.pack.PackWriter.ObjectCountCallback;
+import org.eclipse.jgit.internal.storage.pack.WriteAbortedException;
 import org.eclipse.jgit.lib.NullProgressMonitor;
 import org.eclipse.jgit.lib.ObjectId;
 import org.eclipse.jgit.lib.Ref;
@@ -147,6 +149,18 @@ public class BundleWriterTest extends SampleDataRepositoryTestCase {
 		}
 	}
 
+	@Test
+	public void testAbortWrite() throws Exception {
+		boolean caught = false;
+		try {
+			makeBundleWithCallback(
+					"refs/heads/aa", db.resolve("a").name(), null, false);
+		} catch (WriteAbortedException e) {
+			caught = true;
+		}
+		assertTrue(caught);
+	}
+
 	private static FetchResult fetchFromBundle(final Repository newRepo,
 			final byte[] bundle) throws URISyntaxException,
 			NotSupportedException, TransportException {
@@ -161,9 +175,17 @@ public class BundleWriterTest extends SampleDataRepositoryTestCase {
 	private byte[] makeBundle(final String name,
 			final String anObjectToInclude, final RevCommit assume)
 			throws FileNotFoundException, IOException {
+		return makeBundleWithCallback(name, anObjectToInclude, assume, true);
+	}
+
+	private byte[] makeBundleWithCallback(final String name,
+			final String anObjectToInclude, final RevCommit assume,
+			boolean value)
+			throws FileNotFoundException, IOException {
 		final BundleWriter bw;
 
 		bw = new BundleWriter(db);
+		bw.setObjectCountCallback(new NaiveObjectCountCallback(value));
 		bw.include(name, ObjectId.fromString(anObjectToInclude));
 		if (assume != null)
 			bw.assume(assume);
@@ -172,4 +194,19 @@ public class BundleWriterTest extends SampleDataRepositoryTestCase {
 		return out.toByteArray();
 	}
 
+	private static class NaiveObjectCountCallback
+			implements ObjectCountCallback {
+		private final boolean value;
+
+		NaiveObjectCountCallback(boolean value) {
+			this.value = value;
+		}
+
+		@Override
+		public void setObjectCount(long unused) throws WriteAbortedException {
+			if (!value)
+				throw new WriteAbortedException();
+		}
+	}
+
 }
diff --git a/org.eclipse.jgit/src/org/eclipse/jgit/internal/storage/pack/PackWriter.java b/org.eclipse.jgit/src/org/eclipse/jgit/internal/storage/pack/PackWriter.java
index 85533d6a33..e023f3e4b9 100644
--- a/org.eclipse.jgit/src/org/eclipse/jgit/internal/storage/pack/PackWriter.java
+++ b/org.eclipse.jgit/src/org/eclipse/jgit/internal/storage/pack/PackWriter.java
@@ -167,6 +167,31 @@ public class PackWriter implements AutoCloseable {
 		boolean contains(AnyObjectId objectId);
 	}
 
+	/**
+	 * A callback to tell caller the count of objects ASAP.
+	 *
+	 * @since 4.1
+	 */
+	public interface ObjectCountCallback {
+		/**
+		 * Invoked when the PackWriter has counted the objects to be written
+		 * to pack.
+		 * <p>
+		 * An ObjectCountCallback can use this information to decide whether
+		 * the
+		 * {@link #writePack(ProgressMonitor, ProgressMonitor, OutputStream)}
+		 * operation should be aborted.
+		 * <p>
+		 * This callback will be called exactly once.
+		 *
+		 * @param objectCount
+		 *            the count of the objects.
+		 * @throws WriteAbortedException
+		 *             to indicate that the write operation should be aborted.
+		 */
+		void setObjectCount(long objectCount) throws WriteAbortedException;
+	}
+
 	private static final Map<WeakReference<PackWriter>, Boolean> instances =
 			new ConcurrentHashMap<WeakReference<PackWriter>, Boolean>();
 
@@ -289,6 +314,8 @@ public class PackWriter implements AutoCloseable {
 
 	private CRC32 crc32;
 
+	private ObjectCountCallback callback;
+
 	/**
 	 * Create writer for specified repository.
 	 * <p>
@@ -359,6 +386,20 @@ public class PackWriter implements AutoCloseable {
 	}
 
 	/**
+	 * Set the {@code ObjectCountCallback}.
+	 * <p>
+	 * It should be set before calling
+	 * {@link #writePack(ProgressMonitor, ProgressMonitor, OutputStream)}.
+	 *
+	 * @return this object for chaining.
+	 * @since 4.1
+	 */
+	public PackWriter setObjectCountCallback(ObjectCountCallback callback) {
+		this.callback = callback;
+		return this;
+	}
+
+	/**
 	 * Check whether writer can store delta base as an offset (new style
 	 * reducing pack size) or should store it as an object id (legacy style,
 	 * compatible with old readers).
@@ -906,6 +947,9 @@ public class PackWriter implements AutoCloseable {
 	 *             an error occurred reading a local object's data to include in
 	 *             the pack, or writing compressed object data to the output
 	 *             stream.
+	 * @throws WriteAbortedException
+	 *             the write operation is aborted by
+	 *             {@link PackWriter.ObjectCountCallback}.
 	 */
 	public void writePack(ProgressMonitor compressMonitor,
 			ProgressMonitor writeMonitor, OutputStream packStream)
@@ -947,6 +991,8 @@ public class PackWriter implements AutoCloseable {
 
 		long objCnt = getObjectCount();
 		stats.totalObjects = objCnt;
+		if (callback != null)
+			callback.setObjectCount(objCnt);
 		beginPhase(PackingPhase.WRITING, writeMonitor, objCnt);
 		long writeStart = System.currentTimeMillis();
 		try {
diff --git a/org.eclipse.jgit/src/org/eclipse/jgit/internal/storage/pack/WriteAbortedException.java b/org.eclipse.jgit/src/org/eclipse/jgit/internal/storage/pack/WriteAbortedException.java
new file mode 100644
index 0000000000..1bb8743109
--- /dev/null
+++ b/org.eclipse.jgit/src/org/eclipse/jgit/internal/storage/pack/WriteAbortedException.java
@@ -0,0 +1,83 @@
+/*
+ * Copyright (C) 2015, Google Inc.
+ * and other copyright owners as documented in the project's IP log.
+ *
+ * This program and the accompanying materials are made available
+ * under the terms of the Eclipse Distribution License v1.0 which
+ * accompanies this distribution, is reproduced below, and is
+ * available at http://www.eclipse.org/org/documents/edl-v10.php
+ *
+ * All rights reserved.
+ *
+ * Redistribution and use in source and binary forms, with or
+ * without modification, are permitted provided that the following
+ * conditions are met:
+ *
+ * - Redistributions of source code must retain the above copyright
+ *   notice, this list of conditions and the following disclaimer.
+ *
+ * - Redistributions in binary form must reproduce the above
+ *   copyright notice, this list of conditions and the following
+ *   disclaimer in the documentation and/or other materials provided
+ *   with the distribution.
+ *
+ * - Neither the name of the Eclipse Foundation, Inc. nor the
+ *   names of its contributors may be used to endorse or promote
+ *   products derived from this software without specific prior
+ *   written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND
+ * CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES,
+ * INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES
+ * OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+ * ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR
+ * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
+ * SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT
+ * NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
+ * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+ * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT,
+ * STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+ * ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF
+ * ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ */
+
+package org.eclipse.jgit.internal.storage.pack;
+
+import java.io.IOException;
+
+/**
+ * An exception to be thrown when the write operation is aborted.
+ * <p>
+ * That can be thrown inside
+ * {@link PackWriter.ObjectCountCallback#setObjectCount(long)}.
+ *
+ * @since 4.1
+ */
+public class WriteAbortedException extends IOException {
+	/**
+	 * Construct a {@code WriteAbortedException}.
+	 */
+	public WriteAbortedException() {
+	}
+
+	/**
+	 * Construct a {@code WriteAbortedException}.
+	 *
+	 * @param s message describing the issue
+	 */
+	public WriteAbortedException(String s) {
+		super(s);
+	}
+
+	/**
+	 * Construct a {@code WriteAbortedException}.
+	 *
+	 * @param s
+	 *            message describing the issue
+	 * @param why
+	 *            a lower level implementation specific issue.
+	 */
+	public WriteAbortedException(String s, Throwable why) {
+		super(s, why);
+	}
+}
diff --git a/org.eclipse.jgit/src/org/eclipse/jgit/transport/BundleWriter.java b/org.eclipse.jgit/src/org/eclipse/jgit/transport/BundleWriter.java
index 81ad981918..34fb406e27 100644
--- a/org.eclipse.jgit/src/org/eclipse/jgit/transport/BundleWriter.java
+++ b/org.eclipse.jgit/src/org/eclipse/jgit/transport/BundleWriter.java
@@ -55,6 +55,8 @@ import java.util.TreeMap;
 
 import org.eclipse.jgit.internal.JGitText;
 import org.eclipse.jgit.internal.storage.pack.PackWriter;
+import org.eclipse.jgit.internal.storage.pack.PackWriter.ObjectCountCallback;
+import org.eclipse.jgit.internal.storage.pack.PackWriter.Statistics;
 import org.eclipse.jgit.lib.AnyObjectId;
 import org.eclipse.jgit.lib.Constants;
 import org.eclipse.jgit.lib.ObjectId;
@@ -92,6 +94,8 @@ public class BundleWriter {
 
 	private PackConfig packConfig;
 
+	private ObjectCountCallback callback;
+
 	/**
 	 * Create a writer for a bundle.
 	 *
@@ -188,6 +192,9 @@ public class BundleWriter {
 	 *             an error occurred reading a local object's data to include in
 	 *             the bundle, or writing compressed object data to the output
 	 *             stream.
+	 * @throws WriteAbortedException
+	 *             the write operation is aborted by
+	 *             {@link ObjectCountCallback}.
 	 */
 	public void writeBundle(ProgressMonitor monitor, OutputStream os)
 			throws IOException {
@@ -195,6 +202,8 @@ public class BundleWriter {
 		if (pc == null)
 			pc = new PackConfig(db);
 		try (PackWriter packWriter = new PackWriter(pc, db.newObjectReader())) {
+			packWriter.setObjectCountCallback(callback);
+
 			final HashSet<ObjectId> inc = new HashSet<ObjectId>();
 			final HashSet<ObjectId> exc = new HashSet<ObjectId>();
 			inc.addAll(include.values());
@@ -234,4 +243,21 @@ public class BundleWriter {
 			packWriter.writePack(monitor, monitor, os);
 		}
 	}
+
+	/**
+	 * Set the {@link ObjectCountCallback}.
+	 * <p>
+	 * It should be set before calling
+	 * {@link #writeBundle(ProgressMonitor, OutputStream)}.
+	 * <p>
+	 * This callback will be passed on to
+	 * {@link PackWriter#setObjectCountCallback}.
+	 *
+	 * @return this object for chaining.
+	 * @since 4.1
+	 */
+	public BundleWriter setObjectCountCallback(ObjectCountCallback callback) {
+		this.callback = callback;
+		return this;
+	}
 }