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>

org.eclipse.jgit.pgm/src/org/eclipse/jgit/pgm/debug/ReadReftable.java

@@ -70,7 +70,7 @@ class ReadReftable extends TextBuiltin {
 				BlockSource src = BlockSource.from(in);
 				ReftableReader reader = new ReftableReader(src)) {
 			try (RefCursor rc = ref != null
-					? reader.seekRef(ref)
+					? reader.seekPrefix(ref)
 					: reader.allRefs()) {
 				while (rc.next()) {
 					write(rc.getRef());

org.eclipse.jgit.test/tst/org/eclipse/jgit/internal/storage/reftable/MergedReftableTest.java

@@ -80,7 +80,7 @@ public class MergedReftableTest {
 		try (RefCursor rc = mr.seekRef(HEAD)) {
 			assertFalse(rc.next());
 		}
-		try (RefCursor rc = mr.seekRef(R_HEADS)) {
+		try (RefCursor rc = mr.seekPrefix(R_HEADS)) {
 			assertFalse(rc.next());
 		}
 	}
@@ -94,7 +94,7 @@ public class MergedReftableTest {
 		try (RefCursor rc = mr.seekRef(HEAD)) {
 			assertFalse(rc.next());
 		}
-		try (RefCursor rc = mr.seekRef(R_HEADS)) {
+		try (RefCursor rc = mr.seekPrefix(R_HEADS)) {
 			assertFalse(rc.next());
 		}
 	}
@@ -108,7 +108,7 @@ public class MergedReftableTest {
 		try (RefCursor rc = mr.seekRef(HEAD)) {
 			assertFalse(rc.next());
 		}
-		try (RefCursor rc = mr.seekRef(R_HEADS)) {
+		try (RefCursor rc = mr.seekPrefix(R_HEADS)) {
 			assertFalse(rc.next());
 		}
 	}

org.eclipse.jgit.test/tst/org/eclipse/jgit/internal/storage/reftable/ReftableTest.java

@@ -101,7 +101,7 @@ public class ReftableTest {
 		try (RefCursor rc = t.seekRef(HEAD)) {
 			assertFalse(rc.next());
 		}
-		try (RefCursor rc = t.seekRef(R_HEADS)) {
+		try (RefCursor rc = t.seekPrefix(R_HEADS)) {
 			assertFalse(rc.next());
 		}
 		try (LogCursor rc = t.allLogs()) {
@@ -317,10 +317,10 @@ public class ReftableTest {
 	public void namespaceNotFound() throws IOException {
 		Ref exp = ref(MASTER, 1);
 		ReftableReader t = read(write(exp));
-		try (RefCursor rc = t.seekRef("refs/changes/")) {
+		try (RefCursor rc = t.seekPrefix("refs/changes/")) {
 			assertFalse(rc.next());
 		}
-		try (RefCursor rc = t.seekRef("refs/tags/")) {
+		try (RefCursor rc = t.seekPrefix("refs/tags/")) {
 			assertFalse(rc.next());
 		}
 	}
@@ -332,12 +332,12 @@ public class ReftableTest {
 		Ref v1 = tag(V1_0, 3, 4);
 
 		ReftableReader t = read(write(master, next, v1));
-		try (RefCursor rc = t.seekRef("refs/tags/")) {
+		try (RefCursor rc = t.seekPrefix("refs/tags/")) {
 			assertTrue(rc.next());
 			assertEquals(V1_0, rc.getRef().getName());
 			assertFalse(rc.next());
 		}
-		try (RefCursor rc = t.seekRef("refs/heads/")) {
+		try (RefCursor rc = t.seekPrefix("refs/heads/")) {
 			assertTrue(rc.next());
 			assertEquals(MASTER, rc.getRef().getName());
 
@@ -484,7 +484,7 @@ public class ReftableTest {
 		try (RefCursor rc = t.allRefs()) {
 			assertFalse(rc.next());
 		}
-		try (RefCursor rc = t.seekRef("refs/heads/")) {
+		try (RefCursor rc = t.seekPrefix("refs/heads/")) {
 			assertFalse(rc.next());
 		}
 		try (LogCursor lc = t.allLogs()) {

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) {
@@ -254,6 +258,29 @@ public class DfsReftableDatabase extends DfsRefDatabase {
 		return new RefMap(prefix, all.toRefList(), none, none);
 	}
 
+	/** {@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 {

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);

org.eclipse.jgit/src/org/eclipse/jgit/internal/storage/reftable/MergedReftable.java

@@ -111,6 +111,16 @@ public class MergedReftable extends Reftable {
 		return m;
 	}
 
+	/** {@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 {

org.eclipse.jgit/src/org/eclipse/jgit/internal/storage/reftable/Reftable.java

@@ -108,24 +108,35 @@ 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.
 	 */
 	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.
 	 *
@@ -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()));
 		}
 	}
 

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;
 	}