RepoSequence: Add method to get IDs in bulk

When we know a caller is going to need more than batchSize IDs, we can
request them all in a single increment rather than requiring multiple
writes to the sequence ref.

Change-Id: Iebf09e2d144c021797b70ac6e877894886521167

gerrit-server/src/main/java/com/google/gerrit/server/notedb/RepoSequence.java

@@ -23,6 +23,7 @@ import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.CharMatcher;
 import com.google.common.base.Predicates;
 import com.google.common.base.Throwables;
+import com.google.common.collect.ImmutableList;
 import com.google.common.primitives.Ints;
 import com.google.common.util.concurrent.Runnables;
 import com.google.gerrit.common.Nullable;
@@ -47,6 +48,8 @@ import org.eclipse.jgit.lib.Repository;
 import org.eclipse.jgit.revwalk.RevWalk;
 
 import java.io.IOException;
+import java.util.ArrayList;
+import java.util.List;
 import java.util.concurrent.Callable;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.TimeUnit;
@@ -135,7 +138,7 @@ public class RepoSequence {
     counterLock.lock();
     try {
       if (counter >= limit) {
-        acquire();
+        acquire(batchSize);
       }
       return counter++;
     } finally {
@@ -143,6 +146,30 @@ public class RepoSequence {
     }
   }
 
+  public ImmutableList<Integer> next(int count) throws OrmException {
+    if (count == 0) {
+      return ImmutableList.of();
+    }
+    checkArgument(count > 0, "count is negative: %s", count);
+    counterLock.lock();
+    try {
+      List<Integer> ids = new ArrayList<>(count);
+      while (counter < limit) {
+        ids.add(counter++);
+        if (ids.size() == count) {
+          return ImmutableList.copyOf(ids);
+        }
+      }
+      acquire(Math.max(count - ids.size(), batchSize));
+      while (ids.size() < count) {
+        ids.add(counter++);
+      }
+      return ImmutableList.copyOf(ids);
+    } finally {
+      counterLock.unlock();
+    }
+  }
+
   @VisibleForTesting
   public void set(int val) throws OrmException {
     // Don't bother spinning. This is only for tests, and a test that calls set
@@ -161,13 +188,13 @@ public class RepoSequence {
     }
   }
 
-  private void acquire() throws OrmException {
+  private void acquire(int count) throws OrmException {
     try (Repository repo = repoManager.openRepository(projectName);
         RevWalk rw = new RevWalk(repo)) {
-      TryAcquire attempt = new TryAcquire(repo, rw);
+      TryAcquire attempt = new TryAcquire(repo, rw, count);
       checkResult(retryer.call(attempt));
       counter = attempt.next;
-      limit = counter + batchSize;
+      limit = counter + count;
       acquireCount++;
     } catch (ExecutionException | RetryException e) {
       Throwables.propagateIfInstanceOf(e.getCause(), OrmException.class);
@@ -186,12 +213,14 @@ public class RepoSequence {
   private class TryAcquire implements Callable<RefUpdate.Result> {
     private final Repository repo;
     private final RevWalk rw;
+    private final int count;
 
     private int next;
 
-    private TryAcquire(Repository repo, RevWalk rw) {
+    private TryAcquire(Repository repo, RevWalk rw, int count) {
       this.repo = repo;
       this.rw = rw;
+      this.count = count;
     }
 
     @Override
@@ -206,7 +235,7 @@ public class RepoSequence {
         oldId = ref.getObjectId();
         next = parse(oldId);
       }
-      return store(repo, rw, oldId, next + batchSize);
+      return store(repo, rw, oldId, next + count);
     }
 
     private int parse(ObjectId id) throws IOException, OrmException {

gerrit-server/src/test/java/com/google/gerrit/server/notedb/RepoSequenceTest.java

@@ -77,7 +77,7 @@ public class RepoSequenceTest {
       RepoSequence s = newSequence(name, 1, batchSize);
       for (int i = 1; i <= max; i++) {
         try {
-          assertThat(s.next()).named("next for " + name).isEqualTo(i);
+          assertThat(s.next()).named("i=" + i + " for " + name).isEqualTo(i);
         } catch (OrmException e) {
           throw new AssertionError(
               "failed batchSize=" + batchSize + ", i=" + i, e);
@@ -89,6 +89,36 @@ public class RepoSequenceTest {
     }
   }
 
+  @Test
+  public void oneCallerNoLoop() throws Exception {
+    RepoSequence s = newSequence("id", 1, 3);
+    assertThat(s.acquireCount).isEqualTo(0);
+
+    assertThat(s.next()).isEqualTo(1);
+    assertThat(s.acquireCount).isEqualTo(1);
+    assertThat(s.next()).isEqualTo(2);
+    assertThat(s.acquireCount).isEqualTo(1);
+    assertThat(s.next()).isEqualTo(3);
+    assertThat(s.acquireCount).isEqualTo(1);
+
+    assertThat(s.next()).isEqualTo(4);
+    assertThat(s.acquireCount).isEqualTo(2);
+    assertThat(s.next()).isEqualTo(5);
+    assertThat(s.acquireCount).isEqualTo(2);
+    assertThat(s.next()).isEqualTo(6);
+    assertThat(s.acquireCount).isEqualTo(2);
+
+    assertThat(s.next()).isEqualTo(7);
+    assertThat(s.acquireCount).isEqualTo(3);
+    assertThat(s.next()).isEqualTo(8);
+    assertThat(s.acquireCount).isEqualTo(3);
+    assertThat(s.next()).isEqualTo(9);
+    assertThat(s.acquireCount).isEqualTo(3);
+
+    assertThat(s.next()).isEqualTo(10);
+    assertThat(s.acquireCount).isEqualTo(4);
+  }
+
   @Test
   public void twoCallers() throws Exception {
     RepoSequence s1 = newSequence("id", 1, 3);
@@ -193,6 +223,52 @@ public class RepoSequenceTest {
     s.next();
   }
 
+  @Test
+  public void nextWithCountOneCaller() throws Exception {
+    RepoSequence s = newSequence("id", 1, 3);
+    assertThat(s.next(2)).containsExactly(1, 2).inOrder();
+    assertThat(s.acquireCount).isEqualTo(1);
+    assertThat(s.next(2)).containsExactly(3, 4).inOrder();
+    assertThat(s.acquireCount).isEqualTo(2);
+    assertThat(s.next(2)).containsExactly(5, 6).inOrder();
+    assertThat(s.acquireCount).isEqualTo(2);
+
+    assertThat(s.next(3)).containsExactly(7, 8, 9).inOrder();
+    assertThat(s.acquireCount).isEqualTo(3);
+    assertThat(s.next(3)).containsExactly(10, 11, 12).inOrder();
+    assertThat(s.acquireCount).isEqualTo(4);
+    assertThat(s.next(3)).containsExactly(13, 14, 15).inOrder();
+    assertThat(s.acquireCount).isEqualTo(5);
+
+    assertThat(s.next(7)).containsExactly(16, 17, 18, 19, 20, 21, 22).inOrder();
+    assertThat(s.acquireCount).isEqualTo(6);
+    assertThat(s.next(7)).containsExactly(23, 24, 25, 26, 27, 28, 29).inOrder();
+    assertThat(s.acquireCount).isEqualTo(7);
+    assertThat(s.next(7)).containsExactly(30, 31, 32, 33, 34, 35, 36).inOrder();
+    assertThat(s.acquireCount).isEqualTo(8);
+  }
+
+  @Test
+  public void nextWithCountMultipleCallers() throws Exception {
+    RepoSequence s1 = newSequence("id", 1, 3);
+    RepoSequence s2 = newSequence("id", 1, 4);
+
+    assertThat(s1.next(2)).containsExactly(1, 2).inOrder();
+    assertThat(s1.acquireCount).isEqualTo(1);
+
+    // s1 hasn't exhausted its last batch.
+    assertThat(s2.next(2)).containsExactly(4, 5).inOrder();
+    assertThat(s2.acquireCount).isEqualTo(1);
+
+    // s1 acquires again to cover this request, plus a whole new batch.
+    assertThat(s1.next(3)).containsExactly(3, 8, 9);
+    assertThat(s1.acquireCount).isEqualTo(2);
+
+    // s2 hasn't exhausted its last batch, do so now.
+    assertThat(s2.next(2)).containsExactly(6, 7);
+    assertThat(s2.acquireCount).isEqualTo(1);
+  }
+
   private RepoSequence newSequence(String name, int start, int batchSize) {
     return newSequence(
         name, start, batchSize, Runnables.doNothing(), RETRYER);