Seek references by prefix in reftable

Reftable implementation of RefDatabase.getRefsByPrefix() should be
more performant, as references are filtered directly by prefix;
instead of fetching the whole subtree then filter by prefix.

Change-Id: If4f5f8c08285ea1eaec9efb83c3d864cea7a1321
Signed-off-by: Minh Thai <mthai@google.com>

5 files changed, 71 insertions, 18 deletions

diff --git a/org.eclipse.jgit/src/org/eclipse/jgit/internal/storage/dfs/DfsReftableDatabase.java b/org.eclipse.jgit/src/org/eclipse/jgit/internal/storage/dfs/DfsReftableDatabase.java
index 40cfb71dde..b891910c00 100644
--- a/org.eclipse.jgit/src/org/eclipse/jgit/internal/storage/dfs/DfsReftableDatabase.java
+++ b/org.eclipse.jgit/src/org/eclipse/jgit/internal/storage/dfs/DfsReftableDatabase.java
@@ -45,6 +45,9 @@ package org.eclipse.jgit.internal.storage.dfs;
 
 import java.io.IOException;
 import java.util.Arrays;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
 import java.util.Map;
 import java.util.concurrent.locks.ReentrantLock;
 
@@ -238,7 +241,8 @@ public class DfsReftableDatabase extends DfsRefDatabase {
 		try {
 			Reftable table = reader();
 			try (RefCursor rc = ALL.equals(prefix) ? table.allRefs()
-					: table.seekRef(prefix)) {
+					: (prefix.endsWith("/") ? table.seekPrefix(prefix) //$NON-NLS-1$
+							: table.seekRef(prefix))) {
 				while (rc.next()) {
 					Ref ref = table.resolve(rc.getRef());
 					if (ref != null && ref.getObjectId() != null) {
@@ -256,6 +260,29 @@ public class DfsReftableDatabase extends DfsRefDatabase {
 
 	/** {@inheritDoc} */
 	@Override
+	public List<Ref> getRefsByPrefix(String prefix) throws IOException {
+		List<Ref> all = new ArrayList<>();
+		lock.lock();
+		try {
+			Reftable table = reader();
+			try (RefCursor rc = ALL.equals(prefix) ? table.allRefs()
+					: table.seekPrefix(prefix)) {
+				while (rc.next()) {
+					Ref ref = table.resolve(rc.getRef());
+					if (ref != null && ref.getObjectId() != null) {
+						all.add(ref);
+					}
+				}
+			}
+		} finally {
+			lock.unlock();
+		}
+
+		return Collections.unmodifiableList(all);
+	}
+
+	/** {@inheritDoc} */
+	@Override
 	public Ref peel(Ref ref) throws IOException {
 		Ref oldLeaf = ref.getLeaf();
 		if (oldLeaf.isPeeled() || oldLeaf.getObjectId() == null) {
diff --git a/org.eclipse.jgit/src/org/eclipse/jgit/internal/storage/dfs/ReftableBatchRefUpdate.java b/org.eclipse.jgit/src/org/eclipse/jgit/internal/storage/dfs/ReftableBatchRefUpdate.java
index e0c056a450..47ac4ec72e 100644
--- a/org.eclipse.jgit/src/org/eclipse/jgit/internal/storage/dfs/ReftableBatchRefUpdate.java
+++ b/org.eclipse.jgit/src/org/eclipse/jgit/internal/storage/dfs/ReftableBatchRefUpdate.java
@@ -70,7 +70,6 @@ import org.eclipse.jgit.errors.MissingObjectException;
 import org.eclipse.jgit.internal.storage.dfs.DfsObjDatabase.PackSource;
 import org.eclipse.jgit.internal.storage.io.BlockSource;
 import org.eclipse.jgit.internal.storage.pack.PackExt;
-import org.eclipse.jgit.internal.storage.reftable.RefCursor;
 import org.eclipse.jgit.internal.storage.reftable.Reftable;
 import org.eclipse.jgit.internal.storage.reftable.ReftableCompactor;
 import org.eclipse.jgit.internal.storage.reftable.ReftableConfig;
@@ -240,11 +239,7 @@ public class ReftableBatchRefUpdate extends BatchRefUpdate {
 	private boolean checkExpected(Reftable table, List<ReceiveCommand> pending)
 			throws IOException {
 		for (ReceiveCommand cmd : pending) {
-			Ref ref;
-			try (RefCursor rc = table.seekRef(cmd.getRefName())) {
-				ref = rc.next() ? rc.getRef() : null;
-			}
-			if (!matchOld(cmd, ref)) {
+			if (!matchOld(cmd, table.exactRef(cmd.getRefName()))) {
 				cmd.setResult(LOCK_FAILURE);
 				if (isAtomic()) {
 					ReceiveCommand.abort(pending);
diff --git a/org.eclipse.jgit/src/org/eclipse/jgit/internal/storage/reftable/MergedReftable.java b/org.eclipse.jgit/src/org/eclipse/jgit/internal/storage/reftable/MergedReftable.java
index ef686a3008..7e22420921 100644
--- a/org.eclipse.jgit/src/org/eclipse/jgit/internal/storage/reftable/MergedReftable.java
+++ b/org.eclipse.jgit/src/org/eclipse/jgit/internal/storage/reftable/MergedReftable.java
@@ -113,6 +113,16 @@ public class MergedReftable extends Reftable {
 
 	/** {@inheritDoc} */
 	@Override
+	public RefCursor seekPrefix(String prefix) throws IOException {
+		MergedRefCursor m = new MergedRefCursor();
+		for (int i = 0; i < tables.length; i++) {
+			m.add(new RefQueueEntry(tables[i].seekPrefix(prefix), i));
+		}
+		return m;
+	}
+
+	/** {@inheritDoc} */
+	@Override
 	public RefCursor byObjectId(AnyObjectId name) throws IOException {
 		MergedRefCursor m = new MergedRefCursor();
 		for (int i = 0; i < tables.length; i++) {
diff --git a/org.eclipse.jgit/src/org/eclipse/jgit/internal/storage/reftable/Reftable.java b/org.eclipse.jgit/src/org/eclipse/jgit/internal/storage/reftable/Reftable.java
index 510c1a14e0..0593d94b20 100644
--- a/org.eclipse.jgit/src/org/eclipse/jgit/internal/storage/reftable/Reftable.java
+++ b/org.eclipse.jgit/src/org/eclipse/jgit/internal/storage/reftable/Reftable.java
@@ -108,18 +108,14 @@ public abstract class Reftable implements AutoCloseable {
 	public abstract RefCursor allRefs() throws IOException;
 
 	/**
-	 * Seek either to a reference, or a reference subtree.
+	 * Seek to a reference.
 	 * <p>
-	 * If {@code refName} ends with {@code "/"} the method will seek to the
-	 * subtree of all references starting with {@code refName} as a prefix. If
-	 * no references start with this prefix, an empty cursor is returned.
-	 * <p>
-	 * Otherwise exactly {@code refName} will be looked for. If present, the
+	 * This method will seek to the reference {@code refName}. If present, the
 	 * returned cursor will iterate exactly one entry. If not found, an empty
 	 * cursor is returned.
 	 *
 	 * @param refName
-	 *            reference name or subtree to find.
+	 *            reference name.
 	 * @return cursor to iterate; empty cursor if no references match.
 	 * @throws java.io.IOException
 	 *             if references cannot be read.
@@ -127,6 +123,21 @@ public abstract class Reftable implements AutoCloseable {
 	public abstract RefCursor seekRef(String refName) throws IOException;
 
 	/**
+	 * Seek references with prefix.
+	 * <p>
+	 * The method will seek all the references starting with {@code prefix} as a
+	 * prefix. If no references start with this prefix, an empty cursor is
+	 * returned.
+	 *
+	 * @param prefix
+	 *            prefix to find.
+	 * @return cursor to iterate; empty cursor if no references match.
+	 * @throws java.io.IOException
+	 *             if references cannot be read.
+	 */
+	public abstract RefCursor seekPrefix(String prefix) throws IOException;
+
+	/**
 	 * Match references pointing to a specific object.
 	 *
 	 * @param id
@@ -206,8 +217,9 @@ public abstract class Reftable implements AutoCloseable {
 	 *             if references cannot be read.
 	 */
 	public boolean hasRef(String refName) throws IOException {
-		try (RefCursor rc = seekRef(refName)) {
-			return rc.next();
+		try (RefCursor rc = seekPrefix(refName)) {
+			return rc.next() && (refName.endsWith("/") //$NON-NLS-1$
+					|| refName.equals(rc.getRef().getName()));
 		}
 	}
 
diff --git a/org.eclipse.jgit/src/org/eclipse/jgit/internal/storage/reftable/ReftableReader.java b/org.eclipse.jgit/src/org/eclipse/jgit/internal/storage/reftable/ReftableReader.java
index 5356952b5d..80e386a62c 100644
--- a/org.eclipse.jgit/src/org/eclipse/jgit/internal/storage/reftable/ReftableReader.java
+++ b/org.eclipse.jgit/src/org/eclipse/jgit/internal/storage/reftable/ReftableReader.java
@@ -183,9 +183,18 @@ public class ReftableReader extends Reftable {
 		initRefIndex();
 
 		byte[] key = refName.getBytes(CHARSET);
-		boolean prefix = key[key.length - 1] == '/';
+		RefCursorImpl i = new RefCursorImpl(refEnd, key, false);
+		i.block = seek(REF_BLOCK_TYPE, key, refIndex, 0, refEnd);
+		return i;
+	}
+
+	/** {@inheritDoc} */
+	@Override
+	public RefCursor seekPrefix(String prefix) throws IOException {
+		initRefIndex();
 
-		RefCursorImpl i = new RefCursorImpl(refEnd, key, prefix);
+		byte[] key = prefix.getBytes(CHARSET);
+		RefCursorImpl i = new RefCursorImpl(refEnd, key, true);
 		i.block = seek(REF_BLOCK_TYPE, key, refIndex, 0, refEnd);
 		return i;
 	}